Method mainLoop

plugins/wasi_nn/MLX/model/whisper/decoding.cpp:647–708

}

std::tuple<mx::array, mx::array, mx::array>
DecodingTask::mainLoop(const mx::array &AudioFeatures,
                       const mx::array &Tokens) {

  int NBatch = Tokens.shape(0);
  mx::array CurrentTokens = Tokens;
  mx::array SumLogprobs = mx::zeros({NBatch}, mx::float32);
  bool Completed = false;

  auto StepFunction = [&](const mx::array &Inputs, const mx::array &AudioFeats,
                          const mx::array &TokSeq, const mx::array &SumLogp)
      -> std::tuple<mx::array, bool, mx::array, mx::array> {
    mx::array PreLogits = Inference->logits(Inputs, AudioFeats);
    mx::array Logits = take(PreLogits, PreLogits.shape(1) - 1, 1);
    for (const auto &Filter : LogitFilters) {
      Logits = Filter->apply(Logits, TokSeq);
    }
    auto [NextTokens, CompletedFlag, NextSumLogprobs] =
        Decoder->update(TokSeq, Logits, SumLogp);
    return std::make_tuple(NextTokens, CompletedFlag, NextSumLogprobs,
                           PreLogits);
  };

  auto [NextTokens, CompletedFlag, NextSumLogprobs, PreLogits] =
      StepFunction(CurrentTokens, AudioFeatures, CurrentTokens, SumLogprobs);

  CurrentTokens = NextTokens;
  SumLogprobs = NextSumLogprobs;
  Completed = CompletedFlag;

  mx::array NoSpeechProbs = mx::zeros({NBatch}, mx::float32);
  if (Tokenizer->getNoSpeech() != -1) {
    auto ProbsAtSot = mx::softmax(mx::take(PreLogits, SotIndex, 1), -1);
    NoSpeechProbs = mx::take(ProbsAtSot, Tokenizer->getNoSpeech(), 1);
  } else {
    NoSpeechProbs = mx::full({NBatch}, std::numeric_limits<float>::quiet_NaN(),
                             mx::float32);
  }

  mx::eval(CurrentTokens, SumLogprobs, NoSpeechProbs);
  for (int I = 1; I < SampleLen; ++I) {
    mx::array Inputs =
        take(CurrentTokens, mx::array({CurrentTokens.shape(1) - 1}), 1);

    if (CurrentTokens.shape(-1) > NCtx) {
      break;
    }
    auto [NextToks, NextCompleted, NextSumLogp, _] =
        StepFunction(Inputs, AudioFeatures, CurrentTokens, SumLogprobs);
    mx::eval(NextToks, NextSumLogp);

    if (Completed) {
      break;
    }

    CurrentTokens = NextToks;
    Completed = NextCompleted;
    SumLogprobs = NextSumLogp;