diff --git a/examples/recipes/facebook_mms-1b-all/cpu/cpu/automatic-speech-recognition_fp32_config.json b/examples/recipes/facebook_mms-1b-all/cpu/cpu/automatic-speech-recognition_fp32_config.json new file mode 100644 index 000000000..4b16ad812 --- /dev/null +++ b/examples/recipes/facebook_mms-1b-all/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/models/hf/__init__.py b/src/winml/modelkit/models/hf/__init__.py index 0339b2b8e..8d7b2988e 100644 --- a/src/winml/modelkit/models/hf/__init__.py +++ b/src/winml/modelkit/models/hf/__init__.py @@ -100,6 +100,7 @@ ) from .vision_encoder_decoder import VisionEncoderIOConfig as _VisionEncoderIOConfig from .vitpose import MODEL_CLASS_MAPPING as _VITPOSE_CLASS_MAPPING +from .wav2vec2 import MODEL_CLASS_MAPPING as _WAV2VEC2_CLASS_MAPPING from .zoedepth import ZoeDepthIOConfig as _ZoeDepthIOConfig # triggers registration @@ -133,6 +134,7 @@ _T5_CLASS_MAPPING, _VED_CLASS_MAPPING, _VITPOSE_CLASS_MAPPING, + _WAV2VEC2_CLASS_MAPPING, ) for _key, _model_cls in _sub_mapping.items() } diff --git a/src/winml/modelkit/models/hf/wav2vec2.py b/src/winml/modelkit/models/hf/wav2vec2.py new file mode 100644 index 000000000..8e5d61f73 --- /dev/null +++ b/src/winml/modelkit/models/hf/wav2vec2.py @@ -0,0 +1,21 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- +"""Wav2Vec2 HuggingFace model class routing. + +Wav2Vec2 supports several audio heads. Optimum can infer/export the ONNX config, +but its task default for ASR is ambiguous between seq2seq and CTC loaders. Route +CTC ASR checkpoints through ``AutoModelForCTC`` so recipe-free resolution follows +the architecture head instead of the seq2seq fallback. +""" + +from __future__ import annotations + +from transformers import AutoModelForAudioClassification, AutoModelForCTC + + +MODEL_CLASS_MAPPING: dict[tuple[str, str], type] = { + ("wav2vec2", "audio-classification"): AutoModelForAudioClassification, + ("wav2vec2", "automatic-speech-recognition"): AutoModelForCTC, +} diff --git a/tests/unit/models/test_wav2vec2_mapping.py b/tests/unit/models/test_wav2vec2_mapping.py new file mode 100644 index 000000000..374ceb9ca --- /dev/null +++ b/tests/unit/models/test_wav2vec2_mapping.py @@ -0,0 +1,50 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- +"""Tests for Wav2Vec2 task-specific model-class routing.""" + +from __future__ import annotations + +from typing import ClassVar + +from winml.modelkit.loader import resolve_task +from winml.modelkit.loader.resolution import _get_custom_model_class +from winml.modelkit.models.hf import MODEL_CLASS_MAPPING +from winml.modelkit.models.hf.wav2vec2 import MODEL_CLASS_MAPPING as WAV2VEC2_MAPPING + + +class _Config: + model_type = "wav2vec2" + architectures: ClassVar[list[str]] = ["Wav2Vec2ForCTC"] + _name_or_path = "facebook/mms-1b-all" + is_encoder_decoder = False + + +class TestWav2Vec2Mapping: + """Wav2Vec2 ASR and audio-classification loaders use architecture-appropriate classes.""" + + def test_asr_mapping_returns_ctc_model(self): + """Wav2Vec2 ASR checkpoints should resolve to the CTC loader.""" + result = _get_custom_model_class("wav2vec2", "automatic-speech-recognition") + + assert result is not None + assert result.__name__ == "AutoModelForCTC" + + def test_audio_classification_mapping_returns_audio_model(self): + """Wav2Vec2 sequence-classification checkpoints keep the audio loader.""" + result = _get_custom_model_class("wav2vec2", "audio-classification") + + assert result is not None + assert result.__name__ == "AutoModelForAudioClassification" + + def test_mapping_merged_into_aggregate(self): + """The module-level mapping is included in the aggregated mapping.""" + assert WAV2VEC2_MAPPING.items() <= MODEL_CLASS_MAPPING.items() + + def test_detected_ctc_asr_uses_ctc_loader(self): + """Detected Wav2Vec2ForCTC ASR resolves to AutoModelForCTC.""" + resolution = resolve_task(_Config()) + + assert resolution.task == "automatic-speech-recognition" + assert resolution.model_class.__name__ == "AutoModelForCTC"