model(wav2vec2): support facebook mms 1b asr

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"
+  }
+}

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()
 }

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,
+}

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"