fix(parakeet-mlx): support v2 (English) model — derive vocab/blank-id per model#9
Open
seidnerj wants to merge 2 commits into
Open
fix(parakeet-mlx): support v2 (English) model — derive vocab/blank-id per model#9seidnerj wants to merge 2 commits into
seidnerj wants to merge 2 commits into
Conversation
…ing v3 The MLX engine hardcoded the v3 (multilingual) vocabulary size (8192) for the blank id and the joint-output split. Loading the v2 (English, 1024-vocab) model produces a 1030-class joint output, but greedy TDT decode still sliced the duration head at index 8193 — an empty array — so argmax failed with "Cannot argmax reduce zero size array" on every segment, speech or silence. Derive num_classes from the loaded joint output projection and compute blank_id = num_classes - durations - 1 (= vocab_size) per model. Works for both v2 (->1024) and v3 (->8192). Verified end-to-end (examples/verify_v2_mlx.rs): parakeet-tdt-0.6b-v2-mlx now transcribes English speech correctly on GPU and returns empty without crashing on silence and sub-frame input.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What
Adds support for the English-specialized
mlx-community/parakeet-tdt-0.6b-v2MLX model and fixes a crash when loading any Parakeet TDT model whose
vocabulary differs from v3.
Why
The MLX Parakeet engine hardcoded v3's vocabulary size (8192) in
VOCAB_SIZE/BLANK_ID/NUM_CLASSES. The v2 (English) model has 1024 tokens.Because the blank id and class count were derived from the v3 constant, greedy
TDT decoding sliced the joint output's duration head at the wrong index — on v2
it produced an empty array and crashed every segment with:
Changes
parakeet-tdt-0.6b-v2→mlx-community/parakeet-tdt-0.6b-v2infrom_pretrainedandfrom_pretrained_cache_only.projection (
out_proj.weightfirst dim =num_classes), instead of the v3default constants. The constants now only allocate the modules before
load_weightsreplaces every parameter with the model's own tensor.duration table + blank.
examples/verify_v2_mlx.rsto verify v2 loads and transcribes on GPU.Testing
Verified
parakeet-tdt-0.6b-v2loads and transcribes correctly on Metal GPUwith zero decode errors; the v3 path is unchanged.