diff --git a/.github/workflows/gpu_l4_golden_parity.yml b/.github/workflows/gpu_l4_golden_parity.yml index 65be966b8..91dfb2c8e 100644 --- a/.github/workflows/gpu_l4_golden_parity.yml +++ b/.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 diff --git a/.github/workflows/gpu_l5_generation_e2e.yml b/.github/workflows/gpu_l5_generation_e2e.yml index 13d17d4f8..23808927b 100644 --- a/.github/workflows/gpu_l5_generation_e2e.yml +++ b/.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 diff --git a/scripts/generate_golden.py b/scripts/generate_golden.py index 209e60b3e..338d46eaa 100644 --- a/scripts/generate_golden.py +++ b/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, diff --git a/src/mobius/_builder.py b/src/mobius/_builder.py index 52ed2402e..07ead5de6 100644 --- a/src/mobius/_builder.py +++ b/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) diff --git a/src/mobius/_builder_test.py b/src/mobius/_builder_test.py index 354308ed1..5bfaa2428 100644 --- a/src/mobius/_builder_test.py +++ b/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, + }, + ) + ] diff --git a/src/mobius/_config_resolver.py b/src/mobius/_config_resolver.py index ea24679b7..610c0d842 100644 --- a/src/mobius/_config_resolver.py +++ b/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 diff --git a/src/mobius/_configs/__init__.py b/src/mobius/_configs/__init__.py index f7b2f0968..aa855ff4c 100644 --- a/src/mobius/_configs/__init__.py +++ b/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", diff --git a/src/mobius/_configs/_base.py b/src/mobius/_configs/_base.py index 26c30b3cc..6cd3ef416 100644 --- a/src/mobius/_configs/_base.py +++ b/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), + ) diff --git a/src/mobius/_diffusers_builder.py b/src/mobius/_diffusers_builder.py index f95ee9e20..ed31f8359 100644 --- a/src/mobius/_diffusers_builder.py +++ b/src/mobius/_diffusers_builder.py @@ -108,7 +108,9 @@ def _init_diffusers_class_map() -> None: ) -def _load_diffusers_pipeline_index(model_id: str) -> dict | None: +def _load_diffusers_pipeline_index( + model_id: str, *, revision: str | None = None +) -> dict | None: """Try to load a diffusers ``model_index.json`` from HuggingFace. Returns the parsed JSON dict, or ``None`` if not found. @@ -116,7 +118,11 @@ def _load_diffusers_pipeline_index(model_id: str) -> dict | None: from huggingface_hub import hf_hub_download try: - path = hf_hub_download(repo_id=model_id, filename="model_index.json") + path = hf_hub_download( + repo_id=model_id, + filename="model_index.json", + revision=revision, + ) except (OSError, ValueError) as e: logger.debug("Failed to download model_index.json for %s: %s", model_id, e) return None @@ -126,7 +132,10 @@ def _load_diffusers_pipeline_index(model_id: str) -> dict | None: def _download_diffusers_component_weights( - model_id: str, component_name: str + model_id: str, + component_name: str, + *, + revision: str | None = None, ) -> dict[str, torch.Tensor]: """Download weights for a specific component of a diffusers pipeline. @@ -153,6 +162,7 @@ def _download_diffusers_component_weights( index_path = hf_hub_download( repo_id=model_id, filename=f"{prefix}{basename}.{ext}.index.json", + revision=revision, ) with open(index_path) as f: index = json.load(f) @@ -165,7 +175,11 @@ def _download_diffusers_component_weights( # Single-file weights. for basename in weight_basenames: try: - hf_hub_download(repo_id=model_id, filename=f"{prefix}{basename}.{ext}") + hf_hub_download( + repo_id=model_id, + filename=f"{prefix}{basename}.{ext}", + revision=revision, + ) all_files = [f"{basename}.{ext}"] break except EntryNotFoundError: @@ -182,6 +196,7 @@ def _download_diffusers_component_weights( paths = _parallel_download( model_id, [f"{prefix}{f}" for f in all_files], + revision=revision, desc=f"{component_name} weights", ) @@ -194,11 +209,20 @@ def _download_diffusers_component_weights( return state_dict -def _load_diffusers_component_config(model_id: str, component_name: str) -> dict: +def _load_diffusers_component_config( + model_id: str, + component_name: str, + *, + revision: str | None = None, +) -> dict: """Load the config.json for a specific diffusers pipeline component.""" from huggingface_hub import hf_hub_download - path = hf_hub_download(repo_id=model_id, filename=f"{component_name}/config.json") + path = hf_hub_download( + repo_id=model_id, + filename=f"{component_name}/config.json", + revision=revision, + ) with open(path) as f: return json.load(f) @@ -233,6 +257,7 @@ def _prepare_unet_loras(unet_loras: dict) -> tuple[tuple, dict]: def build_diffusers_pipeline( model_id: str, *, + revision: str | None = None, dtype: str | ir.DataType | None = None, load_weights: bool = True, unet_loras: dict | None = None, @@ -248,6 +273,7 @@ def build_diffusers_pipeline( Args: model_id: HuggingFace model repository ID for a diffusers pipeline. + revision: Optional Hugging Face revision used for all pipeline artifacts. dtype: Override the model dtype. load_weights: Whether to download and apply weights. unet_loras: Optional ``{adapter_name: lora.safetensors}`` map. Each LoRA @@ -264,7 +290,7 @@ def build_diffusers_pipeline( """ _init_diffusers_class_map() - pipeline_index = _load_diffusers_pipeline_index(model_id) + pipeline_index = _load_diffusers_pipeline_index(model_id, revision=revision) if pipeline_index is None: raise ValueError( f"'{model_id}' does not appear to be a diffusers pipeline " @@ -299,7 +325,11 @@ def build_diffusers_pipeline( class_name, ) - component_config_dict = _load_diffusers_component_config(model_id, component_name) + component_config_dict = _load_diffusers_component_config( + model_id, + component_name, + revision=revision, + ) config = config_class.from_diffusers(component_config_dict) if dtype is not None and hasattr(config, "dtype"): @@ -330,7 +360,11 @@ def build_diffusers_pipeline( package[f"{component_name}_{sub_name}"] = sub_model if load_weights: - state_dict = _download_diffusers_component_weights(model_id, component_name) + state_dict = _download_diffusers_component_weights( + model_id, + component_name, + revision=revision, + ) if hasattr(model_module, "preprocess_weights"): state_dict = model_module.preprocess_weights(state_dict) if lora_weights: diff --git a/src/mobius/_diffusers_builder_test.py b/src/mobius/_diffusers_builder_test.py index 8f1cf5dfd..03e1f4d4a 100644 --- a/src/mobius/_diffusers_builder_test.py +++ b/src/mobius/_diffusers_builder_test.py @@ -5,14 +5,17 @@ from __future__ import annotations -from unittest.mock import patch +from unittest.mock import mock_open, patch import onnx_ir as ir import pytest from mobius._diffusers_builder import ( _DIFFUSERS_CLASS_MAP, + _download_diffusers_component_weights, _init_diffusers_class_map, + _load_diffusers_component_config, + _load_diffusers_pipeline_index, build_diffusers_pipeline, ) from mobius._model_package import ModelPackage @@ -133,6 +136,70 @@ def test_raises_when_only_non_nn_components(self, _mock_load): build_diffusers_pipeline("fake/scheduler-only", load_weights=False) +# ── Hub revision propagation ───────────────────────────────────────────── + + +class TestDiffusersHubRevision: + @patch("huggingface_hub.hf_hub_download", return_value="model_index.json") + def test_pipeline_index_download_uses_revision(self, mock_download): + with patch("builtins.open", mock_open(read_data='{"_class_name": "FakePipeline"}')): + result = _load_diffusers_pipeline_index( + "fake/model", + revision="pinned-revision", + ) + + assert result == {"_class_name": "FakePipeline"} + mock_download.assert_called_once_with( + repo_id="fake/model", + filename="model_index.json", + revision="pinned-revision", + ) + + @patch("huggingface_hub.hf_hub_download", return_value="config.json") + def test_component_config_download_uses_revision(self, mock_download): + with patch("builtins.open", mock_open(read_data='{"in_channels": 3}')): + result = _load_diffusers_component_config( + "fake/model", + "vae", + revision="pinned-revision", + ) + + assert result == {"in_channels": 3} + mock_download.assert_called_once_with( + repo_id="fake/model", + filename="vae/config.json", + revision="pinned-revision", + ) + + @patch("mobius._diffusers_builder._parallel_download", return_value=[]) + @patch("huggingface_hub.hf_hub_download", return_value="weights.index.json") + def test_component_weight_downloads_use_revision( + self, + mock_download, + mock_parallel_download, + ): + index = '{"weight_map": {"weight": "model-00001-of-00001.safetensors"}}' + with patch("builtins.open", mock_open(read_data=index)): + result = _download_diffusers_component_weights( + "fake/model", + "vae", + revision="pinned-revision", + ) + + assert result == {} + mock_download.assert_called_once_with( + repo_id="fake/model", + filename="vae/diffusion_pytorch_model.safetensors.index.json", + revision="pinned-revision", + ) + mock_parallel_download.assert_called_once_with( + "fake/model", + ["vae/model-00001-of-00001.safetensors"], + revision="pinned-revision", + desc="vae weights", + ) + + # ── build_diffusers_pipeline component filtering ───────────────────────── @@ -400,6 +467,40 @@ def test_dtype_ir_datatype_passthrough( # Verify build_from_module was called (ir.DataType accepted without error) mock_build_from_module.assert_called_once() + @patch("mobius._diffusers_builder._download_diffusers_component_weights") + @patch("mobius._diffusers_builder.apply_weights") + @patch("mobius._diffusers_builder.build_from_module") + @patch("mobius._diffusers_builder._load_diffusers_component_config") + @patch("mobius._diffusers_builder._load_diffusers_pipeline_index") + def test_revision_propagates_to_all_pipeline_artifacts( + self, + mock_load_index, + mock_load_config, + mock_build_from_module, + mock_apply_weights, + mock_download_weights, + ): + mock_load_index.return_value = _fake_pipeline_index( + {"vae": ["diffusers", "AutoencoderKL"]} + ) + mock_load_config.return_value = {} + graph = ir.Graph([], [], nodes=[], name="vae") + mock_build_from_module.return_value = ModelPackage( + {"model": ir.Model(graph, ir_version=10)} + ) + mock_download_weights.return_value = {} + + build_diffusers_pipeline("fake/model", revision="pinned-revision") + + mock_load_index.assert_called_once_with("fake/model", revision="pinned-revision") + mock_load_config.assert_called_once_with( + "fake/model", "vae", revision="pinned-revision" + ) + mock_download_weights.assert_called_once_with( + "fake/model", "vae", revision="pinned-revision" + ) + mock_apply_weights.assert_called_once() + # ── build_diffusers_pipeline weight loading ────────────────────────────── @@ -438,7 +539,7 @@ def test_load_weights_true_downloads_and_applies( build_diffusers_pipeline("fake/model", load_weights=True) - mock_download_weights.assert_called_once_with("fake/model", "vae") + mock_download_weights.assert_called_once_with("fake/model", "vae", revision=None) mock_apply_weights.assert_called_once() @patch( diff --git a/src/mobius/_registry.py b/src/mobius/_registry.py index f6b704d3c..dd94cc2fc 100644 --- a/src/mobius/_registry.py +++ b/src/mobius/_registry.py @@ -33,6 +33,7 @@ MoonshineConfig, MuseGlimmerConfig, NemotronParseConfig, + ParakeetCTCConfig, WhisperConfig, ) from mobius.models import ( @@ -83,6 +84,7 @@ NemotronParseForConditionalGeneration, OLMo2CausalLMModel, OLMoCausalLMModel, + ParakeetForCTCModel, Phi3CausalLMModel, Phi3MoECausalLMModel, Phi3SmallCausalLMModel, @@ -819,6 +821,11 @@ def _detect_fallback_registration(hf_config) -> ModelRegistration | None: "wav2vec2-conformer": ModelRegistration(Wav2Vec2Model, task="audio-feature-extraction"), "wavlm": ModelRegistration(Wav2Vec2Model, task="audio-feature-extraction"), "mms": ModelRegistration(Wav2Vec2ForCTCModel, task="ctc-asr", config_class=MMSConfig), + "parakeet_ctc": ModelRegistration( + ParakeetForCTCModel, + task="feature-ctc-asr", + config_class=ParakeetCTCConfig, + ), "fastconformer_rnnt": ModelRegistration(EncDecRNNTModel, task="fastconformer-rnnt"), "sortformer": ModelRegistration(SortformerDiarizationModel, task="diarization"), } @@ -1065,6 +1072,7 @@ def _create_default_registry() -> ModelRegistry: "fun_asr": "justinchuby/Fun-ASR-Nano-2512", "sensevoice_small": "mlx-community/SenseVoiceSmall", "mms": "facebook/mms-300m", + "parakeet_ctc": "nvidia/parakeet-ctc-1.1b", "speecht5": "microsoft/speecht5_asr", "sew": "asapp/sew-tiny-100k", "sew-d": "asapp/sew-d-tiny-100k", diff --git a/src/mobius/_weight_loading.py b/src/mobius/_weight_loading.py index 33f1653c7..6ee95e366 100644 --- a/src/mobius/_weight_loading.py +++ b/src/mobius/_weight_loading.py @@ -148,7 +148,11 @@ def apply_weights(model: ir.Model, state_dict: dict[str, torch.Tensor]) -> None: def _parallel_download( - model_id: str, filenames: list[str], *, desc: str = "files" + model_id: str, + filenames: list[str], + *, + revision: str | None = None, + desc: str = "files", ) -> list[str]: """Download files from HuggingFace Hub in parallel. @@ -166,14 +170,19 @@ def _parallel_download( """ if len(filenames) <= 1: # No benefit from parallelism for a single file - return [hf_hub_download(repo_id=model_id, filename=f) for f in filenames] + kwargs = {"repo_id": model_id} + if revision is not None: + kwargs["revision"] = revision + return [hf_hub_download(filename=f, **kwargs) for f in filenames] print(f"Downloading {len(filenames)} {desc} files (parallel)...") path_map: dict[str, str] = {} with concurrent.futures.ThreadPoolExecutor(max_workers=8) as executor: + kwargs = {"repo_id": model_id} + if revision is not None: + kwargs["revision"] = revision futures = { - executor.submit(hf_hub_download, repo_id=model_id, filename=f): f - for f in filenames + executor.submit(hf_hub_download, filename=f, **kwargs): f for f in filenames } for future in tqdm.tqdm( concurrent.futures.as_completed(futures), @@ -295,7 +304,7 @@ def _dequantize_fp8_weights(state_dict: dict[str, torch.Tensor]) -> dict[str, to return {k: v for k, v in result.items() if not any(k.endswith(s) for s in aux_suffixes)} -def _download_weights(model_id: str) -> dict[str, torch.Tensor]: +def _download_weights(model_id: str, revision: str | None = None) -> dict[str, torch.Tensor]: """Download weights from HuggingFace and return as a state dict. Uses local safetensors files when *model_id* is a directory, otherwise @@ -305,17 +314,20 @@ def _download_weights(model_id: str) -> dict[str, torch.Tensor]: paths = _local_weight_paths(pathlib.Path(model_id)) if paths is None: try: - index_path = pathlib.Path( - hf_hub_download( - repo_id=model_id, - filename=_WEIGHT_INDEX_NAME, - ) - ) + kwargs = {"repo_id": model_id, "filename": _WEIGHT_INDEX_NAME} + if revision is not None: + kwargs["revision"] = revision + index_path = pathlib.Path(hf_hub_download(**kwargs)) all_files = _weight_filenames_from_index(index_path) except EntryNotFoundError: all_files = [_SINGLE_WEIGHT_NAME] - paths = _parallel_download(model_id, all_files, desc="safetensors") + paths = _parallel_download( + model_id, + all_files, + revision=revision, + desc="safetensors", + ) state_dict: dict[str, torch.Tensor] = {} for path in tqdm.tqdm(paths, desc="Loading weights"): diff --git a/src/mobius/components/__init__.py b/src/mobius/components/__init__.py index 308596f7e..c7594c615 100644 --- a/src/mobius/components/__init__.py +++ b/src/mobius/components/__init__.py @@ -5,6 +5,7 @@ "AdaLayerNormOutput", "AdaLayerNormZero", "Attention", + "BatchNorm1d", "BatchNorm2d", "BertEmbeddings", "BlockQuantizedLinear", @@ -57,6 +58,7 @@ "OffsetRMSNorm", "PatchEmbed", "PatchEmbedding", + "ParakeetFastConformerEncoder", "PostGatedRMSNorm", "PostNormDecoderLayer", "QuantizedEmbedding", @@ -135,6 +137,7 @@ create_static_cache_attention_bias, ) from mobius.components._conv import ( + BatchNorm1d, BatchNorm2d, CausalDepthwiseConv1d, Conv2d, @@ -197,6 +200,7 @@ MLPMultiModalProjector, ) from mobius.components._muse_glimmer_vision import MuseGlimmerVisionModel +from mobius.components._parakeet_audio import ParakeetFastConformerEncoder from mobius.components._pixtral_vision import ( Mistral3MultiModalProjector as Mistral3MultiModalProjector, ) diff --git a/src/mobius/components/_conv.py b/src/mobius/components/_conv.py index d421fcb12..0f11187e1 100644 --- a/src/mobius/components/_conv.py +++ b/src/mobius/components/_conv.py @@ -162,6 +162,28 @@ def forward(self, op: OpBuilder, x: ir.Value): ) +class BatchNorm1d(nn.Module): + """1D batch normalization with frozen running statistics.""" + + def __init__(self, num_features: int, eps: float = 1e-5): + super().__init__() + self.weight = nn.Parameter((num_features,)) + self.bias = nn.Parameter((num_features,)) + self.running_mean = nn.Parameter((num_features,)) + self.running_var = nn.Parameter((num_features,)) + self._eps = eps + + def forward(self, op: OpBuilder, x: ir.Value): + return op.BatchNormalization( + x, + self.weight, + self.bias, + self.running_mean, + self.running_var, + epsilon=self._eps, + ) + + class RmsNorm2d(nn.Module): """Channel-axis RMS normalization for NCHW tensors, scale-only. diff --git a/src/mobius/components/_conv_test.py b/src/mobius/components/_conv_test.py index de164458b..65e209646 100644 --- a/src/mobius/components/_conv_test.py +++ b/src/mobius/components/_conv_test.py @@ -14,6 +14,7 @@ ) from mobius.components._codec_conv import CausalConvNd from mobius.components._conv import ( + BatchNorm1d, BatchNorm2d, Conv2d, Conv2dNoBias, @@ -166,6 +167,21 @@ def test_forward_builds_graph(self): assert count_op_type(graph, "BatchNormalization") >= 1 +class TestBatchNorm1d: + """Tests for 1D batch normalization.""" + + def test_forward_uses_fused_op(self): + bn = BatchNorm1d(16) + builder, op, graph = create_test_builder() + x = create_test_input(builder, "x", [1, 16, 32]) + + result = bn(op, x) + builder._adapt_outputs([result], "") + + assert count_op_type(graph, "BatchNormalization") == 1 + assert count_op_type(graph, "Sqrt") == 0 + + class TestConvTranspose2d: """Tests for transposed 2D convolution.""" diff --git a/src/mobius/components/_parakeet_audio.py b/src/mobius/components/_parakeet_audio.py new file mode 100644 index 000000000..c4afa8e21 --- /dev/null +++ b/src/mobius/components/_parakeet_audio.py @@ -0,0 +1,522 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Offline FastConformer components used by Hugging Face Parakeet encoders.""" + +from __future__ import annotations + +import math +from typing import TYPE_CHECKING + +import numpy as np +import onnx_ir as ir +from onnxscript import OpBuilder, nn + +from mobius.components._activations import get_activation +from mobius.components._common import LayerNorm, Linear +from mobius.components._conv import BatchNorm1d, Conv2d +from mobius.components._whisper import Conv1d + +if TYPE_CHECKING: + from mobius._configs import ParakeetCTCConfig + + +def _dim(op: OpBuilder, value: ir.Value, axis: int) -> ir.Value: + return op.Shape(value, start=axis, end=axis + 1) + + +def _scalar_like(op: OpBuilder, value: float, reference: ir.Value) -> ir.Value: + scalar = op.Constant(value=ir.tensor(np.float32(value))) + return op.CastLike(scalar, reference) + + +class _ReLU(nn.Module): + def forward(self, op: OpBuilder, value: ir.Value) -> ir.Value: + return op.Relu(value) + + +class _ParakeetSubsampling(nn.Module): + """Symmetric depthwise Conv2d subsampling from mel frames to encoder states.""" + + def __init__(self, config: ParakeetCTCConfig): + super().__init__() + kernel = config.subsampling_conv_kernel_size + stride = config.subsampling_conv_stride + padding = (kernel - 1) // 2 + channels = config.subsampling_conv_channels + num_layers = int(math.log2(config.subsampling_factor)) + + layers: list[nn.Module] = [ + Conv2d( + 1, + channels, + kernel_size=kernel, + stride=stride, + padding=padding, + ), + _ReLU(), + ] + for _ in range(num_layers - 1): + layers.extend( + [ + Conv2d( + channels, + channels, + kernel_size=kernel, + stride=stride, + padding=padding, + groups=channels, + ), + Conv2d(channels, channels, kernel_size=1), + _ReLU(), + ] + ) + self.layers = nn.ModuleList(layers) + self._conv_indices = tuple( + index for index, layer in enumerate(layers) if isinstance(layer, Conv2d) + ) + self._strided_conv_indices = tuple( + index + for index, layer in enumerate(layers) + if isinstance(layer, Conv2d) and layer._stride != 1 + ) + + output_frequency = config.num_mel_bins // config.subsampling_factor + self.linear = Linear(channels * output_frequency, config.hidden_size, bias=True) + self._kernel = kernel + self._stride = stride + self._padding = padding + + def _downsample_lengths(self, op: OpBuilder, lengths: ir.Value) -> ir.Value: + numerator = op.Add( + lengths, + op.Constant(value=ir.tensor(np.int64(2 * self._padding - self._kernel))), + ) + return op.Add( + op.Div(numerator, op.Constant(value=ir.tensor(np.int64(self._stride)))), + op.Constant(value=ir.tensor(np.int64(1))), + ) + + def _mask_conv_output( + self, + op: OpBuilder, + hidden_states: ir.Value, + lengths: ir.Value, + ) -> ir.Value: + frame_ids = op.Range( + op.Constant(value=ir.tensor(np.int64(0))), + op.Squeeze(_dim(op, hidden_states, 2)), + op.Constant(value=ir.tensor(np.int64(1))), + ) + valid = op.Less( + op.Unsqueeze(frame_ids, op.Constant(value_ints=[0])), + op.Unsqueeze(lengths, op.Constant(value_ints=[1])), + ) + valid = op.Unsqueeze(valid, op.Constant(value_ints=[1, 3])) + return op.Where(valid, hidden_states, _scalar_like(op, 0.0, hidden_states)) + + def forward( + self, + op: OpBuilder, + input_features: ir.Value, + attention_mask: ir.Value, + ) -> ir.Value: + # (B, T, mel) -> (B, 1, T, mel) + hidden_states = op.Unsqueeze(input_features, op.Constant(value_ints=[1])) + lengths = op.ReduceSum( + op.Cast(attention_mask, to=ir.DataType.INT64), + axes=[1], + keepdims=0, + ) + + for index, layer in enumerate(self.layers): + hidden_states = layer(op, hidden_states) + if index in self._strided_conv_indices: + lengths = self._downsample_lengths(op, lengths) + if index in self._conv_indices: + hidden_states = self._mask_conv_output(op, hidden_states, lengths) + + # (B, C, T', F') -> (B, T', C*F') -> (B, T', hidden) + hidden_states = op.Transpose(hidden_states, perm=[0, 2, 1, 3]) + hidden_states = op.Reshape(hidden_states, [0, 0, -1]) + return self.linear(op, hidden_states) + + +class _ParakeetRelativePositionEncoding(nn.Module): + """Interleaved sinusoidal relative positions over ``[-T+1, T-1]``.""" + + def __init__(self, hidden_size: int, dtype: ir.DataType): + super().__init__() + self._hidden_size = hidden_size + self._dtype = dtype + self._inv_freq = 1.0 / ( + 10_000.0 ** (np.arange(0, hidden_size, 2, dtype=np.float32) / hidden_size) + ) + + def forward(self, op: OpBuilder, hidden_states: ir.Value) -> ir.Value: + seq_length = _dim(op, hidden_states, 1) + seq_float = op.Cast(op.Squeeze(seq_length), to=ir.DataType.FLOAT) + one = op.Constant(value=ir.tensor(np.float32(1.0))) + positions = op.Range(op.Sub(seq_float, one), op.Neg(seq_float), op.Neg(one)) + frequencies = op.Mul( + op.Unsqueeze(positions, op.Constant(value_ints=[1])), + op.Unsqueeze( + op.Constant(value=ir.tensor(self._inv_freq)), + op.Constant(value_ints=[0]), + ), + ) + sin = op.Unsqueeze(op.Sin(frequencies), op.Constant(value_ints=[-1])) + cos = op.Unsqueeze(op.Cos(frequencies), op.Constant(value_ints=[-1])) + positions = op.Reshape( + op.Concat(sin, cos, axis=-1), + op.Concat( + _dim(op, frequencies, 0), + op.Constant(value_ints=[self._hidden_size]), + axis=0, + ), + ) + positions = op.Unsqueeze(positions, op.Constant(value_ints=[0])) + if self._dtype != ir.DataType.FLOAT: + positions = op.Cast(positions, to=self._dtype) + return positions + + +class _ParakeetFeedForward(nn.Module): + """Macaron feed-forward projection with the configured activation.""" + + def __init__(self, config: ParakeetCTCConfig): + super().__init__() + self.linear1 = Linear( + config.hidden_size, + config.intermediate_size, + bias=config.attention_bias, + ) + self.linear2 = Linear( + config.intermediate_size, + config.hidden_size, + bias=config.attention_bias, + ) + self.act_fn = get_activation(config.hidden_act) + + def forward(self, op: OpBuilder, hidden_states: ir.Value) -> ir.Value: + return self.linear2(op, self.act_fn(op, self.linear1(op, hidden_states))) + + +class _ParakeetAttention(nn.Module): + """Transformer-XL relative-position self-attention used by Parakeet.""" + + def __init__(self, config: ParakeetCTCConfig): + super().__init__() + hidden_size = config.hidden_size + self._num_heads = config.num_attention_heads + self._head_dim = hidden_size // self._num_heads + self.q_proj = Linear(hidden_size, hidden_size, bias=config.attention_bias) + self.k_proj = Linear(hidden_size, hidden_size, bias=config.attention_bias) + self.v_proj = Linear(hidden_size, hidden_size, bias=config.attention_bias) + self.o_proj = Linear(hidden_size, hidden_size, bias=config.attention_bias) + self.relative_k_proj = Linear(hidden_size, hidden_size, bias=False) + self.bias_u = nn.Parameter([self._num_heads, self._head_dim]) + self.bias_v = nn.Parameter([self._num_heads, self._head_dim]) + + def _split_heads(self, op: OpBuilder, value: ir.Value) -> ir.Value: + shape = op.Concat( + _dim(op, value, 0), + _dim(op, value, 1), + op.Constant(value_ints=[self._num_heads, self._head_dim]), + axis=0, + ) + return op.Reshape(value, shape) + + def _relative_shift(self, op: OpBuilder, scores: ir.Value) -> ir.Value: + # (B, H, T, 2T-1) -> pad/reshape/shift -> (B, H, T, 2T-1) + zero_column = op.Expand( + _scalar_like(op, 0.0, scores), + op.Concat( + _dim(op, scores, 0), + _dim(op, scores, 1), + _dim(op, scores, 2), + op.Constant(value_ints=[1]), + axis=0, + ), + ) + scores = op.Concat(zero_column, scores, axis=-1) + scores = op.Reshape( + scores, + op.Concat( + _dim(op, scores, 0), + _dim(op, scores, 1), + op.Constant(value_ints=[-1]), + _dim(op, scores, 2), + axis=0, + ), + ) + scores = op.Slice( + scores, + op.Constant(value_ints=[1]), + _dim(op, scores, 2), + op.Constant(value_ints=[2]), + ) + return op.Reshape( + scores, + op.Concat( + _dim(op, scores, 0), + _dim(op, scores, 1), + _dim(op, scores, 3), + op.Constant(value_ints=[-1]), + axis=0, + ), + ) + + def forward( + self, + op: OpBuilder, + hidden_states: ir.Value, + position_embeddings: ir.Value, + attention_mask: ir.Value, + ) -> ir.Value: + scale = float(self._head_dim**-0.5) + query = self.q_proj(op, hidden_states) # (B, T, H*D) + key = self.k_proj(op, hidden_states) + value = self.v_proj(op, hidden_states) + + query_u = op.Add( + query, + op.Reshape( + self.bias_u, + op.Constant(value_ints=[self._num_heads * self._head_dim]), + ), + ) + query_v = self._split_heads(op, query) + query_v = op.Add( + query_v, + op.Reshape( + self.bias_v, + op.Constant(value_ints=[1, 1, self._num_heads, self._head_dim]), + ), + ) + query_v = op.Transpose(query_v, perm=[0, 2, 1, 3]) + + relative_key = self._split_heads(op, self.relative_k_proj(op, position_embeddings)) + relative_key = op.Transpose(relative_key, perm=[0, 2, 1, 3]) + relative_scores = op.MatMul( + query_v, + op.Transpose(relative_key, perm=[0, 1, 3, 2]), + ) + relative_scores = self._relative_shift(op, relative_scores) + relative_scores = op.Slice( + relative_scores, + op.Constant(value_ints=[0]), + _dim(op, key, 1), + op.Constant(value_ints=[3]), + ) + relative_scores = op.Mul(relative_scores, _scalar_like(op, scale, relative_scores)) + + additive_mask = op.Where( + attention_mask, + _scalar_like(op, 0.0, relative_scores), + _scalar_like(op, float("-inf"), relative_scores), + ) + attention_bias = op.Add(relative_scores, additive_mask) + output = op.Attention( + query_u, + key, + value, + attention_bias, + q_num_heads=self._num_heads, + kv_num_heads=self._num_heads, + scale=scale, + ) + return self.o_proj(op, output) + + +class _ParakeetConvolution(nn.Module): + """Bi-directional depthwise Conformer convolution with BatchNorm.""" + + def __init__(self, config: ParakeetCTCConfig): + super().__init__() + hidden_size = config.hidden_size + kernel_size = config.conv_kernel_size + self.pointwise_conv1 = Conv1d( + hidden_size, + 2 * hidden_size, + kernel_size=1, + bias=config.convolution_bias, + ) + self.depthwise_conv = Conv1d( + hidden_size, + hidden_size, + kernel_size=kernel_size, + padding=(kernel_size - 1) // 2, + groups=hidden_size, + bias=config.convolution_bias, + ) + self.norm = BatchNorm1d(hidden_size) + self.act_fn = get_activation(config.hidden_act) + self.pointwise_conv2 = Conv1d( + hidden_size, + hidden_size, + kernel_size=1, + bias=config.convolution_bias, + ) + + def forward( + self, + op: OpBuilder, + hidden_states: ir.Value, + valid_frames: ir.Value, + ) -> ir.Value: + # (B, T, C) -> pointwise GLU -> depthwise convolution in (B, C, T). + hidden_states = op.Transpose(hidden_states, perm=[0, 2, 1]) + hidden_states = self.pointwise_conv1(op, hidden_states) + first, gate = op.Split(hidden_states, axis=1, num_outputs=2, _outputs=2) + hidden_states = op.Mul(first, op.Sigmoid(gate)) + hidden_states = op.Where( + op.Unsqueeze(valid_frames, op.Constant(value_ints=[1])), + hidden_states, + _scalar_like(op, 0.0, hidden_states), + ) + hidden_states = self.depthwise_conv(op, hidden_states) + hidden_states = self.norm(op, hidden_states) + hidden_states = self.act_fn(op, hidden_states) + hidden_states = self.pointwise_conv2(op, hidden_states) + return op.Transpose(hidden_states, perm=[0, 2, 1]) + + +class _ParakeetEncoderLayer(nn.Module): + """Macaron FastConformer block with relative attention and convolution.""" + + def __init__(self, config: ParakeetCTCConfig): + super().__init__() + hidden_size = config.hidden_size + eps = config.layer_norm_eps + self.feed_forward1 = _ParakeetFeedForward(config) + self.self_attn = _ParakeetAttention(config) + self.conv = _ParakeetConvolution(config) + self.feed_forward2 = _ParakeetFeedForward(config) + self.norm_feed_forward1 = LayerNorm(hidden_size, eps=eps) + self.norm_self_att = LayerNorm(hidden_size, eps=eps) + self.norm_conv = LayerNorm(hidden_size, eps=eps) + self.norm_feed_forward2 = LayerNorm(hidden_size, eps=eps) + self.norm_out = LayerNorm(hidden_size, eps=eps) + + def forward( + self, + op: OpBuilder, + hidden_states: ir.Value, + position_embeddings: ir.Value, + attention_mask: ir.Value, + valid_frames: ir.Value, + ) -> ir.Value: + half = _scalar_like(op, 0.5, hidden_states) + hidden_states = op.Add( + hidden_states, + op.Mul( + self.feed_forward1(op, self.norm_feed_forward1(op, hidden_states)), + half, + ), + ) + hidden_states = op.Add( + hidden_states, + self.self_attn( + op, + self.norm_self_att(op, hidden_states), + position_embeddings, + attention_mask, + ), + ) + hidden_states = op.Add( + hidden_states, + self.conv(op, self.norm_conv(op, hidden_states), valid_frames), + ) + hidden_states = op.Add( + hidden_states, + op.Mul( + self.feed_forward2(op, self.norm_feed_forward2(op, hidden_states)), + half, + ), + ) + return self.norm_out(op, hidden_states) + + +class ParakeetFastConformerEncoder(nn.Module): + """Hugging Face Parakeet offline FastConformer audio encoder. + + Inputs are normalized log-mel features ``(batch, frames, mel_bins)`` and a + boolean valid-frame mask. The output is ``(batch, ceil(frames / 8), hidden)``. + """ + + def __init__(self, config: ParakeetCTCConfig): + super().__init__() + self.subsampling = _ParakeetSubsampling(config) + self.encode_positions = _ParakeetRelativePositionEncoding( + config.hidden_size, config.dtype + ) + self.layers = nn.ModuleList( + [_ParakeetEncoderLayer(config) for _ in range(config.num_hidden_layers)] + ) + self._num_subsampling_layers = int(math.log2(config.subsampling_factor)) + self._subsampling_kernel = config.subsampling_conv_kernel_size + self._subsampling_stride = config.subsampling_conv_stride + self._subsampling_padding = (config.subsampling_conv_kernel_size - 1) // 2 + self._input_scale = math.sqrt(config.hidden_size) if config.scale_input else 1.0 + + def _output_mask( + self, + op: OpBuilder, + attention_mask: ir.Value, + target_length: ir.Value, + ) -> ir.Value: + lengths = op.ReduceSum( + op.Cast(attention_mask, to=ir.DataType.INT64), + axes=[1], + keepdims=0, + ) + add_pad = 2 * self._subsampling_padding - self._subsampling_kernel + stride = op.Constant(value=ir.tensor(np.int64(self._subsampling_stride))) + for _ in range(self._num_subsampling_layers): + lengths = op.Add( + op.Div( + op.Add( + lengths, + op.Constant(value=ir.tensor(np.int64(add_pad))), + ), + stride, + ), + op.Constant(value=ir.tensor(np.int64(1))), + ) + frame_ids = op.Range( + op.Constant(value=ir.tensor(np.int64(0))), + op.Squeeze(target_length), + op.Constant(value=ir.tensor(np.int64(1))), + ) + return op.Less( + op.Unsqueeze(frame_ids, op.Constant(value_ints=[0])), + op.Unsqueeze(lengths, op.Constant(value_ints=[1])), + ) + + def forward( + self, + op: OpBuilder, + input_features: ir.Value, + attention_mask: ir.Value, + ) -> tuple[ir.Value, ir.Value]: + hidden_states = self.subsampling(op, input_features, attention_mask) + hidden_states = op.Mul( + hidden_states, _scalar_like(op, self._input_scale, hidden_states) + ) + position_embeddings = self.encode_positions(op, hidden_states) + output_mask = self._output_mask(op, attention_mask, _dim(op, hidden_states, 1)) + attention_mask_4d = op.And( + op.Unsqueeze(output_mask, op.Constant(value_ints=[1, 2])), + op.Unsqueeze(output_mask, op.Constant(value_ints=[1, 3])), + ) + + for layer in self.layers: + hidden_states = layer( + op, + hidden_states, + position_embeddings, + attention_mask_4d, + output_mask, + ) + return hidden_states, output_mask diff --git a/src/mobius/components/_whisper.py b/src/mobius/components/_whisper.py index 005b719bf..4c058d9ac 100644 --- a/src/mobius/components/_whisper.py +++ b/src/mobius/components/_whisper.py @@ -33,13 +33,15 @@ def __init__( stride: int = 1, padding: int = 0, bias: bool = True, + groups: int = 1, ): super().__init__() - self.weight = nn.Parameter([out_channels, in_channels, kernel_size]) + self.weight = nn.Parameter([out_channels, in_channels // groups, kernel_size]) self.bias = nn.Parameter([out_channels]) if bias else None self._kernel_shape = [kernel_size] self._strides = [stride] self._pads = [padding, padding] + self._groups = groups def forward(self, op: OpBuilder, x: ir.Value): # x: [batch, in_channels, seq_len] @@ -50,6 +52,7 @@ def forward(self, op: OpBuilder, x: ir.Value): kernel_shape=self._kernel_shape, strides=self._strides, pads=self._pads, + group=self._groups, ) diff --git a/src/mobius/components/_whisper_test.py b/src/mobius/components/_whisper_test.py index 72a5de43f..b8f6d5284 100644 --- a/src/mobius/components/_whisper_test.py +++ b/src/mobius/components/_whisper_test.py @@ -34,6 +34,13 @@ def test_parameter_count(self): params = list(conv.parameters()) assert len(params) == 2 # weight + bias + def test_positional_bias_and_grouped_weight_shape(self): + biasless = Conv1d(4, 4, 3, 1, 1, False) + grouped = Conv1d(4, 4, 3, groups=4) + + assert biasless.bias is None + assert list(grouped.weight.shape) == [4, 1, 3] + def test_forward_builds_graph(self): conv = Conv1d(80, 512, kernel_size=3, padding=1) builder, op, graph = create_test_builder() diff --git a/src/mobius/integrations/ort_genai/auto_export.py b/src/mobius/integrations/ort_genai/auto_export.py index 500815138..8f1741630 100644 --- a/src/mobius/integrations/ort_genai/auto_export.py +++ b/src/mobius/integrations/ort_genai/auto_export.py @@ -1190,6 +1190,11 @@ def write_ort_genai_config( "This is set automatically when building with mobius.build(). " "Diffusion models (which have no config) are not supported." ) + if config.model_type == "parakeet_ctc": + raise ValueError( + "ORT GenAI does not define a feature-input CTC ASR pipeline; " + "export Parakeet CTC as ONNX and run it directly with ONNX Runtime." + ) if {"vision_encoder", "decoder"}.issubset(pkg) and "embedding" not in pkg: model_type = getattr(config, "model_type", "unknown") raise NotImplementedError( diff --git a/src/mobius/models/__init__.py b/src/mobius/models/__init__.py index 431cc8367..ded44e522 100644 --- a/src/mobius/models/__init__.py +++ b/src/mobius/models/__init__.py @@ -96,6 +96,7 @@ "OLMo2CausalLMModel", "OLMoCausalLMModel", "OPTCausalLMModel", + "ParakeetForCTCModel", "PersimmonCausalLMModel", "Phi3CausalLMModel", "Phi3MoECausalLMModel", @@ -259,6 +260,7 @@ from mobius.models.nemotron_parse import NemotronParseForConditionalGeneration from mobius.models.olmo import OLMo2CausalLMModel, OLMoCausalLMModel from mobius.models.opt import OPTCausalLMModel +from mobius.models.parakeet_ctc import ParakeetForCTCModel from mobius.models.persimmon import PersimmonCausalLMModel from mobius.models.phi import ( Phi3SmallCausalLMModel, diff --git a/src/mobius/models/parakeet_ctc.py b/src/mobius/models/parakeet_ctc.py new file mode 100644 index 000000000..afbbabed2 --- /dev/null +++ b/src/mobius/models/parakeet_ctc.py @@ -0,0 +1,55 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Hugging Face Parakeet FastConformer model for CTC speech recognition.""" + +from __future__ import annotations + +import onnx_ir as ir +from onnxscript import OpBuilder, nn + +from mobius._configs import ParakeetCTCConfig +from mobius.components import Conv1d, ParakeetFastConformerEncoder + + +class ParakeetForCTCModel(nn.Module): + """Replicate ``transformers.ParakeetForCTC`` as a single ONNX graph. + + Inputs are normalized log-mel features and their valid-frame mask. The + FastConformer encoder downsamples time by 8 and the CTC head returns + per-frame vocabulary logits. + """ + + default_task = "feature-ctc-asr" + category = "Speech-to-Text" + config_class = ParakeetCTCConfig + + def __init__(self, config: ParakeetCTCConfig): + super().__init__() + if config.dtype == ir.DataType.BFLOAT16: + raise ValueError( + "Parakeet CTC bf16 is disabled because ONNX Runtime produces " + "incorrect CTC logits; use dtype='f16' or dtype='f32'." + ) + self.encoder = ParakeetFastConformerEncoder(config) + self.ctc_head = Conv1d(config.hidden_size, config.vocab_size, kernel_size=1) + + def forward( + self, + op: OpBuilder, + input_features: ir.Value, + attention_mask: ir.Value, + ) -> ir.Value: + hidden_states, _ = self.encoder(op, input_features, attention_mask) + # The checkpoint stores the CTC projection as Conv1d (vocab, hidden, 1). + hidden_states = op.Transpose(hidden_states, perm=[0, 2, 1]) + logits = self.ctc_head(op, hidden_states) + return op.Transpose(logits, perm=[0, 2, 1]) + + def preprocess_weights(self, state_dict: dict[str, object]) -> dict[str, object]: + """Drop PyTorch-only BatchNorm counters; all tensor names align directly.""" + return { + name: value + for name, value in state_dict.items() + if not name.endswith(".num_batches_tracked") + } diff --git a/src/mobius/models/parakeet_ctc_test.py b/src/mobius/models/parakeet_ctc_test.py new file mode 100644 index 000000000..929a68a5b --- /dev/null +++ b/src/mobius/models/parakeet_ctc_test.py @@ -0,0 +1,178 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +from __future__ import annotations + +import dataclasses + +import numpy as np +import onnx_ir as ir +import pytest +import torch +from transformers import ( + ParakeetCTCConfig as HFParakeetCTCConfig, +) +from transformers import ( + ParakeetEncoderConfig as HFParakeetEncoderConfig, +) +from transformers import ( + ParakeetForCTC as HFParakeetForCTC, +) + +from mobius import build_from_module +from mobius._configs import ParakeetCTCConfig +from mobius._testing.ort_inference import OnnxModelSession +from mobius._weight_loading import apply_weights +from mobius.integrations.ort_genai import write_ort_genai_config +from mobius.models import ParakeetForCTCModel +from mobius.tasks import FeatureCTCAsrTask + + +def _hf_config() -> HFParakeetCTCConfig: + encoder = HFParakeetEncoderConfig( + hidden_size=16, + intermediate_size=32, + num_hidden_layers=1, + num_attention_heads=2, + num_key_value_heads=2, + num_mel_bins=8, + subsampling_conv_channels=4, + conv_kernel_size=5, + dropout=0.0, + dropout_positions=0.0, + layerdrop=0.0, + activation_dropout=0.0, + attention_dropout=0.0, + max_position_embeddings=64, + ) + return HFParakeetCTCConfig( + encoder_config=encoder, + vocab_size=17, + pad_token_id=16, + ) + + +def _build_tiny(): + hf_config = _hf_config() + config = ParakeetCTCConfig.from_transformers(hf_config) + config.dtype = ir.DataType.FLOAT + module = ParakeetForCTCModel(config) + package = build_from_module(module, config, task=FeatureCTCAsrTask()) + return hf_config, config, module, package + + +def test_parakeet_config_extracts_nested_encoder_fields(): + hf_config = _hf_config() + hf_config.dtype = torch.bfloat16 + config = ParakeetCTCConfig.from_transformers(hf_config) + + assert config.model_type == "parakeet_ctc" + assert config.vocab_size == 17 + assert config.pad_token_id == 16 + assert config.hidden_size == 16 + assert config.num_mel_bins == 8 + assert config.subsampling_factor == 8 + assert config.conv_kernel_size == 5 + assert config.dtype == ir.DataType.FLOAT + + +def test_parakeet_rejects_bfloat16(): + config = ParakeetCTCConfig.from_transformers(_hf_config()) + config.dtype = ir.DataType.BFLOAT16 + + with pytest.raises(ValueError, match="bf16 is disabled"): + ParakeetForCTCModel(config) + + +def test_parakeet_graph_io_and_hf_weight_names_align(): + hf_config, _, module, package = _build_tiny() + hf_model = HFParakeetForCTC(hf_config) + processed = module.preprocess_weights(dict(hf_model.state_dict())) + model = package["model"] + + assert [value.name for value in model.graph.inputs] == [ + "input_features", + "attention_mask", + ] + assert [value.name for value in model.graph.outputs] == ["logits"] + parameter_names = { + name + for name, initializer in model.graph.initializers.items() + if initializer.const_value is None + } + assert set(processed) == parameter_names + assert not any(name.endswith(".num_batches_tracked") for name in processed) + + +def test_parakeet_graph_uses_fused_encoder_ops(): + _, config, _, package = _build_tiny() + model = package["model"] + op_types = [node.op_type for node in model.graph.all_nodes()] + + assert op_types.count("Attention") == config.num_hidden_layers + assert op_types.count("BatchNormalization") == config.num_hidden_layers + assert op_types.count("Swish") == 3 * config.num_hidden_layers + assert op_types.count("SkipLayerNormalization") == 4 * config.num_hidden_layers + assert "Sqrt" not in op_types + + +def test_parakeet_honors_non_silu_activation_in_ffn_and_convolution(): + hf_config = _hf_config() + config = dataclasses.replace( + ParakeetCTCConfig.from_transformers(hf_config), + hidden_act="relu", + dtype=ir.DataType.FLOAT, + ) + module = ParakeetForCTCModel(config) + model = build_from_module(module, config, task=FeatureCTCAsrTask())["model"] + + op_types = [node.op_type for node in model.graph] + # The three configurable activations are the two Macaron FFNs and the + # convolution module. Subsampling contributes three fixed ReLUs. + assert op_types.count("Relu") == 6 + assert "Swish" not in op_types + + +def test_parakeet_synthetic_parity_with_padding(): + torch.manual_seed(42) + hf_config, _, module, package = _build_tiny() + hf_model = HFParakeetForCTC(hf_config).float().eval() + apply_weights( + package["model"], + module.preprocess_weights(dict(hf_model.state_dict())), + ) + + rng = np.random.default_rng(123) + input_features = rng.standard_normal((2, 24, 8)).astype(np.float32) + attention_mask = np.ones((2, 24), dtype=bool) + attention_mask[1, 17:] = False + input_features[1, 17:] = 0.0 + + with torch.no_grad(): + expected = hf_model( + input_features=torch.from_numpy(input_features), + attention_mask=torch.from_numpy(attention_mask), + ).logits.numpy() + + session = OnnxModelSession(package["model"], device="cpu") + try: + actual = session.run( + { + "input_features": input_features, + "attention_mask": attention_mask, + } + )["logits"] + finally: + session.close() + + np.testing.assert_allclose(actual, expected, atol=1e-5, rtol=1e-5) + + +def test_parakeet_rejects_unsupported_ort_genai_export(tmp_path): + _, _, _, package = _build_tiny() + + with pytest.raises( + ValueError, + match="does not define a feature-input CTC ASR pipeline", + ): + write_ort_genai_config(package, str(tmp_path)) diff --git a/src/mobius/tasks/__init__.py b/src/mobius/tasks/__init__.py index 53626674d..145658f3a 100644 --- a/src/mobius/tasks/__init__.py +++ b/src/mobius/tasks/__init__.py @@ -24,6 +24,7 @@ "AudioFeatureExtractionTask", "CausalLMTask", "CTCAsrTask", + "FeatureCTCAsrTask", "RNNTTask", "CodecTask", "ComponentSpec", @@ -92,7 +93,7 @@ ) from mobius.tasks._codec import CodecTask from mobius.tasks._controlnet import ControlNetTask -from mobius.tasks._ctc_asr import CTCAsrTask +from mobius.tasks._ctc_asr import CTCAsrTask, FeatureCTCAsrTask from mobius.tasks._deepseek_v4 import DeepSeekV4Task from mobius.tasks._denoising import DenoisingTask from mobius.tasks._dflash import DFlashDraftTask @@ -146,6 +147,7 @@ "audio-ctc": AudioCTCTask, "audio-feature-extraction": AudioFeatureExtractionTask, "ctc-asr": CTCAsrTask, + "feature-ctc-asr": FeatureCTCAsrTask, "codec": CodecTask, "controlnet": ControlNetTask, "denoising": DenoisingTask, diff --git a/src/mobius/tasks/_ctc_asr.py b/src/mobius/tasks/_ctc_asr.py index a45fd2773..87d756b2d 100644 --- a/src/mobius/tasks/_ctc_asr.py +++ b/src/mobius/tasks/_ctc_asr.py @@ -61,3 +61,46 @@ def build( builder.add_output(logits, "logits") return ModelPackage({"model": _make_model(graph)}, config=config) + + +class FeatureCTCAsrTask(ModelTask): + """Build feature-input CTC ASR (log-mel features → frame logits). + + Inputs: + ``input_features`` — (batch, frames, mel_bins) normalized log-mel values + ``attention_mask`` — (batch, frames) BOOL valid-frame mask + + Output: + ``logits`` — (batch, subsampled_frames, vocab_size) CTC scores + """ + + name = "feature-ctc-asr" + model_roles: ClassVar[dict[str, str]] = {"model": "encoder"} + + def build( + self, + module, + config: ArchitectureConfig, + ) -> ModelPackage: + batch = ir.SymbolicDim("batch") + frames = ir.SymbolicDim("frames") + + graph, builder = _make_graph(name="feature_ctc_asr") + input_features = builder.input( + "input_features", + dtype=config.dtype, + shape=[batch, frames, config.num_mel_bins], + ) + attention_mask = builder.input( + "attention_mask", + dtype=ir.DataType.BOOL, + shape=[batch, frames], + ) + + logits = module( + builder.op, + input_features=input_features, + attention_mask=attention_mask, + ) + builder.add_output(logits, "logits") + return ModelPackage({"model": _make_model(graph)}, config=config) diff --git a/testdata/cases/audio/parakeet-ctc-1.1b.yaml b/testdata/cases/audio/parakeet-ctc-1.1b.yaml new file mode 100644 index 000000000..563fd8dca --- /dev/null +++ b/testdata/cases/audio/parakeet-ctc-1.1b.yaml @@ -0,0 +1,18 @@ +model_id: "nvidia/parakeet-ctc-1.1b" +model_type: "parakeet_ctc" +revision: "20e63a0fed6aedba145b74b826dbd41df0941730" +task_type: "feature-ctc-asr" +dtype: "float32" + +inputs: + audio: + - "652-129742-0006.flac" + +level: "L4+L5" + +notes: > + NVIDIA Parakeet 1.1B offline FastConformer CTC ASR. The Hugging Face + processor converts real 16 kHz LibriSpeech audio to normalized 80-bin + log-mel features; L4 compares the final valid frame logits and L5 compares + every deterministic CTC frame argmax before collapse, with the decoded + transcript retained in the generation golden for human inspection. diff --git a/testdata/cases/schema.json b/testdata/cases/schema.json index aea8a6fbb..c6f3c3c17 100644 --- a/testdata/cases/schema.json +++ b/testdata/cases/schema.json @@ -48,6 +48,7 @@ "depth-estimation", "dflash-draft", "feature-extraction", + "feature-ctc-asr", "fun-asr-speech-language", "gemma4-assistant", "image-classification", diff --git a/testdata/golden/audio/parakeet-ctc-1.1b.json b/testdata/golden/audio/parakeet-ctc-1.1b.json new file mode 100644 index 000000000..fb4da6126 --- /dev/null +++ b/testdata/golden/audio/parakeet-ctc-1.1b.json @@ -0,0 +1,37 @@ +{ + "top1_id": 1024, + "top2_id": 25, + "top10_ids": [ + 1024, + 25, + 1003, + 3, + 5, + 997, + 45, + 22, + 998, + 4 + ], + "top10_logits": [ + "0x1.4f1f600000000p+6", + "0x1.1e17e80000000p+6", + "0x1.1cee4a0000000p+6", + "0x1.1bd21c0000000p+6", + "0x1.1952000000000p+6", + "0x1.19516e0000000p+6", + "0x1.1834e80000000p+6", + "0x1.17b0440000000p+6", + "0x1.16dfe80000000p+6", + "0x1.16b8e80000000p+6" + ], + "logits_summary": [ + "0x1.4f1f600000000p+6", + "-0x1.8503ae0000000p+2", + "0x1.ff7a2df811fb8p+5", + "0x1.c60f78f1a14d3p+1" + ], + "input_ids": [ + 0 + ] +} diff --git a/testdata/golden/audio/parakeet-ctc-1.1b_generation.json b/testdata/golden/audio/parakeet-ctc-1.1b_generation.json new file mode 100644 index 000000000..cbad02c4d --- /dev/null +++ b/testdata/golden/audio/parakeet-ctc-1.1b_generation.json @@ -0,0 +1,121 @@ +{ + "model_id": "nvidia/parakeet-ctc-1.1b", + "prompt": "testdata\\652-129742-0006.flac", + "generated_tokens": [ + 1024, + 1024, + 1024, + 1024, + 1024, + 1024, + 658, + 148, + 1024, + 159, + 1024, + 1024, + 1006, + 1024, + 632, + 1024, + 377, + 1024, + 1024, + 636, + 1024, + 1024, + 1024, + 530, + 1024, + 1024, + 1024, + 1024, + 1024, + 1024, + 1024, + 1024, + 1024, + 1024, + 1024, + 432, + 1024, + 1024, + 1024, + 1024, + 246, + 76, + 1024, + 1024, + 1024, + 470, + 1024, + 128, + 33, + 1024, + 658, + 1024, + 148, + 1024, + 159, + 1024, + 1006, + 1024, + 632, + 1024, + 1024, + 1024, + 851, + 450, + 1024, + 1024, + 308, + 1024, + 1024, + 369, + 29, + 85, + 1024, + 1024, + 27, + 1024, + 1024, + 505, + 1024, + 21, + 1024, + 1024, + 1024, + 8, + 53, + 998, + 1024, + 1024, + 1024, + 1024, + 1024, + 216, + 151, + 1024, + 12, + 1024, + 25, + 1024, + 1024, + 150, + 263, + 1024, + 1013, + 40, + 1024, + 22, + 1024, + 101, + 1024, + 667, + 1024, + 1024, + 1024, + 1024 + ], + "generated_text": "cauliflower mayonnaise take cold boiled cauliflower break into branches adding salt pepper and vinegar to season" +} diff --git a/tests/_test_configs.py b/tests/_test_configs.py index bc4d3aa9a..751050de8 100644 --- a/tests/_test_configs.py +++ b/tests/_test_configs.py @@ -42,6 +42,7 @@ NanoChatConfig, NemotronHConfig, NemotronParseConfig, + ParakeetCTCConfig, Sam2Config, SegformerConfig, VisionConfig, @@ -2617,6 +2618,20 @@ def vl_overrides(model_type: str) -> dict: # Speech / TTS / Codec configs # --------------------------------------------------------------------------- SPEECH_CONFIGS: list[tuple[str, dict, bool]] = [ + # --- Parakeet CTC (feature-input offline FastConformer) --- + ( + "parakeet_ctc", + { + "_config_cls": ParakeetCTCConfig, + "num_mel_bins": 16, + "subsampling_conv_channels": 8, + "conv_kernel_size": 5, + "attention_bias": True, + "convolution_bias": True, + "scale_input": True, + }, + True, + ), # --- Moonshine (raw-waveform RoPE encoder-decoder ASR) --- ( "moonshine", diff --git a/tests/build_graph_test.py b/tests/build_graph_test.py index 33a05b131..deab80ba4 100644 --- a/tests/build_graph_test.py +++ b/tests/build_graph_test.py @@ -6060,6 +6060,7 @@ def test_outputs_have_shapes_and_dtypes(self, model_type: str, config_overrides: "speech-language": {"audio_encoder", "embedding", "decoder"}, "codec": {"decoder", "encoder"}, "audio-feature-extraction": {"model"}, + "feature-ctc-asr": {"model"}, } diff --git a/tests/e2e_golden_test.py b/tests/e2e_golden_test.py index a253f6575..ad380ad2f 100644 --- a/tests/e2e_golden_test.py +++ b/tests/e2e_golden_test.py @@ -24,13 +24,16 @@ from __future__ import annotations +import ast import dataclasses import functools import os +import shutil import warnings from pathlib import Path from unittest import mock +import huggingface_hub.constants as hf_constants import numpy as np import pytest @@ -107,7 +110,11 @@ def _get_test_device_kwargs() -> dict[str, str]: _IN_CI = os.environ.get("GITHUB_ACTIONS") == "true" -def _load_suppress_token_ids(model_id: str, trust_remote_code: bool = False) -> list[int]: +def _load_suppress_token_ids( + model_id: str, + revision: str | None = None, + trust_remote_code: bool = False, +) -> list[int]: """Return ``generation_config.suppress_tokens`` for a model (empty if none). Mirrors HuggingFace ``generate()``: tokens in ``suppress_tokens`` are forced @@ -122,7 +129,9 @@ def _load_suppress_token_ids(model_id: str, trust_remote_code: bool = False) -> try: gen_config = transformers.GenerationConfig.from_pretrained( - model_id, trust_remote_code=trust_remote_code + model_id, + revision=revision, + trust_remote_code=trust_remote_code, ) except Exception: return [] @@ -279,24 +288,37 @@ def _make_empty_kv_cache( @pytest.fixture(autouse=True) -def _use_temp_hf_cache(tmp_path): +def _use_temp_hf_cache(tmp_path, monkeypatch): """Redirect HuggingFace downloads to a per-test temp dir. - Each test gets a fresh cache that is deleted when the test finishes, - so only one model's weights are on disk at a time. This prevents - unbounded disk growth across the full test suite. + Hugging Face resolves cache constants when its module is imported, so + changing only ``HF_HOME`` is too late for this test module. Patch the + runtime constants as well, including the Xet chunk cache used by large + checkpoints, and eagerly delete the cache after each test. This keeps + all-model GPU golden runs within the hosted runner's disk limit. Each pytest-xdist worker gets its own ``tmp_path``, so parallel workers don't collide. """ - cache_dir = str(tmp_path / "hf_cache") - old = os.environ.get("HF_HOME") - os.environ["HF_HOME"] = cache_dir - yield - if old is None: - os.environ.pop("HF_HOME", None) - else: - os.environ["HF_HOME"] = old + cache_root = tmp_path / "hf_cache" + hub_cache = cache_root / "hub" + assets_cache = cache_root / "assets" + xet_cache = cache_root / "xet" + + monkeypatch.setenv("HF_HOME", str(cache_root)) + monkeypatch.setenv("HF_HUB_CACHE", str(hub_cache)) + monkeypatch.setenv("HF_ASSETS_CACHE", str(assets_cache)) + monkeypatch.setenv("HF_XET_CACHE", str(xet_cache)) + monkeypatch.setattr(hf_constants, "HF_HOME", str(cache_root)) + monkeypatch.setattr(hf_constants, "HF_HUB_CACHE", str(hub_cache)) + monkeypatch.setattr(hf_constants, "HF_ASSETS_CACHE", str(assets_cache)) + monkeypatch.setattr(hf_constants, "HF_XET_CACHE", str(xet_cache)) + + try: + yield + finally: + if cache_root.exists(): + shutil.rmtree(cache_root) # --------------------------------------------------------------------------- @@ -393,6 +415,23 @@ def _discover_cases( # --------------------------------------------------------------------------- +def test_huggingface_artifact_loads_are_revision_pinned(): + """Every Hub-backed test artifact must resolve from the case revision.""" + tree = ast.parse(Path(__file__).read_text()) + unpinned_lines = [ + node.lineno + for node in ast.walk(tree) + if isinstance(node, ast.Call) + and isinstance(node.func, ast.Attribute) + and node.func.attr == "from_pretrained" + and not any(keyword.arg == "revision" for keyword in node.keywords) + ] + + assert not unpinned_lines, ( + f"from_pretrained calls missing revision= at lines {unpinned_lines}" + ) + + def _build_model_package(case: GoldenTestCase) -> ModelPackage: """Build an ONNX ModelPackage with real weights from HuggingFace.""" module_class = None @@ -408,6 +447,7 @@ def _build_model_package(case: GoldenTestCase) -> ModelPackage: task = reg.task or getattr(module_class, "default_task", None) return build( case.model_id, + revision=case.revision, task=task, module_class=module_class, dtype=case.dtype, @@ -615,6 +655,7 @@ def _prepare_prefill_feeds( # comparator below slices to the last frame so the shape matches # the saved golden's per-token vector. "ctc-asr", + "feature-ctc-asr", } ) @@ -681,7 +722,9 @@ def _prepare_vision_feeds( from PIL import Image processor = transformers.AutoImageProcessor.from_pretrained( - case.model_id, trust_remote_code=case.trust_remote_code + case.model_id, + revision=case.revision, + trust_remote_code=case.trust_remote_code, ) image = Image.open(_TESTDATA_DIR / case.images[0]) proc_kwargs: dict = {"images": image, "return_tensors": "np"} @@ -704,7 +747,9 @@ def _detection_forced_size(case: GoldenTestCase) -> dict | None: import transformers config = transformers.AutoConfig.from_pretrained( - case.model_id, trust_remote_code=case.trust_remote_code + case.model_id, + revision=case.revision, + trust_remote_code=case.trust_remote_code, ) image_size = getattr(config, "image_size", None) if isinstance(image_size, (list, tuple)) and len(image_size) == 2: @@ -728,11 +773,15 @@ def _prepare_audio_feeds( # Fall back to AutoFeatureExtractor for models without a tokenizer try: processor = transformers.AutoProcessor.from_pretrained( - case.model_id, trust_remote_code=case.trust_remote_code + case.model_id, + revision=case.revision, + trust_remote_code=case.trust_remote_code, ) except (TypeError, OSError): processor = transformers.AutoFeatureExtractor.from_pretrained( - case.model_id, trust_remote_code=case.trust_remote_code + case.model_id, + revision=case.revision, + trust_remote_code=case.trust_remote_code, ) audio_path = _TESTDATA_DIR / case.audio[0] audio_array, _sr = librosa.load(str(audio_path), sr=16000) @@ -745,6 +794,31 @@ def _prepare_audio_feeds( return feeds +def _prepare_feature_ctc_feeds( + case: GoldenTestCase, +) -> dict[str, np.ndarray]: + """Prepare processor-generated log-mel features for feature-input CTC.""" + import librosa + import transformers + + processor = transformers.AutoProcessor.from_pretrained( + case.model_id, + revision=case.revision, + trust_remote_code=case.trust_remote_code, + ) + audio_path = _TESTDATA_DIR / case.audio[0] + audio_array, sample_rate = librosa.load(str(audio_path), sr=16000) + processed = processor( + audio_array, + sampling_rate=sample_rate, + return_tensors="np", + ) + return { + "input_features": processed["input_features"].astype(np.float32), + "attention_mask": processed["attention_mask"].astype(bool), + } + + def _compute_mrope_position_ids( input_ids: np.ndarray, image_grid_thw: np.ndarray, @@ -919,7 +993,9 @@ def _run_vision_language_prefill( # --- Step 0: Preprocess image with HF processor --- processor = transformers.AutoProcessor.from_pretrained( - case.model_id, trust_remote_code=case.trust_remote_code + case.model_id, + revision=case.revision, + trust_remote_code=case.trust_remote_code, ) image = Image.open(_TESTDATA_DIR / case.images[0]) @@ -1087,12 +1163,18 @@ def _run_vl_generation( # --- Step 0: prepare multimodal inputs --- processor = transformers.AutoProcessor.from_pretrained( - case.model_id, trust_remote_code=case.trust_remote_code + case.model_id, + revision=case.revision, + trust_remote_code=case.trust_remote_code, ) image = Image.open(_TESTDATA_DIR / case.images[0]) prompt_text = _build_mm_prompt(processor, case.prompts[0], case.images, "image") - suppress_ids = _load_suppress_token_ids(case.model_id, case.trust_remote_code) + suppress_ids = _load_suppress_token_ids( + case.model_id, + revision=case.revision, + trust_remote_code=case.trust_remote_code, + ) processed_pt = processor(text=prompt_text, images=[image], return_tensors="pt") processed: dict[str, np.ndarray] = { @@ -1285,7 +1367,9 @@ def _run_speech_to_text_prefill( # Load audio and extract features processor = transformers.AutoProcessor.from_pretrained( - case.model_id, trust_remote_code=case.trust_remote_code + case.model_id, + revision=case.revision, + trust_remote_code=case.trust_remote_code, ) audio_path = _TESTDATA_DIR / case.audio[0] audio_array, _sr = librosa.load(str(audio_path), sr=16000) @@ -1440,7 +1524,9 @@ def _run_phi4mm_multimodal_prefill( from PIL import Image processor = transformers.AutoProcessor.from_pretrained( - case.model_id, trust_remote_code=True + case.model_id, + revision=case.revision, + trust_remote_code=True, ) images = [Image.open(_TESTDATA_DIR / img_path) for img_path in case.images] img_inputs = processor.image_processor(images=images, return_tensors="np") @@ -1484,7 +1570,9 @@ def _run_phi4mm_multimodal_prefill( import librosa processor = transformers.AutoProcessor.from_pretrained( - case.model_id, trust_remote_code=True + case.model_id, + revision=case.revision, + trust_remote_code=True, ) audios = [] for audio_path in case.audio: @@ -1578,7 +1666,9 @@ def _run_speech_language_prefill( # Load audio and extract features processor = transformers.AutoProcessor.from_pretrained( - case.model_id, trust_remote_code=case.trust_remote_code + case.model_id, + revision=case.revision, + trust_remote_code=case.trust_remote_code, ) audio_path = _TESTDATA_DIR / case.audio[0] audio_array, _sr = librosa.load(str(audio_path), sr=16000) @@ -1590,7 +1680,10 @@ def _run_speech_language_prefill( # actually uses under the hood. fe = getattr(processor, "feature_extractor", None) if fe is None or not hasattr(fe, "sampling_rate"): - fe = transformers.WhisperFeatureExtractor.from_pretrained(case.model_id) + fe = transformers.WhisperFeatureExtractor.from_pretrained( + case.model_id, + revision=case.revision, + ) audio_processed = fe( [audio_array], sampling_rate=16000, @@ -1971,6 +2064,12 @@ def test_prefill_argmax_matches_golden(self, case: GoldenTestCase) -> None: outputs = session.run(feeds) finally: session.close() + elif case.task_type == "feature-ctc-asr": + session = _open_decoder_session(pkg) + try: + outputs = session.run(_prepare_feature_ctc_feeds(case)) + finally: + session.close() elif len(pkg) > 1 and "embedding" in pkg: # Multi-model text-generation (e.g. Gemma4 VL text-only) outputs = _run_text_only_multimodel_prefill(pkg, golden, config) @@ -2017,10 +2116,30 @@ def test_prefill_argmax_matches_golden(self, case: GoldenTestCase) -> None: "speech-to-text", "speech-language", "gemma4-assistant", + "ctc-asr", + "feature-ctc-asr", } ) +def _run_ctc_generation( + pkg: ModelPackage, + case: GoldenTestCase, +) -> list[int]: + """Run one CTC forward pass and return deterministic frame argmax IDs.""" + session = _open_decoder_session(pkg) + try: + if case.task_type == "feature-ctc-asr": + feeds = _prepare_feature_ctc_feeds(case) + else: + feeds = _prepare_audio_feeds(case) + feeds["attention_mask"] = np.ones_like(feeds["input_values"], dtype=np.int64) + logits = session.run(feeds)["logits"] + finally: + session.close() + return np.argmax(logits[0], axis=-1).astype(np.int64).tolist() + + def _validate_greedy(case: GoldenTestCase) -> None: """Ensure the test case uses deterministic (greedy) decoding. @@ -2095,7 +2214,11 @@ def _run_multimodel_text_generation( encoder, using 1D position IDs. Returns newly generated token IDs (prompt excluded). """ - suppress_ids = _load_suppress_token_ids(case.model_id, case.trust_remote_code) + suppress_ids = _load_suppress_token_ids( + case.model_id, + revision=case.revision, + trust_remote_code=case.trust_remote_code, + ) device_kwargs = _get_test_device_kwargs() dec_key = "decoder" if "decoder" in pkg else "model" @@ -2300,6 +2423,7 @@ def _run_speech_to_text_generation( # Load audio and extract features (same as L4 prefill) processor = transformers.AutoProcessor.from_pretrained( case.model_id, + revision=case.revision, trust_remote_code=case.trust_remote_code, ) audio_path = _TESTDATA_DIR / case.audio[0] @@ -2401,14 +2525,19 @@ def _run_speech_language_generation( # --- Load audio and extract features --- processor = transformers.AutoProcessor.from_pretrained( - case.model_id, trust_remote_code=case.trust_remote_code + case.model_id, + revision=case.revision, + trust_remote_code=case.trust_remote_code, ) audio_path = _TESTDATA_DIR / case.audio[0] audio_array, _sr = librosa.load(str(audio_path), sr=16000) fe = getattr(processor, "feature_extractor", None) if fe is None or not hasattr(fe, "sampling_rate"): - fe = transformers.WhisperFeatureExtractor.from_pretrained(case.model_id) + fe = transformers.WhisperFeatureExtractor.from_pretrained( + case.model_id, + revision=case.revision, + ) audio_processed = fe( [audio_array], sampling_rate=16000, @@ -2417,7 +2546,11 @@ def _run_speech_language_generation( ) # --- Step 1: audio encoder --- - suppress_ids = _load_suppress_token_ids(case.model_id, case.trust_remote_code) + suppress_ids = _load_suppress_token_ids( + case.model_id, + revision=case.revision, + trust_remote_code=case.trust_remote_code, + ) audio_session = OnnxModelSession(pkg["audio_encoder"], **device_kwargs) try: audio_feeds: dict[str, np.ndarray] = {} @@ -2697,6 +2830,8 @@ def test_generation_matches_golden(self, case: GoldenTestCase) -> None: golden, max_new_tokens=case.generation_params.get("max_new_tokens", 50), ) + elif case.task_type in {"ctc-asr", "feature-ctc-asr"}: + new_tokens = _run_ctc_generation(pkg, case) elif len(pkg) > 1 and "embedding" in pkg: # Multi-model text-generation (e.g. Gemma4 any-to-any text path): # embedding model maps input_ids -> inputs_embeds (+ extra decoder @@ -2716,6 +2851,11 @@ def test_generation_matches_golden(self, case: GoldenTestCase) -> None: expected_tokens = np.array(expected_token_ids, dtype=np.int64) expected_len = len(expected_tokens) actual_len = len(new_tokens) + if case.task_type in {"ctc-asr", "feature-ctc-asr"} and actual_len != expected_len: + pytest.fail( + f"L5 FAIL: CTC frame count changed for {case.case_id}: " + f"expected {expected_len}, got {actual_len}" + ) if actual_len != expected_len: warnings.warn( f"Length mismatch for {case.case_id}: " diff --git a/tests/parakeet_ctc_integration_test.py b/tests/parakeet_ctc_integration_test.py new file mode 100644 index 000000000..945cc8de2 --- /dev/null +++ b/tests/parakeet_ctc_integration_test.py @@ -0,0 +1,135 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Real-checkpoint CUDA parity for NVIDIA Parakeet CTC.""" + +from __future__ import annotations + +import gc +import json +from pathlib import Path + +import librosa +import numpy as np +import onnx_ir as ir +import pytest +import torch +import transformers + +from mobius import build, build_from_module +from mobius._configs import ParakeetCTCConfig +from mobius._testing.ort_inference import OnnxModelSession +from mobius._weight_loading import apply_weights +from mobius.models import ParakeetForCTCModel +from mobius.tasks import FeatureCTCAsrTask + +_MODEL_ID = "nvidia/parakeet-ctc-1.1b" +_REVISION = "20e63a0fed6aedba145b74b826dbd41df0941730" + + +@pytest.mark.integration +def test_parakeet_real_audio_real_weight_cuda_parity(): + if not torch.cuda.is_available(): + pytest.skip("CUDA is required for Parakeet 1.1B real-weight parity") + + processor = transformers.AutoProcessor.from_pretrained(_MODEL_ID, revision=_REVISION) + audio, sample_rate = librosa.load( + str(Path("testdata") / "652-129742-0006.flac"), + sr=16_000, + ) + assert np.any(audio != 0) + inputs = processor( + audio, + sampling_rate=sample_rate, + return_tensors="pt", + ) + + hf_model = transformers.AutoModelForCTC.from_pretrained( + _MODEL_ID, + revision=_REVISION, + torch_dtype=torch.float32, + ).eval() + hf_model.cuda() + with torch.no_grad(): + expected = ( + hf_model(**{name: value.cuda() for name, value in inputs.items()}) + .logits.cpu() + .numpy() + ) + hf_model.cpu() + torch.cuda.empty_cache() + + config = ParakeetCTCConfig.from_transformers(hf_model.config) + config.dtype = ir.DataType.FLOAT + module = ParakeetForCTCModel(config) + package = build_from_module(module, config, task=FeatureCTCAsrTask()) + apply_weights( + package["model"], + module.preprocess_weights(dict(hf_model.state_dict())), + ) + del hf_model + gc.collect() + + session = OnnxModelSession(package["model"], device="cuda") + try: + actual = session.run( + { + "input_features": inputs["input_features"].numpy(), + "attention_mask": inputs["attention_mask"].numpy().astype(bool), + } + )["logits"] + finally: + session.close() + + absolute_error = np.abs(actual - expected) + print( + "Parakeet CUDA parity: " + f"max_abs_diff={absolute_error.max():.6f}, " + f"mean_abs_diff={absolute_error.mean():.6f}" + ) + np.testing.assert_allclose(actual, expected, atol=2e-3, rtol=2e-3) + + +@pytest.mark.integration +def test_parakeet_real_audio_fp16_cuda_matches_ctc_frames(): + if not torch.cuda.is_available(): + pytest.skip("CUDA is required for Parakeet 1.1B fp16 validation") + + processor = transformers.AutoProcessor.from_pretrained(_MODEL_ID, revision=_REVISION) + audio, sample_rate = librosa.load( + str(Path("testdata") / "652-129742-0006.flac"), + sr=16_000, + ) + inputs = processor( + audio, + sampling_rate=sample_rate, + return_tensors="np", + ) + package = build( + _MODEL_ID, + revision=_REVISION, + dtype="f16", + execution_provider="cuda", + ) + session = OnnxModelSession(package["model"], device="cuda") + try: + logits = session.run( + { + "input_features": inputs["input_features"].astype(np.float16), + "attention_mask": inputs["attention_mask"].astype(bool), + } + )["logits"] + finally: + session.close() + + with open( + Path("testdata") / "golden" / "audio" / "parakeet-ctc-1.1b_generation.json" + ) as golden_file: + expected_ids = np.array(json.load(golden_file)["generated_tokens"], dtype=np.int64) + actual_ids = np.argmax(logits[0], axis=-1) + np.testing.assert_array_equal(actual_ids, expected_ids) + assert ( + processor.batch_decode(actual_ids[np.newaxis, :])[0] + == "cauliflower mayonnaise take cold boiled cauliflower break into " + "branches adding salt pepper and vinegar to season" + )