Add ASR fp16 recipe for MahmoudAshraf/mms-300m-1130-forced-aligner

examples/recipes/MahmoudAshraf_mms-300m-1130-forced-aligner/cpu/cpu/automatic-speech-recognition_fp16_config.json

@@ -0,0 +1,58 @@
+{
+  "export": {
+    "opset_version": 17,
+    "batch_size": 1,
+    "export_params": true,
+    "do_constant_folding": true,
+    "verbose": false,
+    "dynamo": false,
+    "enable_hierarchy_tags": true,
+    "clean_onnx": false,
+    "hierarchy_tag_format": "full",
+    "input_tensors": [
+      {
+        "name": "input_values",
+        "dtype": "float32",
+        "shape": [
+          1,
+          16000
+        ],
+        "value_range": [
+          -1,
+          1
+        ]
+      }
+    ],
+    "output_tensors": [
+      {
+        "name": "logits"
+      }
+    ]
+  },
+  "optim": {},
+  "quant": {
+    "mode": "fp16",
+    "samples": 10,
+    "calibration_method": "minmax",
+    "weight_type": "uint8",
+    "activation_type": "uint8",
+    "per_channel": false,
+    "symmetric": false,
+    "weight_symmetric": null,
+    "activation_symmetric": null,
+    "save_calibration": false,
+    "distribution": "uniform",
+    "seed": null,
+    "calibration_load_path": null,
+    "calibration_save_path": null,
+    "op_types_to_quantize": null,
+    "nodes_to_exclude": null,
+    "fp16_keep_io_types": true,
+    "fp16_op_block_list": null
+  },
+  "compile": null,
+  "loader": {
+    "task": "automatic-speech-recognition",
+    "model_type": "wav2vec2"
+  }
+}

src/winml/modelkit/loader/resolution.py

@@ -132,6 +132,38 @@ class from the ``transformers`` package.
         ) from e
 
 
+# Some pipeline tasks span incompatible model families whose AutoModel classes
+# cannot load one another. Keep the mapping architecture-driven: the checkpoint's
+# published ``architectures`` metadata selects the specialized AutoModel, never a
+# model-id allowlist. Add one entry for each future ambiguous task family rather
+# than broadening this into a global preference over TasksManager defaults.
+_ARCHITECTURE_AUTO_CLASS_OVERRIDES: dict[tuple[str, str], str] = {
+    ("automatic-speech-recognition", "ForCTC"): "AutoModelForCTC",
+}
+
+
+def _resolve_architecture_auto_class(config: PretrainedConfig, task: str) -> type | None:
+    """Return an architecture-specific AutoModel class when a task is ambiguous.
+
+    ``automatic-speech-recognition`` covers both encoder-only CTC models and
+    encoder-decoder speech models. Optimum's task-level default is
+    ``AutoModelForSpeechSeq2Seq``, which cannot load a ``Wav2Vec2Config``. A
+    published ``*ForCTC`` architecture is an architecture-wide signal that the
+    correct loader is ``AutoModelForCTC``; seq2seq architectures continue to the
+    existing TasksManager fallback.
+    """
+    architectures = getattr(config, "architectures", None)
+    if not architectures:
+        return None
+
+    architecture = architectures[0]
+    for (mapped_task, suffix), auto_class_name in _ARCHITECTURE_AUTO_CLASS_OVERRIDES.items():
+        if task == mapped_task and architecture.endswith(suffix):
+            transformers_module = importlib.import_module("transformers")
+            return cast("type", getattr(transformers_module, auto_class_name))
+    return None
+
+
 def _detect_task_from_model_class(model_class: type) -> str:
     """Detect task from a model class via TasksManager.
 
@@ -515,6 +547,8 @@ def resolve_task(
             resolved = _get_custom_model_class(
                 model_type_norm, original
             ) or _get_custom_model_class(model_type_norm, normalized)
+        if resolved is None:
+            resolved = _resolve_architecture_auto_class(config, normalized)
         if resolved is None:
             try:
                 resolved = TasksManager.get_model_class_for_task(normalized, framework="pt")
@@ -580,6 +614,8 @@ def resolve_task(
     # --- Stage 2: model class (if not already resolved in 1b) -------------
     if resolved is None:
         resolved = _get_custom_model_class(model_type_norm, opt_task)
+        if resolved is None:
+            resolved = _resolve_architecture_auto_class(config, opt_task)
         if resolved is None:
             try:
                 resolved = TasksManager.get_model_class_for_task(opt_task, framework="pt")

tests/unit/loader/test_config_resolution.py

@@ -84,6 +84,34 @@ def test_auto_detect_new_architectures(
         assert hasattr(r.model_class, "__name__")
 
 
+class TestResolveAmbiguousASRModelClass:
+    """Resolve ASR's CTC and seq2seq families from architecture metadata."""
+
+    @pytest.mark.parametrize("explicit_task", [False, True], ids=["auto-task", "explicit-task"])
+    def test_ctc_architecture_uses_auto_model_for_ctc(
+        self,
+        explicit_task: bool,
+        make_mock_config,
+    ) -> None:
+        config = make_mock_config("wav2vec2", ["Wav2Vec2ForCTC"])
+
+        r = resolve_task(
+            config,
+            task="automatic-speech-recognition" if explicit_task else None,
+        )
+
+        assert r.task == "automatic-speech-recognition"
+        assert r.model_class.__name__ == "AutoModelForCTC"
+
+    def test_seq2seq_architecture_keeps_speech_seq2seq_default(self, make_mock_config) -> None:
+        config = make_mock_config("whisper", ["WhisperForConditionalGeneration"])
+
+        r = resolve_task(config)
+
+        assert r.task == "automatic-speech-recognition"
+        assert r.model_class.__name__ == "AutoModelForSpeechSeq2Seq"
+
+
 class TestResolveTaskAliasPreservation:
     """User-task path: original task is returned, not the normalized form.