Add MMS (Massively Multilingual Speech) CTC ASR model - #96
Merged
Conversation
Performance Comparison
|
🏗️ Architecture Diff
No architecture changes detected. ✅ Legend: ⚪ No change · 🔵 Minor (attrs/inits) · 🟡 Moderate (nodes added/removed) · 🔴 Major (interface changed) |
justinchuby
marked this pull request as draft
April 2, 2026 15:03
![github-code-quality[bot]](https://avatars.githubusercontent.com/u/9919?s=60&v=4){kind=small}
Implements facebook/mms-300m and facebook/mms-1b-all as ONNX models using the mobius 4-layer stack. Architecture: - Wav2Vec2ForCTCModel: feature extractor (7-layer CNN) + transformer encoder + optional language adapter + CTC projection head - CTCAsrTask: single-model task (input_values + attention_mask → logits) - MMSConfig: ArchitectureConfig subclass with adapter fields (add_adapter, output_hidden_size, adapter_kernel_size, adapter_stride, num_adapter_layers) - _AdapterLayer: strided Conv1d → GLU using bare nn.Parameter (matches onnxscript path-resolution semantics) with preprocess_weights renames for HF's .conv.weight / .conv.bias → .conv / .conv_bias Weight name alignment: - HF wav2vec2.feature_extractor.conv_layers.N.conv.weight → feature_extractor.conv_layers.N.conv - HF wav2vec2.encoder.layer_norm.* → top-level layer_norm.* - HF wav2vec2.adapter.layers.N.conv.weight → adapter.layers.N.conv - HF intermediate_dense/output_dense → up_proj/down_proj Testing: - TestBuildMMSGraph: 6 tests (package builds, IO contract, CTC head initializers, adapter variant, registry lookup, ORT inference) - MMSConfig added to _COVERAGE_SKIP with reason Example: - examples/mms.py: CLI for multilingual CTC ASR with language adapter loading, audio file / mic input, greedy CTC decoding, multi-language demo Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
examples/mms.py: - Replace build(hf_model) with build_from_module() pattern: build() takes a model_id string, not a PyTorch model object. Because language adapter weights are loaded into the HF model via hf_model.load_adapter(lang), we must apply the resulting state_dict directly rather than re-downloading from Hub (which would reset the adapter). - Add build_mms_package() helper that correctly wires config extraction, graph construction, and state-dict application. - Import build_from_module + MMSConfig instead of build. src/mobius/models/wav2vec2_ctc.py: - Remove dead 'if TYPE_CHECKING: pass' block. - Expand attention_mask docstring: note that callers can pass all-ones for non-padded inputs and clarify the 1=valid/0=padding convention. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
Add three new modes to the MMS ASR example: --stream: Process audio in short windows (default 1s) and display partial transcript in real-time with a progress bar, overwriting the current terminal line as each window decodes. --live: Capture from the default microphone continuously, buffer audio in configurable windows, decode each window immediately and print the timestamped result. Stops on Ctrl-C. Requires sounddevice. --play: Play audio through the default output device in a background thread while transcribing. Requires sounddevice; silently skipped if not installed. --window SECONDS: Shared window size for --stream and --live (default 1.0s). New helpers: play_audio() — background playback via sounddevice _clear_line() — ANSI terminal line-erase utility transcribe_streaming() — windowed decode with live progress bar live_microphone_transcription() — real-time mic loop with queue Batch mode and all existing flags (--audio, --mic, --chunk, --save-to, --list-langs) are unchanged. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
justinchuby
force-pushed
the
justinchu/mms-models
branch
from
June 2, 2026 20:35
2fff1f7 to
d0cc0cd
Compare
Member
Author
|
Rebased on current main and re-verified. The branch was 156 commits behind and the merge state was
Also updated Implementation review:
Lint fixes: replaced ambiguous Verification on H200:
Implementation is simple, focused, and consistent with the existing model/task patterns in the repo. Ready for review. |
Codecov Report❌ Patch coverage is 📢 Thoughts on this report? Let us know! |
- wav2vec2_ctc.py: fix Wav2Vec2AdapterLayer Conv1d padding — was hard-coded to 1, which only matched HF for kernel_size=3. Now uses kernel_size // 2 (matching HF's Wav2Vec2AdapterLayer(padding=kernel_size // 2)). Real correctness bug for any non-default adapter_kernel_size. - _ctc_asr.py: fix CTCAsrTask docstring to reflect that attention_mask is a required graph input, not optional. Pass an all-ones mask when there is no padding. - examples/mms.py: add explanatory comment to the live-mic Ctrl-C handler so the "empty except" finding goes away while preserving the existing cleanup flow. - examples/mms.py: fix transcribe_streaming docstring vs behaviour mismatch. Old docstring claimed the overlap tail was "trimmed from decoded output", but the implementation appended every segment verbatim, duplicating tokens at chunk boundaries when overlap_seconds > 0. Set the default to 0 and document the trade-off honestly so callers opting in know what to expect. Verified locally on H200: 8/8 TestBuildMMSGraph pass, full suite still 2773 passed + 43 skipped, ruff check + format clean. Signed-off-by: Justin Chu <11205048+justinchuby@users.noreply.github.com>
Member
Author
|
Addressed the Copilot review (commit 80f73da):
Local re-verification on H200:
|
- New _generate_ctc_asr generator in scripts/generate_golden.py: loads Wav2Vec2ForCTC + AutoProcessor with target_lang, calls load_adapter, runs a forward pass and saves top-K of the last frame (L4) plus a CTC-greedy-decoded transcript sidecar (L5). - Wire 'ctc-asr' into _HIDDEN_STATE_TASKS and the e2e dispatcher in tests/e2e_golden_test.py; adds an all-ones attention_mask alongside input_values since _prepare_audio_feeds only emits the latter. - Add testdata/cases/audio/mms-1b-all.yaml (L4+L5) using the standard testdata 652-129742-0006.flac fixture. Currently skipped: the full 1B-param graph build OOMs host RAM during ONNX construction; goldens are committed for future regression coverage once the build path is fixed. - Auto-dispatch wav2vec2/hubert/wavlm checkpoints whose architecture contains 'ForCTC' to the 'mms' (Wav2Vec2ForCTCModel + ctc-asr task) registration so HF model_type='wav2vec2' + Wav2Vec2ForCTC reaches the CTC code path. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <11205048+justinchuby@users.noreply.github.com>
Member
Author
|
Added L4 + L5 golden test infrastructure for CTC-ASR:
The L4 e2e is currently marked |
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
Member
Author
|
Addressed the actionable Copilot review items in db077ed:
Remaining Copilot comments are not actionable here:
|
- Switch Apache-2.0/ONNX-Project-Contributors headers to Microsoft MIT
on the three new files (src/mobius/tasks/_ctc_asr.py,
src/mobius/models/wav2vec2_ctc.py, examples/mms.py) to match the
rest of the repo.
- Add model_roles = {"model": "encoder"} to CTCAsrTask so
build_from_module() classifies the graph as an encoder and skips
decoder-only optimization passes.
(Copilot's other new findings are stale or false positives:
- 'UNet tests indented inside TestBuildMMSGraph' — not present, all
8 MMS tests still pass; UNet/DiT classes are separate.
- 'mms registry key ineffective' — already fixed in 1cf7cfc by adding
the wav2vec2/hubert/wavlm + ForCTC → mms architecture override in
_builder.py.
- LFM2/audio-to-audio comments target files dropped during rebase.)
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Signed-off-by: Justin Chu <11205048+justinchuby@users.noreply.github.com>
- testdata/cases/schema.json: register 'ctc-asr' in the task_type enum and add 'lang' under generation properties so the new mms-1b-all case validates. yaml_schema_test now passes (225/225). - Drop suppressible try/except in _generate_ctc_asr (RUFF/SIM105) in favour of contextlib.suppress. - Ruff-format blank line in _registry.py. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <11205048+justinchuby@users.noreply.github.com>
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
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.
Summary
Adds support for facebook/mms-300m and facebook/mms-1b-all — Meta's Massively Multilingual Speech models covering 1,100+ languages.
Architecture
Wav2Vec2Model)add_adapter=True): N strided Conv1d + GLU layers for language-specific adaptationImplementation
src/mobius/models/wav2vec2_ctc.py—Wav2Vec2ForCTCModel(extendsWav2Vec2Model),_AdapterLayer,_Adaptersrc/mobius/tasks/_ctc_asr.py—CTCAsrTask(single-model,input_values+attention_mask→logits)src/mobius/_configs.py—MMSConfig(subclass ofArchitectureConfigwith adapter fields)mms→Wav2Vec2ForCTCModel/ctc-asrWeight name alignment
Uses bare
nn.Parameterfor adapter conv weights (matching the wav2vec2 feature extractor pattern) withpreprocess_weightsrenames for HF'snn.Conv1dlayout.Testing
6 new tests in
TestBuildMMSGraph— all pass. Full suite: 2676 passed.