Add audeering wav2vec2 dimensional emotion (speech regression) support

examples/recipes/audeering_wav2vec2-large-robust-12-ft-emotion-msp-dim/cpu/cpu/audio-classification_fp16_config.json

@@ -0,0 +1,43 @@
+{
+  "export": {
+    "opset_version": 17,
+    "batch_size": 1,
+    "export_params": true,
+    "do_constant_folding": true,
+    "verbose": false,
+    "dynamo": false,
+    "enable_hierarchy_tags": false,
+    "clean_onnx": false,
+    "hierarchy_tag_format": "full",
+    "input_tensors": [
+      {
+        "name": "input_values",
+        "dtype": "float32",
+        "shape": [
+          1,
+          16000
+        ],
+        "value_range": [
+          -1,
+          1
+        ]
+      }
+    ],
+    "output_tensors": [
+      {
+        "name": "hidden_states"
+      },
+      {
+        "name": "logits"
+      }
+    ]
+  },
+  "optim": {},
+  "quant": null,
+  "compile": null,
+  "loader": {
+    "task": "audio-classification",
+    "model_class": "EmotionModel",
+    "model_type": "wav2vec2_emotion_regression"
+  }
+}

examples/recipes/audeering_wav2vec2-large-robust-12-ft-emotion-msp-dim/cpu/cpu/audio-classification_fp32_config.json

@@ -0,0 +1,43 @@
+{
+  "export": {
+    "opset_version": 17,
+    "batch_size": 1,
+    "export_params": true,
+    "do_constant_folding": true,
+    "verbose": false,
+    "dynamo": false,
+    "enable_hierarchy_tags": false,
+    "clean_onnx": false,
+    "hierarchy_tag_format": "full",
+    "input_tensors": [
+      {
+        "name": "input_values",
+        "dtype": "float32",
+        "shape": [
+          1,
+          16000
+        ],
+        "value_range": [
+          -1,
+          1
+        ]
+      }
+    ],
+    "output_tensors": [
+      {
+        "name": "hidden_states"
+      },
+      {
+        "name": "logits"
+      }
+    ]
+  },
+  "optim": {},
+  "quant": null,
+  "compile": null,
+  "loader": {
+    "task": "audio-classification",
+    "model_class": "EmotionModel",
+    "model_type": "wav2vec2_emotion_regression"
+  }
+}

src/winml/modelkit/models/hf/__init__.py

@@ -95,6 +95,11 @@
 )
 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 .wav2vec2 import (
+    # triggers registration
+    Wav2Vec2EmotionRegressionIOConfig as _Wav2Vec2EmotionRegressionIOConfig,
+)
 from .zoedepth import ZoeDepthIOConfig as _ZoeDepthIOConfig  # triggers registration
 
 
@@ -128,6 +133,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,112 @@
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+# --------------------------------------------------------------------------
+"""Audeering wav2vec2 emotion-regression export variant."""
+
+from __future__ import annotations
+
+from typing import Any, cast
+
+import torch
+import torch.nn as nn
+from optimum.exporters.onnx import OnnxConfig
+from optimum.utils import NormalizedConfig
+from optimum.utils.input_generators import DummyAudioInputGenerator
+from transformers.models.wav2vec2.modeling_wav2vec2 import (
+    Wav2Vec2Model,
+    Wav2Vec2PreTrainedModel,
+)
+
+from ...export import register_onnx_overwrite
+from ..winml import register_specialization
+
+
+EMOTION_REGRESSION_MODEL_TYPE = "wav2vec2_emotion_regression"
+
+
+class RegressionHead(nn.Module):
+    """Audeering dimensional-emotion regression head."""
+
+    def __init__(self, config: Any) -> None:
+        super().__init__()
+        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
+        self.dropout = nn.Dropout(config.final_dropout)
+        self.out_proj = nn.Linear(config.hidden_size, config.num_labels)
+
+    def forward(self, features: torch.Tensor) -> torch.Tensor:  # noqa: D102
+        x = features
+        x = self.dropout(x)
+        x = self.dense(x)
+        x = torch.tanh(x)
+        x = self.dropout(x)
+        return cast("torch.Tensor", self.out_proj(x))
+
+
+class EmotionModel(Wav2Vec2PreTrainedModel):
+    """Audeering wav2vec2 mean-pooling regression model."""
+
+    def __init__(self, config: Any) -> None:
+        super().__init__(config)
+        self.config = config
+        self.wav2vec2 = Wav2Vec2Model(config)
+        self.classifier = RegressionHead(config)
+        self.config.model_type = EMOTION_REGRESSION_MODEL_TYPE
+        self.init_weights()
+
+    def forward(self, input_values: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:  # noqa: D102
+        outputs = self.wav2vec2(input_values)
+        hidden_states = outputs[0]
+        hidden_states = torch.mean(hidden_states, dim=1)
+        logits = self.classifier(hidden_states)
+        return hidden_states, logits
+
+
+_EMOTION_NORMALIZED_CONFIG = NormalizedConfig.with_args(
+    hidden_size="hidden_size",
+    num_labels="num_labels",
+    allow_new=True,
+)
+
+
+@register_onnx_overwrite(
+    EMOTION_REGRESSION_MODEL_TYPE,
+    "audio-classification",
+    library_name="transformers",
+)
+class Wav2Vec2EmotionRegressionIOConfig(OnnxConfig):  # type: ignore[misc]  # optimum base is untyped
+    """ONNX config for raw waveform -> hidden_states, logits."""
+
+    NORMALIZED_CONFIG_CLASS = _EMOTION_NORMALIZED_CONFIG
+    DUMMY_INPUT_GENERATOR_CLASSES = (DummyAudioInputGenerator,)
+
+    @property
+    def inputs(self) -> dict[str, dict[int, str]]:  # noqa: D102
+        return {"input_values": {0: "batch_size", 1: "audio_sequence_length"}}
+
+    @property
+    def outputs(self) -> dict[str, dict[int, str]]:  # noqa: D102
+        return {
+            "hidden_states": {0: "batch_size"},
+            "logits": {0: "batch_size"},
+        }
+
+
+MODEL_CLASS_MAPPING: dict[tuple[str, str], type] = {
+    ("wav2vec2-emotion-regression", "audio-classification"): EmotionModel,
+}
+
+register_specialization(
+    EMOTION_REGRESSION_MODEL_TYPE,
+    "audio-classification",
+    "WinMLModelForGenericTask",
+)
+
+
+__all__ = [
+    "EMOTION_REGRESSION_MODEL_TYPE",
+    "MODEL_CLASS_MAPPING",
+    "EmotionModel",
+    "RegressionHead",
+    "Wav2Vec2EmotionRegressionIOConfig",
+]

tests/unit/models/wav2vec2/test_onnx_config.py

@@ -0,0 +1,142 @@
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+# --------------------------------------------------------------------------
+"""Tests for the audeering wav2vec2 dimensional-emotion export variant.
+
+The audeering ``wav2vec2-large-robust-12-ft-emotion-msp-dim`` checkpoint ships a
+custom mean-pooling ``RegressionHead`` that transformers/Optimum cannot route on
+their own. Support hinges on three contracts these tests lock in:
+
+- the ``("wav2vec2-emotion-regression", "audio-classification")`` entry in
+  ``MODEL_CLASS_MAPPING`` resolving to ``EmotionModel``;
+- ``resolve_task`` honouring the build variant's underscore ``model_type``
+  (``wav2vec2_emotion_regression``) via the ``_`` -> ``-`` normalization the
+  mapping key depends on;
+- ``Wav2Vec2EmotionRegressionIOConfig`` registering for the
+  ``audio-classification`` ONNX export with the expected input/output axes.
+
+A rename of ``EMOTION_REGRESSION_MODEL_TYPE`` or a resolver normalization change
+would otherwise break routing with nothing failing in CI.
+"""
+
+from __future__ import annotations
+
+import pytest
+from optimum.exporters.tasks import TasksManager
+from optimum.utils.input_generators import DummyAudioInputGenerator
+from transformers import Wav2Vec2Config
+
+from winml.modelkit.loader import resolve_task
+from winml.modelkit.models.hf import MODEL_CLASS_MAPPING
+from winml.modelkit.models.hf.wav2vec2 import (
+    EMOTION_REGRESSION_MODEL_TYPE,
+    EmotionModel,
+    Wav2Vec2EmotionRegressionIOConfig,
+)
+from winml.modelkit.models.hf.wav2vec2 import MODEL_CLASS_MAPPING as WAV2VEC2_MAPPING
+
+
+# =============================================================================
+# Test Constants
+# =============================================================================
+
+TASK = "audio-classification"
+MAPPING_KEY = ("wav2vec2-emotion-regression", TASK)
+
+
+# =============================================================================
+# Fixtures
+# =============================================================================
+
+
+@pytest.fixture(scope="module")
+def emotion_config():
+    """Minimal Wav2Vec2Config exercising the emotion-regression head dimensions."""
+    return Wav2Vec2Config(hidden_size=16, num_labels=3)
+
+
+# =============================================================================
+# MODEL_CLASS_MAPPING — routing to the custom regression wrapper
+# =============================================================================
+
+
+class TestWav2Vec2EmotionModelClassMapping:
+    """The emotion-regression variant routes to EmotionModel."""
+
+    def test_mapping_entry_registered(self):
+        """The aggregated mapping exposes the emotion-regression entry."""
+        assert MAPPING_KEY in MODEL_CLASS_MAPPING
+        assert MODEL_CLASS_MAPPING[MAPPING_KEY] is EmotionModel
+
+    def test_module_mapping_merged_into_aggregate(self):
+        """The module-level mapping is included in the aggregated mapping."""
+        assert WAV2VEC2_MAPPING.items() <= MODEL_CLASS_MAPPING.items()
+
+
+# =============================================================================
+# resolve_task — underscore model_type normalization contract
+# =============================================================================
+
+
+class TestWav2Vec2EmotionTaskResolution:
+    """resolve_task honours the underscore build-variant model_type."""
+
+    def test_model_type_override_resolves_emotion_model(self, emotion_config):
+        """An underscore model_type override normalizes and resolves EmotionModel."""
+        resolution = resolve_task(
+            emotion_config,
+            task=TASK,
+            model_class="EmotionModel",
+            model_type_override=EMOTION_REGRESSION_MODEL_TYPE,
+        )
+
+        assert resolution.task == TASK
+        assert resolution.model_class is EmotionModel
+
+    def test_registered_model_type_uses_underscores(self):
+        """The registered model_type uses underscores; the mapping key uses dashes."""
+        assert EMOTION_REGRESSION_MODEL_TYPE == "wav2vec2_emotion_regression"
+        assert EMOTION_REGRESSION_MODEL_TYPE.replace("_", "-") == MAPPING_KEY[0]
+
+
+# =============================================================================
+# Wav2Vec2EmotionRegressionIOConfig — registration and I/O axes
+# =============================================================================
+
+
+class TestWav2Vec2EmotionRegressionIOConfig:
+    """The ONNX export config is registered with the expected axes."""
+
+    def test_onnx_config_registered(self):
+        """Config is registered with TasksManager for audio-classification."""
+        config_cls = TasksManager.get_exporter_config_constructor(
+            model_type=EMOTION_REGRESSION_MODEL_TYPE,
+            exporter="onnx",
+            task=TASK,
+            library_name="transformers",
+        )
+        assert config_cls.func.__name__ == Wav2Vec2EmotionRegressionIOConfig.__name__
+
+    def test_inputs_axes(self, emotion_config):
+        """Inputs expose input_values with batch and audio-length dynamic axes."""
+        io_config = Wav2Vec2EmotionRegressionIOConfig(emotion_config, task=TASK)
+
+        assert io_config.inputs == {
+            "input_values": {0: "batch_size", 1: "audio_sequence_length"},
+        }
+
+    def test_outputs_axes(self, emotion_config):
+        """Outputs expose hidden_states and logits with a batch dynamic axis."""
+        io_config = Wav2Vec2EmotionRegressionIOConfig(emotion_config, task=TASK)
+
+        assert io_config.outputs == {
+            "hidden_states": {0: "batch_size"},
+            "logits": {0: "batch_size"},
+        }
+
+    def test_dummy_input_generator_class(self):
+        """Uses DummyAudioInputGenerator for raw-waveform dummy inputs."""
+        assert (
+            DummyAudioInputGenerator,
+        ) == Wav2Vec2EmotionRegressionIOConfig.DUMMY_INPUT_GENERATOR_CLASSES