From 63b905a3491991f5665b6ad01b97bb73e6766625 Mon Sep 17 00:00:00 2001 From: Yong Yue Date: Tue, 21 Jul 2026 11:16:53 +0800 Subject: [PATCH 1/3] Pass model_type to TasksManager for accurate model class resolution When a task like automatic-speech-recognition maps to multiple model classes (AutoModelForSpeechSeq2Seq, AutoModelForCTC), Optimum defaults to the first one without model_type context. This caused CTC-based ASR models (e.g. Wav2Vec2ForCTC) to fail with an incorrect class. Passing model_type_norm enables Optimum to select the correct class per architecture family. Add regression recipe for jonatasgrosman/wav2vec2-large-xlsr-53-russian (cpu/fp32 and cpu/fp16). --- ...omatic-speech-recognition_fp16_config.json | 60 +++++++++++++++++++ ...omatic-speech-recognition_fp32_config.json | 41 +++++++++++++ src/winml/modelkit/loader/resolution.py | 4 +- 3 files changed, 104 insertions(+), 1 deletion(-) create mode 100644 examples/recipes/jonatasgrosman_wav2vec2-large-xlsr-53-russian/cpu/cpu/automatic-speech-recognition_fp16_config.json create mode 100644 examples/recipes/jonatasgrosman_wav2vec2-large-xlsr-53-russian/cpu/cpu/automatic-speech-recognition_fp32_config.json diff --git a/examples/recipes/jonatasgrosman_wav2vec2-large-xlsr-53-russian/cpu/cpu/automatic-speech-recognition_fp16_config.json b/examples/recipes/jonatasgrosman_wav2vec2-large-xlsr-53-russian/cpu/cpu/automatic-speech-recognition_fp16_config.json new file mode 100644 index 000000000..55acb2cb3 --- /dev/null +++ b/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 + } +} diff --git a/examples/recipes/jonatasgrosman_wav2vec2-large-xlsr-53-russian/cpu/cpu/automatic-speech-recognition_fp32_config.json b/examples/recipes/jonatasgrosman_wav2vec2-large-xlsr-53-russian/cpu/cpu/automatic-speech-recognition_fp32_config.json new file mode 100644 index 000000000..467298225 --- /dev/null +++ b/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 + } +} diff --git a/src/winml/modelkit/loader/resolution.py b/src/winml/modelkit/loader/resolution.py index 548848175..9d82d11d6 100644 --- a/src/winml/modelkit/loader/resolution.py +++ b/src/winml/modelkit/loader/resolution.py @@ -644,7 +644,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_norm or None + ) except Exception: resolved = _resolve_model_class_from_config(config) # arch fallback From 4948419d678ec1a197b1438ca9e336b5e2602d47 Mon Sep 17 00:00:00 2001 From: Yong Yue Date: Wed, 22 Jul 2026 14:12:05 +0800 Subject: [PATCH 2/3] fix(loader): pass model_type in USER_TASK resolution branch The USER_TASK branch called get_model_class_for_task without model_type, causing models like Wav2Vec2 (CTC) to resolve to AutoModelForSpeechSeq2Seq instead of AutoModelForCTC when task is explicit but model_class is not. Add regression test for explicit-task resolution without model-class override. --- src/winml/modelkit/loader/resolution.py | 4 +- tests/unit/loader/test_config_resolution.py | 47 +++++++++++++++++++++ 2 files changed, 50 insertions(+), 1 deletion(-) diff --git a/src/winml/modelkit/loader/resolution.py b/src/winml/modelkit/loader/resolution.py index 9d82d11d6..846a084ae 100644 --- a/src/winml/modelkit/loader/resolution.py +++ b/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_norm or None + ) except KeyError as e: if composite is not None: # Pure composite (e.g. table-question-answering): no single model class diff --git a/tests/unit/loader/test_config_resolution.py b/tests/unit/loader/test_config_resolution.py index f11b89492..845d94ae6 100644 --- a/tests/unit/loader/test_config_resolution.py +++ b/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." + ) From 2d76b80020d2789c7ef17d123973c35269ed095d Mon Sep 17 00:00:00 2001 From: Yong Yue Date: Thu, 23 Jul 2026 20:15:44 +0800 Subject: [PATCH 3/3] fix(loader): pass raw model_type to TasksManager instead of normalized Optimum registers some model types with underscores (e.g. speech_to_text). Passing the hyphen-normalized variant prevents Optimum from matching the correct AutoModel class. Keep model_type_norm for WinML internal mapping keys, pass the original model_type to both TasksManager calls. --- src/winml/modelkit/loader/resolution.py | 4 ++-- .../test_resolve_task_and_model_class.py | 23 +++++++++++++++++++ 2 files changed, 25 insertions(+), 2 deletions(-) diff --git a/src/winml/modelkit/loader/resolution.py b/src/winml/modelkit/loader/resolution.py index 846a084ae..d64100264 100644 --- a/src/winml/modelkit/loader/resolution.py +++ b/src/winml/modelkit/loader/resolution.py @@ -558,7 +558,7 @@ def resolve_task( if resolved is None: try: resolved = TasksManager.get_model_class_for_task( - normalized, framework="pt", model_type=model_type_norm or None + normalized, framework="pt", model_type=model_type or None ) except KeyError as e: if composite is not None: @@ -647,7 +647,7 @@ def resolve_task( if resolved is None: try: resolved = TasksManager.get_model_class_for_task( - opt_task, framework="pt", model_type=model_type_norm or None + opt_task, framework="pt", model_type=model_type or None ) except Exception: resolved = _resolve_model_class_from_config(config) # arch fallback diff --git a/tests/unit/loader/test_resolve_task_and_model_class.py b/tests/unit/loader/test_resolve_task_and_model_class.py index 610677a16..48454682d 100644 --- a/tests/unit/loader/test_resolve_task_and_model_class.py +++ b/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__