Skip to content

fix(TextDecoder): don't early-exit the decode loop on predictions sampled during promptTokens prefill#497

Open
yangzichao wants to merge 1 commit into
argmaxinc:mainfrom
yangzichao:upstream-pr/prompt-tokens-prefill-early-exit
Open

fix(TextDecoder): don't early-exit the decode loop on predictions sampled during promptTokens prefill#497
yangzichao wants to merge 1 commit into
argmaxinc:mainfrom
yangzichao:upstream-pr/prompt-tokens-prefill-early-exit

Conversation

@yangzichao

Copy link
Copy Markdown

Problem

Fixes #372.

With DecodingOptions.promptTokens, the initial prompt becomes [<|startofprev|>, ...prompt..., <|startoftranscript|>, ...] and decodeText force-feeds every one of those tokens through the main loop. The model's sampled prediction at each forced step is discarded (the forced token wins), but it still drives the early-exit checks:

  1. If the model's top prediction at any mid-prompt position is <|endoftext|>, sampleResult.completed ends the segment before a single audio token is decoded → empty transcript.
  2. isFirstToken points at prefilledIndex (the <|startofprev|> step), so firstTokenLogProbThreshold is evaluated against a forced-prompt placeholder prediction instead of the first real token — this is the spurious Fallback #1 (firstTokenLogProbThreshold) reported in Using promptTokens causes the Transcription to return empty result. #372.

Whether a model predicts EOT at those positions is model-dependent, which is why large-v3-v20240930_turbo and base.en reproduce while base appears fine — but the checks shouldn't be consulting forced-prefill predictions on any model. The reference implementation feeds the prompt without sampling checks and only starts them at sample_begin.

Fix

When promptTokens are present, ignore sampleResult.completed during prefill steps and evaluate the first-token check at the first genuinely sampled position (the last prefill index). The no-prompt path is untouched — checks keep their existing indices, so default behavior is unchanged.

Verification

Standalone harness against the real openai_whisper-large-v3-v20240930_turbo_632MB CoreML model on recordings that reproduced the empty decode:

case before after
no prompt ok ok (unchanged)
19-token prompt (stripped to content tokens) empty ok
single content token empty (2/3 clips) ok

With the fix, the prompt also demonstrably conditions the output (a punctuated prompt yields punctuated transcripts), i.e. the channel now works as intended rather than merely not crashing.

…pled during promptTokens prefill

With DecodingOptions.promptTokens, the initial prompt becomes
[<|startofprev|>, ...prompt..., <|startoftranscript|>, ...] and decodeText
force-feeds every one of those tokens through the main loop. The model's
sampled prediction at each forced step is discarded (the forced token wins),
but it still drove the early-exit checks:

1. If the model's top prediction at ANY mid-prompt position is
   <|endoftext|>, sampleResult.completed ended the segment before a single
   audio token was decoded -> empty transcript.
2. isFirstToken pointed at prefilledIndex (the <|startofprev|> step), so
   firstTokenLogProbThreshold was evaluated against a forced-prompt
   placeholder prediction instead of the first real token.

On whisper-large-v3-v20240930_turbo this made ANY non-empty promptTokens
(even a single content token) return an empty result, while the same audio
decoded fine without a prompt (argmaxinc#372). Whether the
model predicts EOT at those positions is model-dependent, which is why some
models appeared unaffected.

Fix: when promptTokens are present, ignore sampleResult.completed during
prefill steps and evaluate the first-token check at the first genuinely
sampled position (the last prefill index), matching the reference
implementation's sample_begin semantics. The no-prompt path is untouched.

Verified against the real 632MB turbo model on recordings that previously
reproduced the empty decode: with this change, no-prompt / 19-token prompt /
single-token prompt all return correct transcripts.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Using promptTokens causes the Transcription to return empty result.

2 participants