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..28ea6d121 --- /dev/null +++ b/examples/recipes/jonatasgrosman_wav2vec2-large-xlsr-53-russian/cpu/cpu/automatic-speech-recognition_fp16_config.json @@ -0,0 +1,44 @@ +{ + "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": { + "precision": "fp16", + "op_types_to_quantize": null, + "nodes_to_exclude": null + }, + "compile": null, + "loader": { + "task": "automatic-speech-recognition", + "model_class": "AutoModelForCTC", + "model_type": "wav2vec2" + } +} 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..4b16ad812 --- /dev/null +++ b/examples/recipes/jonatasgrosman_wav2vec2-large-xlsr-53-russian/cpu/cpu/automatic-speech-recognition_fp32_config.json @@ -0,0 +1,40 @@ +{ + "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" + } +} 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 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..52e5ff0af 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,44 @@ 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 TestCTCModelClassResolution: + """CTC ASR models resolve to AutoModelForCTC, not AutoModelForSpeechSeq2Seq. + + When automatic-speech-recognition maps to multiple candidate classes, + passing model_type enables Optimum to select the correct class per + architecture family. + """ + + @pytest.mark.parametrize( + "model_type,arch_name", + [ + ("wav2vec2", "Wav2Vec2ForCTC"), + ("hubert", "HubertForCTC"), + ("data2vec-audio", "Data2VecAudioForCTC"), + ], + ) + def test_ctc_architecture_resolves_to_auto_model_for_ctc(self, model_type, arch_name): + """ForCTC architectures resolve to AutoModelForCTC via model_type.""" + config = MagicMock() + config.model_type = model_type + config.architectures = [arch_name] + config._name_or_path = "" + + r = resolve_task(config) + + assert r.task == "automatic-speech-recognition" + assert r.model_class.__name__ == "AutoModelForCTC" + + def test_whisper_still_resolves_to_speech_seq2seq(self): + """Non-CTC ASR models (Whisper) still get AutoModelForSpeechSeq2Seq.""" + config = MagicMock() + config.model_type = "whisper" + config.architectures = ["WhisperForConditionalGeneration"] + config._name_or_path = "" + + r = resolve_task(config) + + assert r.task == "automatic-speech-recognition" + assert r.model_class.__name__ == "AutoModelForSpeechSeq2Seq"