Skip to content

Add Whisper beam-search decoding - #1439

Open
sfeng01-tech wants to merge 1 commit into
ml-explore:mainfrom
sfeng01-tech:whisper-beam-search
Open

sfeng01-tech wants to merge 1 commit into
ml-explore:mainfrom
sfeng01-tech:whisper-beam-search

Conversation

@sfeng01-tech

Copy link
Copy Markdown

Summary

Adds beam-search decoding to the MLX Whisper example.

  • Implements beam candidate expansion, ranking, and finalization.
  • Reorders decoder KV caches after each beam selection.
  • Repeats audio features per beam so cross-attention caches remain aligned.
  • Adds --beam-size CLI support and API/README examples.
  • Adds regression coverage for beam decoding and option validation.
  • Adds an acknowledgment for Sahil Faizal.

Testing

  • PYTHONPATH=. ../../.venv-whisper/bin/python -m unittest test.TestWhisper.test_decode_beam_search test.TestWhisper.test_beam_search_options

    • Passed locally on Metal.
  • PYTHONPATH=. ../../.venv-whisper/bin/python -m unittest test.TestWhisper

    • The beam-search tests pass.
    • Existing numerical-parity assertions fail with MLX 0.32.0 in test_torch_mlx, test_decode_lang, and test_decode_greedy.

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.

1 participant