Add Parakeet CTC automatic speech recognition support

.github/workflows/gpu_l4_golden_parity.yml

@@ -5,6 +5,13 @@ on:
     # Run nightly at 4am UTC (after golden regen on Sunday)
     - cron: '0 4 * * *'
   workflow_dispatch:
+    inputs:
+      affected_models:
+        description: >-
+          JSON array of model names. Empty string or empty array runs all models.
+        required: false
+        type: string
+        default: ""
   workflow_call:
     inputs:
       affected_models:
@@ -48,8 +55,6 @@ jobs:
         with:
           path: ~/.cache/huggingface
           key: hf-gpu-${{ hashFiles('testdata/cases/**/*.yaml') }}
-          restore-keys: |
-            hf-gpu-
 
       - name: Install PyTorch (CUDA)
         run: pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu128

.github/workflows/gpu_l5_generation_e2e.yml

@@ -5,6 +5,13 @@ on:
     # Run nightly at 5am UTC (after L4 golden tests)
     - cron: '0 5 * * *'
   workflow_dispatch:
+    inputs:
+      affected_models:
+        description: >-
+          JSON array of model names. Empty string or empty array runs all models.
+        required: false
+        type: string
+        default: ""
   workflow_call:
     inputs:
       affected_models:
@@ -48,8 +55,6 @@ jobs:
         with:
           path: ~/.cache/huggingface
           key: hf-gpu-${{ hashFiles('testdata/cases/**/*.yaml') }}
-          restore-keys: |
-            hf-gpu-
 
       - name: Install PyTorch (CUDA)
         run: pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu128

scripts/generate_golden.py

@@ -963,7 +963,7 @@ 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).
+    """Generate golden data for raw-waveform or feature-input CTC ASR.
 
     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
@@ -984,21 +984,29 @@ def _generate_ctc_asr(case: TestCase, json_path: Path, device: str) -> None:
 
     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
-    )
+    processor_kwargs: dict[str, object] = {
+        "revision": case.revision,
+        "trust_remote_code": case.trust_remote_code,
+    }
+    model_kwargs: dict[str, object] = {
+        "revision": case.revision,
+        "torch_dtype": torch.float32,
+        "device_map": device,
+        "trust_remote_code": case.trust_remote_code,
+    }
+    if case.model_type == "mms":
+        processor_kwargs["target_lang"] = lang
+        model_kwargs.update(
+            target_lang=lang,
+            ignore_mismatched_sizes=True,
+        )
+
+    processor = transformers.AutoProcessor.from_pretrained(case.model_id, **processor_kwargs)
+    model = transformers.AutoModelForCTC.from_pretrained(case.model_id, **model_kwargs)
     # 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"):
+    if case.model_type == "mms" and hasattr(model, "load_adapter"):
         with contextlib.suppress(ValueError, KeyError, OSError):
             model.load_adapter(lang)
     model.eval()
@@ -1771,6 +1779,7 @@ def _hook(_module, _args, kwargs, output):
     "speech-language": _generate_speech_language,
     "audio-feature-extraction": _generate_audio_feature_extraction,
     "ctc-asr": _generate_ctc_asr,
+    "feature-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,

src/mobius/_builder.py

@@ -411,6 +411,7 @@ def build(
     model_id: str,
     task: str | ModelTask | None = None,
     *,
+    revision: str | None = None,
     module_class: type[nn.Module] | None = None,
     dtype: str | ir.DataType | None = None,
     output_layer_indices: list[int] | None = None,
@@ -443,6 +444,8 @@ def build(
         task: The model task. Either a task name string
             (e.g. ``"text-generation"``) or a :class:`ModelTask` instance.
             When ``None``, the task is auto-detected from the model type.
+        revision: Optional Hugging Face revision used consistently for the
+            configuration and every weight shard.
         module_class: Custom module class to use instead of the auto-detected
             one. The class must accept an :class:`ArchitectureConfig` as its
             constructor argument and have a ``forward()`` method compatible
@@ -537,14 +540,15 @@ def build(
     from mobius._diffusers_builder import build_diffusers_pipeline
 
     try:
-        hf_config = transformers.AutoConfig.from_pretrained(
-            model_id, trust_remote_code=trust_remote_code
-        )
+        config_kwargs = {"trust_remote_code": trust_remote_code}
+        if revision is not None:
+            config_kwargs["revision"] = revision
+        hf_config = transformers.AutoConfig.from_pretrained(model_id, **config_kwargs)
     except (ValueError, KeyError, OSError):
         # AutoConfig failed — the model_type may not be in transformers,
         # or the HF config class has a bug (e.g. NemotronH with '-' pattern).
         # Try loading config.json directly if the model is in our registry.
-        hf_config = _try_load_config_json(model_id)
+        hf_config = _try_load_config_json(model_id, revision=revision)
         if hf_config is None or hf_config.model_type not in registry:
             if text_only:
                 raise ValueError(
@@ -555,6 +559,7 @@ def build(
             # Not a model we support — try diffusers pipeline
             return build_diffusers_pipeline(
                 model_id,
+                revision=revision,
                 dtype=dtype,
                 load_weights=load_weights,
             )
@@ -687,7 +692,7 @@ def build(
         model.graph.name = f"{model_id}/{name}"
 
     if load_weights:
-        state_dict = _download_weights(model_id)
+        state_dict = _download_weights(model_id, revision=revision)
         if hasattr(model_module, "preprocess_weights"):
             state_dict = model_module.preprocess_weights(state_dict)
         prefix_map = getattr(model_module, "weight_prefix_map", None)

src/mobius/_builder_test.py

@@ -17,6 +17,7 @@
     _enable_prefill_prefix_pruning_task,
     _graph_requires_opset24,
     _maybe_apply_opset_lowering,
+    build,
     flags,
 )
 from mobius._model_package import ModelPackage
@@ -183,3 +184,47 @@ def test_maybe_apply_opset_lowering_skipped_when_flag_disabled(
     _maybe_apply_opset_lowering(pkg, execution_provider="cuda")
 
     assert pkg["embedding"].graph.opset_imports[""] == 24
+
+
+def test_build_threads_revision_to_diffusers_fallback(monkeypatch: pytest.MonkeyPatch) -> None:
+    import transformers
+
+    import mobius._config_resolver as config_resolver
+    import mobius._diffusers_builder as diffusers_builder
+
+    monkeypatch.setattr(
+        transformers.AutoConfig,
+        "from_pretrained",
+        lambda *args, **kwargs: (_ for _ in ()).throw(OSError("not transformers")),
+    )
+    monkeypatch.setattr(config_resolver, "_try_load_config_json", lambda *args, **kwargs: None)
+    expected = ModelPackage({})
+    calls: list[tuple[tuple, dict]] = []
+
+    def fake_build_diffusers(*args, **kwargs):
+        calls.append((args, kwargs))
+        return expected
+
+    monkeypatch.setattr(
+        diffusers_builder,
+        "build_diffusers_pipeline",
+        fake_build_diffusers,
+    )
+
+    result = build(
+        "fake/diffusers",
+        revision="pinned-revision",
+        load_weights=False,
+    )
+
+    assert result is expected
+    assert calls == [
+        (
+            ("fake/diffusers",),
+            {
+                "revision": "pinned-revision",
+                "dtype": None,
+                "load_weights": False,
+            },
+        )
+    ]

src/mobius/_config_resolver.py

@@ -76,7 +76,7 @@ def _default_task_for_model(model_type: str) -> str:
     return getattr(cls, "default_task", "text-generation")
 
 
-def _try_load_config_json(model_id: str):
+def _try_load_config_json(model_id: str, revision: str | None = None):
     """Try to load config.json directly for models not in transformers.
 
     Returns a ``PretrainedConfig``-like object with attribute access,
@@ -87,7 +87,10 @@ def _try_load_config_json(model_id: str):
     from huggingface_hub import hf_hub_download
 
     try:
-        path = hf_hub_download(repo_id=model_id, filename="config.json")
+        kwargs = {"repo_id": model_id, "filename": "config.json"}
+        if revision is not None:
+            kwargs["revision"] = revision
+        path = hf_hub_download(**kwargs)
     except (OSError, ValueError) as e:
         logger.debug("Failed to download config.json for %s: %s", model_id, e)
         return None

src/mobius/_configs/__init__.py

@@ -50,6 +50,7 @@
     NanoChatConfig,
     NemotronHConfig,
     NemotronParseConfig,
+    ParakeetCTCConfig,
     Qwen35MtpConfig,
     Sam2Config,
     SegformerConfig,
@@ -121,6 +122,7 @@
     "NanoChatConfig",
     "NemotronParseConfig",
     "NemotronHConfig",
+    "ParakeetCTCConfig",
     "QuantizationConfig",
     "Qwen35MtpConfig",
     "RoPEConfig",

src/mobius/_configs/_base.py

@@ -2930,3 +2930,51 @@ def from_transformers(cls, config, parent_config=None) -> MMSConfig:
             adapter_stride=getattr(config, "adapter_stride", 2),
             num_adapter_layers=getattr(config, "num_adapter_layers", 3),
         )
+
+
+@dataclasses.dataclass
+class ParakeetCTCConfig(ArchitectureConfig):
+    """Configuration for Hugging Face Parakeet FastConformer CTC models."""
+
+    num_mel_bins: int = 80
+    subsampling_factor: int = 8
+    subsampling_conv_channels: int = 256
+    subsampling_conv_kernel_size: int = 3
+    subsampling_conv_stride: int = 2
+    conv_kernel_size: int = 9
+    attention_bias: bool = True
+    convolution_bias: bool = True
+    scale_input: bool = True
+    layer_norm_eps: float = 1e-5
+
+    @classmethod
+    def from_transformers(cls, config, parent_config=None) -> ParakeetCTCConfig:
+        """Extract the nested Parakeet encoder and parent CTC vocabulary fields."""
+        encoder = getattr(config, "encoder_config", config)
+        parent = config if encoder is not config else parent_config
+        base = ArchitectureConfig.from_transformers(encoder, parent_config=parent)
+        fields = _shallow_fields(base)
+        fields.update(
+            vocab_size=getattr(parent, "vocab_size", fields["vocab_size"]),
+            pad_token_id=getattr(parent, "pad_token_id", fields["pad_token_id"]),
+            model_type=getattr(parent, "model_type", fields["model_type"]),
+        )
+        resolved_dtype = _resolve_dtype(parent)
+        # ORT CUDA executes this architecture in bf16 but diverges enough to
+        # collapse real CTC output to blanks. The checkpoint weights are fp32,
+        # so keep the safe fp32 default; callers may explicitly select fp16.
+        if resolved_dtype is not None and resolved_dtype != ir.DataType.BFLOAT16:
+            fields["dtype"] = resolved_dtype
+        return cls(
+            **fields,
+            num_mel_bins=getattr(encoder, "num_mel_bins", 80),
+            subsampling_factor=getattr(encoder, "subsampling_factor", 8),
+            subsampling_conv_channels=getattr(encoder, "subsampling_conv_channels", 256),
+            subsampling_conv_kernel_size=getattr(encoder, "subsampling_conv_kernel_size", 3),
+            subsampling_conv_stride=getattr(encoder, "subsampling_conv_stride", 2),
+            conv_kernel_size=getattr(encoder, "conv_kernel_size", 9),
+            attention_bias=getattr(encoder, "attention_bias", True),
+            convolution_bias=getattr(encoder, "convolution_bias", True),
+            scale_input=getattr(encoder, "scale_input", True),
+            layer_norm_eps=getattr(encoder, "layer_norm_eps", 1e-5),
+        )