Add MMS (Massively Multilingual Speech) CTC ASR model (#96)

## Summary
Adds support for facebook/mms-300m and facebook/mms-1b-all — Meta's
Massively Multilingual Speech models covering 1,100+ languages.

## Architecture

- **Feature extractor**: 7-layer causal CNN (reuses existing
`Wav2Vec2Model`)
- **Encoder**: transformer encoder with positional convolutions
- **Language adapter** (`add_adapter=True`): N strided Conv1d + GLU
layers for language-specific adaptation
- **CTC head**: linear projection → per-frame logits for connectionist
temporal classification

## Implementation

- `src/mobius/models/wav2vec2_ctc.py` — `Wav2Vec2ForCTCModel` (extends
`Wav2Vec2Model`), `_AdapterLayer`, `_Adapter`
- `src/mobius/tasks/_ctc_asr.py` — `CTCAsrTask` (single-model,
`input_values` + `attention_mask` → `logits`)
- `src/mobius/_configs.py` — `MMSConfig` (subclass of
`ArchitectureConfig` with adapter fields)
- Registry: `mms` → `Wav2Vec2ForCTCModel` / `ctc-asr`

## Weight name alignment

Uses bare `nn.Parameter` for adapter conv weights (matching the wav2vec2
feature extractor pattern) with `preprocess_weights` renames for HF's
`nn.Conv1d` layout.

## Testing

6 new tests in `TestBuildMMSGraph` — all pass. Full suite: 2676 passed.

---------

Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
Signed-off-by: Justin Chu <11205048+justinchuby@users.noreply.github.com>
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Co-authored-by: Justin Chu <11205048+justinchuby@users.noreply.github.com>
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>

Author: 3 people

examples/mms.py

@@ -38,6 +38,7 @@
 from __future__ import annotations
 
 import argparse
+import contextlib
 import fnmatch
 import sys
 import time
@@ -804,6 +805,88 @@ def _generate_audio_feature_extraction(case: TestCase, json_path: Path, device:
     )
 
 
+def _generate_ctc_asr(case: TestCase, json_path: Path, device: str) -> None:
+    """Generate golden data for CTC-based ASR (Wav2Vec2ForCTC / MMS).
+
+    The model output is per-frame logits over a vocabulary; we save the
+    top-K over the final frame's logit vector (matching the existing
+    audio-feature-extraction pattern), and when L5 is requested we also
+    save the CTC-greedy-decoded transcript as a token-id sequence so the
+    end-to-end test can compare against the runtime's greedy decode.
+
+    MMS specifically requires picking a target language adapter via
+    ``processor.tokenizer.set_target_lang(lang)`` and
+    ``model.load_adapter(lang)`` before the forward pass. The language is
+    read from ``case.generation_params['lang']`` (default ``"eng"``).
+    """
+    import librosa
+    import torch
+    import transformers
+
+    from mobius._testing.golden import save_generation_json, save_golden_ref
+
+    lang = case.generation_params.get("lang", "eng")
+
+    processor = transformers.AutoProcessor.from_pretrained(
+        case.model_id, trust_remote_code=case.trust_remote_code, target_lang=lang
+    )
+    model = transformers.Wav2Vec2ForCTC.from_pretrained(
+        case.model_id,
+        torch_dtype=torch.float32,
+        device_map=device,
+        trust_remote_code=case.trust_remote_code,
+        target_lang=lang,
+        ignore_mismatched_sizes=True,  # MMS lm_head shape changes per language
+    )
+    # For MMS, switching languages also requires loading the per-language adapter.
+    # Non-MMS Wav2Vec2ForCTC checkpoints don't have language adapters;
+    # the missing-adapter case is expected and harmless there.
+    if hasattr(model, "load_adapter"):
+        with contextlib.suppress(ValueError, KeyError, OSError):
+            model.load_adapter(lang)
+    model.eval()
+
+    audio_path = Path("testdata") / case.audio[0]
+    audio_array, sample_rate = librosa.load(str(audio_path), sr=16000)
+    processed = processor(audio_array, sampling_rate=sample_rate, return_tensors="pt").to(
+        next(model.parameters()).device
+    )
+
+    with torch.no_grad():
+        outputs = model(**processed)
+
+    # CTC logits: (batch, num_frames, vocab_size). Use last frame for top-K.
+    logits = outputs.logits[0]  # (num_frames, vocab_size)
+    last_logits = logits[-1].cpu().numpy()
+    golden = _extract_logits_golden(last_logits)
+
+    save_golden_ref(
+        json_path,
+        top1_id=golden["top1_id"],
+        top2_id=golden["top2_id"],
+        top10_ids=golden["top10_ids"],
+        top10_logits=golden["top10_logits"],
+        logits_summary=golden["logits_summary"],
+        input_ids=np.array([[0]], dtype=np.int64),  # placeholder
+    )
+
+    if "L5" in case.level:
+        # CTC greedy decode: argmax over vocab per frame, then collapse
+        # repeats and remove blanks. Save the post-collapse token IDs (and
+        # the decoded text for human inspection) into the standard
+        # ``*_generation.json`` sidecar.
+        predicted_ids = torch.argmax(logits, dim=-1).cpu().numpy()
+        transcript = processor.batch_decode(predicted_ids[np.newaxis, :])[0]
+        gen_path = json_path.with_name(json_path.stem + "_generation.json")
+        save_generation_json(
+            gen_path,
+            model_id=case.model_id,
+            prompt=str(audio_path),
+            generated_tokens=predicted_ids.tolist(),
+            generated_text=transcript,
+        )
+
+
 def _generate_image_classification(case: TestCase, json_path: Path, device: str) -> None:
     """Generate golden data for image classification (ViT, CLIP, etc.).
 
@@ -1186,6 +1269,7 @@ def _generate_phi4mm_multimodal(case: TestCase, json_path: Path, device: str) ->
     "speech-to-text": _generate_speech_to_text,
     "speech-language": _generate_speech_language,
     "audio-feature-extraction": _generate_audio_feature_extraction,
+    "ctc-asr": _generate_ctc_asr,
     # Vision tasks that produce last_hidden_state — reuse image classification.
     "depth-estimation": _generate_image_classification,
     "image-segmentation": _generate_image_classification,

scripts/generate_golden.py

@@ -22,6 +22,7 @@
     "ModelRegistration",
     "ModelRegistry",
     "ModelTask",
+    "MMSConfig",
     "OPSET_VERSION",
     "Sam2Config",
     "SegformerConfig",
@@ -67,6 +68,7 @@
     Gemma4Config,
     MambaConfig,
     MllamaConfig,
+    MMSConfig,
     Sam2Config,
     SegformerConfig,
     VisionConfig,

src/mobius/__init__.py

@@ -370,6 +370,15 @@ def build(
             model_type = "qwen3_5_moe_vl"
         hf_config = hf_config.text_config
 
+    # Wav2Vec2 / HuBERT / WavLM ship ``model_type="wav2vec2"`` (etc.) for both
+    # feature-extraction and CTC checkpoints. Switch to the ``mms`` registration
+    # (Wav2Vec2ForCTCModel + ctc-asr task) when the architecture indicates a
+    # CTC head — this covers both MMS and vanilla Wav2Vec2ForCTC fine-tunes.
+    if model_type in ("wav2vec2", "hubert", "wavlm"):
+        architectures = getattr(parent_config, "architectures", None) or []
+        if any("ForCTC" in arch for arch in architectures):
+            model_type = "mms"
+
     if module_class is None:
         if model_type in registry:
             module_class = registry.get(model_type)

src/mobius/_builder.py

@@ -39,6 +39,7 @@
     Mamba2Config,
     MambaConfig,
     MllamaConfig,
+    MMSConfig,
     NanoChatConfig,
     NemotronHConfig,
     Sam2Config,
@@ -98,6 +99,7 @@
     "Mamba2Config",
     "MambaConfig",
     "MllamaConfig",
+    "MMSConfig",
     "NanoChatConfig",
     "NemotronHConfig",
     "QuantizationConfig",

src/mobius/_configs/__init__.py

@@ -2024,3 +2024,43 @@ def from_transformers(cls, config) -> WhisperConfig:
             options["dtype"] = resolved
 
         return cls(**options)
+
+
+@dataclasses.dataclass
+class MMSConfig(ArchitectureConfig):
+    """Configuration for MMS (Massively Multilingual Speech) CTC models.
+
+    Extends ``ArchitectureConfig`` with the adapter parameters used in
+    ``facebook/mms-1b-all`` and related checkpoints.  When ``add_adapter=True``
+    the adapter layers are included in the exported ONNX graph; set this after
+    calling ``model.load_adapter(lang_code)`` to bake a specific language's
+    weights into the model.
+
+    HuggingFace class: ``Wav2Vec2ForCTC`` with ``config.model_type == "wav2vec2"``
+    """
+
+    add_adapter: bool = False
+    output_hidden_size: int = 0  # 0 → use hidden_size
+    adapter_kernel_size: int = 3
+    adapter_stride: int = 2
+    num_adapter_layers: int = 3
+
+    def __post_init__(self):
+        if self.output_hidden_size == 0:
+            self.output_hidden_size = self.hidden_size
+
+    @classmethod
+    def from_transformers(cls, config, parent_config=None) -> MMSConfig:
+        """Extract MMSConfig from a HuggingFace Wav2Vec2Config."""
+        base = ArchitectureConfig.from_transformers(config, parent_config=parent_config)
+        base_fields = _shallow_fields(base)
+        return cls(
+            **base_fields,
+            add_adapter=getattr(config, "add_adapter", False),
+            output_hidden_size=getattr(
+                config, "output_hidden_size", base_fields["hidden_size"]
+            ),
+            adapter_kernel_size=getattr(config, "adapter_kernel_size", 3),
+            adapter_stride=getattr(config, "adapter_stride", 2),
+            num_adapter_layers=getattr(config, "num_adapter_layers", 3),
+        )

src/mobius/_configs/_base.py

@@ -25,6 +25,7 @@
 from mobius._configs import (
     BaseModelConfig,
     Gemma4Config,
+    MMSConfig,
     WhisperConfig,
 )
 from mobius.models import (
@@ -126,6 +127,7 @@
 from mobius.models.trocr import TrOCRForConditionalGeneration
 from mobius.models.vit import ViTModel
 from mobius.models.wav2vec2 import Wav2Vec2Model
+from mobius.models.wav2vec2_ctc import Wav2Vec2ForCTCModel
 from mobius.models.xlm import XLMCausalLMModel
 from mobius.models.yolos import YolosForObjectDetection
 from mobius.models.zamba2 import Zamba2CausalLMModel
@@ -661,6 +663,7 @@ def _detect_fallback_registration(hf_config) -> ModelRegistration | None:
     "wav2vec2-bert": ModelRegistration(Wav2Vec2Model, task="audio-feature-extraction"),
     "wav2vec2-conformer": ModelRegistration(Wav2Vec2Model, task="audio-feature-extraction"),
     "wavlm": ModelRegistration(Wav2Vec2Model, task="audio-feature-extraction"),
+    "mms": ModelRegistration(Wav2Vec2ForCTCModel, task="ctc-asr", config_class=MMSConfig),
 }
 
 
@@ -862,6 +865,7 @@ def _create_default_registry() -> ModelRegistry:
     "qwen3_asr": "Qwen/Qwen3-ASR-0.6B",
     "fun_asr": "justinchuby/Fun-ASR-Nano-2512",
     "sensevoice_small": "mlx-community/SenseVoiceSmall",
+    "mms": "facebook/mms-300m",
     "speecht5": "microsoft/speecht5_asr",
     "sew": "asapp/sew-tiny-100k",
     "sew-d": "asapp/sew-d-tiny-100k",
@@ -1080,6 +1084,7 @@ def _create_default_registry() -> ModelRegistry:
     "wav2vec2-conformer": "wav2vec2",
     "hubert": "wav2vec2",
     "wavlm": "wav2vec2",
+    "mms": "wav2vec2",
     "vit": "vit",
     "vit_hybrid": "vit",
     "vit_mae": "vit",

src/mobius/_registry.py

@@ -126,6 +126,7 @@
     "UNet2DConditionModel",
     "ViTModel",
     "VideoAutoencoderModel",
+    "Wav2Vec2ForCTCModel",
     "Wav2Vec2Model",
     "WhisperForConditionalGeneration",
     "XLMCausalLMModel",
@@ -259,6 +260,7 @@
 from mobius.models.video_vae import VideoAutoencoderModel
 from mobius.models.vit import ViTModel
 from mobius.models.wav2vec2 import Wav2Vec2Model
+from mobius.models.wav2vec2_ctc import Wav2Vec2ForCTCModel
 from mobius.models.whisper import WhisperForConditionalGeneration
 from mobius.models.xlm import XLMCausalLMModel
 from mobius.models.zamba2 import Zamba2CausalLMModel