Fix CTC model class resolution by passing model_type to TasksManager

examples/recipes/jonatasgrosman_wav2vec2-large-xlsr-53-russian/cpu/cpu/automatic-speech-recognition_fp16_config.json

@@ -0,0 +1,60 @@
+{
+  "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_class": "AutoModelForCTC",
+    "model_type": "wav2vec2",
+    "trust_remote_code": true
+  }
+}

examples/recipes/jonatasgrosman_wav2vec2-large-xlsr-53-russian/cpu/cpu/automatic-speech-recognition_fp32_config.json

@@ -0,0 +1,41 @@
+{
+  "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": null,
+  "compile": null,
+  "loader": {
+    "task": "automatic-speech-recognition",
+    "model_class": "AutoModelForCTC",
+    "model_type": "wav2vec2",
+    "trust_remote_code": true
+  }
+}

src/winml/modelkit/loader/resolution.py

@@ -557,7 +557,9 @@ def resolve_task(
             ) or _get_custom_model_class(model_type_norm, normalized)
         if resolved is None:
             try:
-                resolved = TasksManager.get_model_class_for_task(normalized, framework="pt")
+                resolved = TasksManager.get_model_class_for_task(
+                    normalized, framework="pt", model_type=model_type or None
+                )
             except KeyError as e:
                 if composite is not None:
                     # Pure composite (e.g. table-question-answering): no single model class
@@ -644,7 +646,9 @@ def resolve_task(
         resolved = _get_custom_model_class(model_type_norm, opt_task)
         if resolved is None:
             try:
-                resolved = TasksManager.get_model_class_for_task(opt_task, framework="pt")
+                resolved = TasksManager.get_model_class_for_task(
+                    opt_task, framework="pt", model_type=model_type or None
+                )
             except Exception:
                 resolved = _resolve_model_class_from_config(config)  # arch fallback
 

tests/unit/loader/test_config_resolution.py

@@ -225,3 +225,50 @@ def test_model_type_normalization(
             f"Expected {expected_class_name} for model_type='{raw_model_type}', "
             f"got {r.model_class.__name__}. model_type normalization may be broken."
         )
+
+
+class TestExplicitTaskWithModelType:
+    """Explicit task without model_class must still use model_type for disambiguation.
+
+    Regression: when task is supplied but model_class is not, the USER_TASK branch
+    must pass model_type to get_model_class_for_task so that architectures sharing
+    a task (e.g. Wav2Vec2 CTC vs Whisper Seq2Seq for automatic-speech-recognition)
+    resolve to the correct model class.
+    """
+
+    @pytest.mark.parametrize(
+        "model_type,arch,task,expected_class",
+        [
+            pytest.param(
+                "wav2vec2",
+                "Wav2Vec2ForCTC",
+                "automatic-speech-recognition",
+                "AutoModelForCTC",
+                id="wav2vec2-asr-resolves-to-ctc",
+            ),
+            pytest.param(
+                "whisper",
+                "WhisperForConditionalGeneration",
+                "automatic-speech-recognition",
+                "AutoModelForSpeechSeq2Seq",
+                id="whisper-asr-resolves-to-seq2seq",
+            ),
+        ],
+    )
+    def test_explicit_task_respects_model_type(
+        self,
+        model_type: str,
+        arch: str,
+        task: str,
+        expected_class: str,
+        make_mock_config,
+    ) -> None:
+        """Explicit task + model_type resolves the model-type-specific class."""
+        config = make_mock_config(model_type, [arch])
+
+        r = resolve_task(config, task=task)
+
+        assert r.model_class.__name__ == expected_class, (
+            f"Expected {expected_class} for model_type='{model_type}' with task='{task}', "
+            f"got {r.model_class.__name__}. USER_TASK branch may not be passing model_type."
+        )

tests/unit/loader/test_resolve_task_and_model_class.py

@@ -117,3 +117,26 @@ def test_task_normalized_with_model_class(self):
         # Task is normalized (unlike the user-task path which preserves the original)
         assert r.task == "fill-mask"
         assert r.source == TaskSource.USER_CLASS
+
+
+class TestUnderscoreModelTypePassedToTasksManager:
+    """Regression: model_type with underscores must reach TasksManager un-normalized.
+
+    Optimum registers some model types with underscores (e.g. speech_to_text).
+    Converting to hyphens (speech-to-text) prevents Optimum from matching the
+    correct AutoModel class. See PR review comment on explicit ASR resolution.
+    """
+
+    def test_speech_to_text_resolves_seq2seq_class(self):
+        """speech_to_text (underscore) resolves AutoModelForSpeechSeq2Seq via TasksManager."""
+        config = MagicMock()
+        config.model_type = "speech_to_text"
+        config.architectures = ["Speech2TextForConditionalGeneration"]
+        config._name_or_path = ""
+
+        r = resolve_task(config, task="automatic-speech-recognition")
+
+        assert r.task == "automatic-speech-recognition"
+        assert r.source == TaskSource.USER_TASK
+        # TasksManager should resolve to SpeechSeq2Seq for speech_to_text
+        assert "Seq2Seq" in r.model_class.__name__ or "Speech" in r.model_class.__name__