diff --git a/examples/recipes/audeering_wav2vec2-large-robust-12-ft-emotion-msp-dim/cpu/cpu/audio-classification_fp16_config.json b/examples/recipes/audeering_wav2vec2-large-robust-12-ft-emotion-msp-dim/cpu/cpu/audio-classification_fp16_config.json new file mode 100644 index 000000000..07673bafa --- /dev/null +++ b/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" + } +} diff --git a/examples/recipes/audeering_wav2vec2-large-robust-12-ft-emotion-msp-dim/cpu/cpu/audio-classification_fp32_config.json b/examples/recipes/audeering_wav2vec2-large-robust-12-ft-emotion-msp-dim/cpu/cpu/audio-classification_fp32_config.json new file mode 100644 index 000000000..07673bafa --- /dev/null +++ b/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" + } +} diff --git a/src/winml/modelkit/models/hf/__init__.py b/src/winml/modelkit/models/hf/__init__.py index 094714494..7298b2341 100644 --- a/src/winml/modelkit/models/hf/__init__.py +++ b/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() } diff --git a/src/winml/modelkit/models/hf/wav2vec2.py b/src/winml/modelkit/models/hf/wav2vec2.py new file mode 100644 index 000000000..869734272 --- /dev/null +++ b/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", +] diff --git a/tests/unit/models/wav2vec2/test_onnx_config.py b/tests/unit/models/wav2vec2/test_onnx_config.py new file mode 100644 index 000000000..dab7a16a0 --- /dev/null +++ b/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 diff --git a/tests/unit/recipes/test_cpu_recipes.py b/tests/unit/recipes/test_cpu_recipes.py new file mode 100644 index 000000000..9884dde5b --- /dev/null +++ b/tests/unit/recipes/test_cpu_recipes.py @@ -0,0 +1,74 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- + +import json +from pathlib import Path + +import pytest + +from winml.modelkit.config import WinMLBuildConfig + + +REPO_ROOT = Path(__file__).resolve().parents[3] + +recipes = [ + { + "path": REPO_ROOT + / "examples" + / "recipes" + / "audeering_wav2vec2-large-robust-12-ft-emotion-msp-dim" + / "cpu" + / "cpu" + / "audio-classification_fp32_config.json", + "loader_task": "audio-classification", + "loader_model_class": "EmotionModel", + "loader_model_type": "wav2vec2_emotion_regression", + "opset_version": 17, + }, + { + "path": REPO_ROOT + / "examples" + / "recipes" + / "audeering_wav2vec2-large-robust-12-ft-emotion-msp-dim" + / "cpu" + / "cpu" + / "audio-classification_fp16_config.json", + "loader_task": "audio-classification", + "loader_model_class": "EmotionModel", + "loader_model_type": "wav2vec2_emotion_regression", + "opset_version": 17, + }, +] + + +@pytest.mark.parametrize( + "rec", + recipes, + ids=["audeering-wav2vec2-emotion-fp32", "audeering-wav2vec2-emotion-fp16"], +) +def test_cpu_recipes(rec): + path: Path = rec["path"] + assert path.exists(), f"Recipe file missing: {path}" + + # EP/device is encoded by folder layout: ///.json + assert path.parent.name == "cpu" # device + assert path.parent.parent.name == "cpu" # ep + + data = json.loads(path.read_text(encoding="utf-8")) + + # Construct the validated config from the recipe dict + config = WinMLBuildConfig.from_dict(data) + + # export.opset_version exact + assert config.export is not None + assert config.export.opset_version == rec["opset_version"] + + # loader routes to the emotion-regression head + assert config.loader.task == rec["loader_task"] + assert config.loader.model_class == rec["loader_model_class"] + assert config.loader.model_type == rec["loader_model_type"] + + # float CPU recipe (fp32/fp16): no quantization + assert config.quant is None