diff --git a/examples/recipes/MahmoudAshraf_mms-300m-1130-forced-aligner/cpu/cpu/automatic-speech-recognition_fp16_config.json b/examples/recipes/MahmoudAshraf_mms-300m-1130-forced-aligner/cpu/cpu/automatic-speech-recognition_fp16_config.json index cc15d0ecd..80d6f2efd 100644 --- a/examples/recipes/MahmoudAshraf_mms-300m-1130-forced-aligner/cpu/cpu/automatic-speech-recognition_fp16_config.json +++ b/examples/recipes/MahmoudAshraf_mms-300m-1130-forced-aligner/cpu/cpu/automatic-speech-recognition_fp16_config.json @@ -21,6 +21,18 @@ -1, 1 ] + }, + { + "name": "attention_mask", + "dtype": "int64", + "shape": [ + 1, + 16000 + ], + "value_range": [ + 0, + 2 + ] } ], "output_tensors": [ diff --git a/examples/recipes/MahmoudAshraf_mms-300m-1130-forced-aligner/cpu/cpu/automatic-speech-recognition_fp32_config.json b/examples/recipes/MahmoudAshraf_mms-300m-1130-forced-aligner/cpu/cpu/automatic-speech-recognition_fp32_config.json index 92aecbc3c..6ecdddc57 100644 --- a/examples/recipes/MahmoudAshraf_mms-300m-1130-forced-aligner/cpu/cpu/automatic-speech-recognition_fp32_config.json +++ b/examples/recipes/MahmoudAshraf_mms-300m-1130-forced-aligner/cpu/cpu/automatic-speech-recognition_fp32_config.json @@ -21,6 +21,18 @@ -1, 1 ] + }, + { + "name": "attention_mask", + "dtype": "int64", + "shape": [ + 1, + 16000 + ], + "value_range": [ + 0, + 2 + ] } ], "output_tensors": [ diff --git a/examples/recipes/facebook_mms-1b-all/cpu/cpu/automatic-speech-recognition_fp16_config.json b/examples/recipes/facebook_mms-1b-all/cpu/cpu/automatic-speech-recognition_fp16_config.json index f2d556b8f..aba2be9a1 100644 --- a/examples/recipes/facebook_mms-1b-all/cpu/cpu/automatic-speech-recognition_fp16_config.json +++ b/examples/recipes/facebook_mms-1b-all/cpu/cpu/automatic-speech-recognition_fp16_config.json @@ -21,6 +21,18 @@ -1, 1 ] + }, + { + "name": "attention_mask", + "dtype": "int64", + "shape": [ + 1, + 16000 + ], + "value_range": [ + 0, + 2 + ] } ], "output_tensors": [ @@ -57,6 +69,24 @@ "loader": { "task": "automatic-speech-recognition", "model_class": "AutoModelForCTC", - "model_type": "wav2vec2" + "model_type": "wav2vec2", + "target_lang": "eng" + }, + "eval": { + "task": "automatic-speech-recognition", + "dataset": { + "path": "google/fleurs", + "name": "en_us", + "split": "validation", + "samples": 2, + "shuffle": false, + "seed": 42, + "streaming": false, + "revision": "70bb2e84b976b7e960aa89f1c648e09c59f894dd", + "columns_mapping": { + "input_column": "audio", + "label_column": "transcription" + } + } } } \ No newline at end of file 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 index bc5c3f1d5..bf54233be 100644 --- 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 @@ -21,6 +21,18 @@ -1, 1 ] + }, + { + "name": "attention_mask", + "dtype": "int64", + "shape": [ + 1, + 16000 + ], + "value_range": [ + 0, + 2 + ] } ], "output_tensors": [ @@ -38,6 +50,24 @@ "loader": { "task": "automatic-speech-recognition", "model_class": "AutoModelForCTC", - "model_type": "wav2vec2" + "model_type": "wav2vec2", + "target_lang": "eng" + }, + "eval": { + "task": "automatic-speech-recognition", + "dataset": { + "path": "google/fleurs", + "name": "en_us", + "split": "validation", + "samples": 2, + "shuffle": false, + "seed": 42, + "streaming": false, + "revision": "70bb2e84b976b7e960aa89f1c648e09c59f894dd", + "columns_mapping": { + "input_column": "audio", + "label_column": "transcription" + } + } } } diff --git a/src/winml/modelkit/build/hf.py b/src/winml/modelkit/build/hf.py index 4094b8938..ab8e5e811 100644 --- a/src/winml/modelkit/build/hf.py +++ b/src/winml/modelkit/build/hf.py @@ -422,6 +422,7 @@ def _load_model( hf_config=hf_config, model_type=model_type, attn_implementation=attn_implementation, + target_lang=config.loader.target_lang, ) return pytorch_model diff --git a/src/winml/modelkit/commands/build.py b/src/winml/modelkit/commands/build.py index a56d9f52e..8de07d6d0 100644 --- a/src/winml/modelkit/commands/build.py +++ b/src/winml/modelkit/commands/build.py @@ -2068,6 +2068,7 @@ def _build_hf_pipeline( max_iters: int = extra_kwargs.pop("hack_max_optim_iterations", 3) allow_unsupported_nodes: bool = extra_kwargs.pop("allow_unsupported_nodes", False) + skip_optimize: bool = extra_kwargs.pop("skip_optimize", config.skip_optimize) model_label = model_id or "random-init" # ── Validate + setup ───────────────────────────────────────── @@ -2153,6 +2154,7 @@ def _name(base: str) -> str: show_io_first=False, analyze_output_path=analyze_result_path, allow_unsupported_nodes=allow_unsupported_nodes, + skip_optimize=skip_optimize, ) # Persist config after autoconf diff --git a/src/winml/modelkit/eval/__init__.py b/src/winml/modelkit/eval/__init__.py index 435601a15..74eb4f949 100644 --- a/src/winml/modelkit/eval/__init__.py +++ b/src/winml/modelkit/eval/__init__.py @@ -19,6 +19,7 @@ if TYPE_CHECKING: + from .ctc_asr_evaluator import WinMLCTCASREvaluator from .depth_estimation_evaluator import WinMLDepthEstimationEvaluator from .feature_extraction_evaluator import WinMLFeatureExtractionEvaluator from .fill_mask_evaluator import WinMLFillMaskEvaluator @@ -47,6 +48,7 @@ _LAZY_ATTRS: dict[str, str] = { # Evaluators + "WinMLCTCASREvaluator": ".ctc_asr_evaluator:WinMLCTCASREvaluator", "WinMLDepthEstimationEvaluator": ".depth_estimation_evaluator:WinMLDepthEstimationEvaluator", "WinMLFeatureExtractionEvaluator": ( ".feature_extraction_evaluator:WinMLFeatureExtractionEvaluator" @@ -126,6 +128,7 @@ def __dir__() -> list[str]: "SpearmanCorrelationMetric", "TensorSimilarityEvaluator", "TopKAccuracyMetric", + "WinMLCTCASREvaluator", "WinMLDepthEstimationEvaluator", "WinMLEvaluationConfig", "WinMLEvaluator", diff --git a/src/winml/modelkit/eval/ctc_asr_evaluator.py b/src/winml/modelkit/eval/ctc_asr_evaluator.py new file mode 100644 index 000000000..ecb4ec1ce --- /dev/null +++ b/src/winml/modelkit/eval/ctc_asr_evaluator.py @@ -0,0 +1,586 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- + +"""Metadata-driven CTC automatic speech recognition evaluation.""" + +from __future__ import annotations + +import math +import unicodedata +from collections import Counter +from io import BytesIO +from pathlib import Path +from typing import TYPE_CHECKING, Any + +import numpy as np + +from ..utils.eval_utils import DatasetValidationError, validate_dataset_columns +from .base_evaluator import WinMLEvaluator + + +if TYPE_CHECKING: + from datasets import Dataset + from transformers.pipelines.base import Pipeline + + from .config import WinMLEvaluationConfig + +_MAX_WINDOWS_PER_UTTERANCE = 64 + + +class _RejectedSampleError(ValueError): + """A dataset row that cannot be scored under the ASR contract.""" + + +def _normalize_transcript(value: Any) -> str: + """Normalize Unicode and collapse whitespace without changing text semantics.""" + if not isinstance(value, str): + raise _RejectedSampleError("transcription is not a string") + return " ".join(unicodedata.normalize("NFKC", value).split()) + + +def _edit_distance(reference: list[str], prediction: list[str]) -> int: + """Return Levenshtein distance with memory bounded by the shorter sequence.""" + if len(reference) < len(prediction): + reference, prediction = prediction, reference + previous = list(range(len(prediction) + 1)) + for row, reference_item in enumerate(reference, start=1): + current = [row] + for column, prediction_item in enumerate(prediction, start=1): + current.append( + min( + current[-1] + 1, + previous[column] + 1, + previous[column - 1] + (reference_item != prediction_item), + ) + ) + previous = current + return previous[-1] + + +def _corpus_error_rate(references: list[str], predictions: list[str], *, words: bool) -> float: + """Compute corpus word or character error rate.""" + if len(references) != len(predictions) or not references: + raise ValueError("WER/CER require equally sized, non-empty reference and prediction lists.") + reference_units = [text.split() if words else list(text) for text in references] + prediction_units = [text.split() if words else list(text) for text in predictions] + denominator = sum(len(units) for units in reference_units) + if denominator == 0: + raise ValueError("WER/CER cannot score empty normalized references.") + errors = sum( + _edit_distance(reference, prediction) + for reference, prediction in zip(reference_units, prediction_units, strict=True) + ) + return errors / denominator + + +def _decode_audio(value: Any) -> tuple[np.ndarray, int]: + """Decode a datasets ``Audio(decode=False)`` bytes/path value with SoundFile.""" + if not isinstance(value, dict): + raise _RejectedSampleError("audio must be a decode=False bytes/path mapping") + source: BytesIO | str + audio_bytes = value.get("bytes") + audio_path = value.get("path") + if isinstance(audio_bytes, bytes): + source = BytesIO(audio_bytes) + elif isinstance(audio_path, str) and audio_path: + source = audio_path + else: + raise _RejectedSampleError("audio has neither bytes nor path") + + try: + import soundfile as sf + + waveform, sampling_rate = sf.read(source, dtype="float32", always_2d=False) + except Exception as error: + raise _RejectedSampleError(f"SoundFile decode failed: {error}") from error + waveform = np.asarray(waveform, dtype=np.float32) + if waveform.ndim == 2: + waveform = waveform.mean(axis=1, dtype=np.float32) + if waveform.ndim != 1 or waveform.size == 0: + raise _RejectedSampleError(f"decoded audio has invalid shape {waveform.shape}") + return waveform, int(sampling_rate) + + +def _require_soundfile() -> None: + """Fail before dataset work when the optional audio decoder is unavailable.""" + try: + import soundfile # noqa: F401 + except ImportError as error: + raise ImportError( + "CTC ASR evaluation requires SoundFile; install the audio extra with " + "`pip install 'winml-cli[audio]'`." + ) from error + + +def _resample_audio( + waveform: np.ndarray, + source_rate: int, + target_rate: int, +) -> np.ndarray: + """Resample mono audio only when processor metadata requests another rate.""" + if source_rate == target_rate: + return waveform + if source_rate <= 0 or target_rate <= 0: + raise _RejectedSampleError("audio sampling rates must be positive") + from scipy.signal import resample_poly + + divisor = math.gcd(source_rate, target_rate) + return np.asarray( + resample_poly(waveform, target_rate // divisor, source_rate // divisor), + dtype=np.float32, + ) + + +def _is_ctc_config(config: Any) -> bool: + architectures = getattr(config, "architectures", None) or [] + return any( + isinstance(architecture, str) + and (architecture.endswith("ForCTC") or architecture == "AutoModelForCTC") + for architecture in architectures + ) + + +def _validate_ctc_model_inputs(model: Any) -> None: + """Reject exported CTC input contracts this evaluator cannot preprocess.""" + io_config = getattr(model, "io_config", None) + if not isinstance(io_config, dict): + return + input_names = io_config.get("input_names") + if not isinstance(input_names, list) or not input_names: + return + supported_names = {"input_values", "attention_mask"} + if "input_values" not in input_names or not set(input_names).issubset(supported_names): + raise ValueError( + "CTC ASR evaluation supports only rank-2 input_values with an optional " + "attention_mask; exported input contract is " + f"{input_names!r}." + ) + input_shapes = io_config.get("input_shapes") + if not isinstance(input_shapes, list): + return + for name in input_names: + index = input_names.index(name) + if index >= len(input_shapes) or not isinstance(input_shapes[index], list): + continue + if len(input_shapes[index]) != 2: + raise ValueError( + f"CTC ASR evaluation requires rank-2 {name}; exported shape is " + f"{input_shapes[index]!r}." + ) + + +def _load_ctc_processor(model_id: str, *, trust_remote_code: bool) -> Any: + """Load the published processor, falling back only from an unavailable optional LM.""" + from transformers import AutoProcessor + + try: + return AutoProcessor.from_pretrained(model_id, trust_remote_code=trust_remote_code) + except ImportError as error: + message = str(error) + if "Wav2Vec2ProcessorWithLM" not in message or "pyctcdecode" not in message: + raise + + from transformers import AutoFeatureExtractor, AutoTokenizer, Wav2Vec2Processor + + try: + feature_extractor = AutoFeatureExtractor.from_pretrained( + model_id, + trust_remote_code=trust_remote_code, + ) + tokenizer = AutoTokenizer.from_pretrained( + model_id, + trust_remote_code=trust_remote_code, + ) + return Wav2Vec2Processor(feature_extractor=feature_extractor, tokenizer=tokenizer) + except Exception as error: + raise ValueError( + "The checkpoint declares Wav2Vec2ProcessorWithLM, but optional pyctcdecode is " + "unavailable and its feature extractor/tokenizer cannot form a greedy CTC processor." + ) from error + + +def _configure_processor_language(processor: Any, model_config: Any) -> str | None: + """Select and validate a tokenizer language using checkpoint metadata only.""" + tokenizer = getattr(processor, "tokenizer", processor) + configured = getattr(model_config, "target_lang", None) or getattr( + model_config, + "adapter_lang", + None, + ) + active = getattr(tokenizer, "target_lang", None) + adapter_attn_dim = getattr(model_config, "adapter_attn_dim", None) + adapter_capable = ( + isinstance(adapter_attn_dim, int) and adapter_attn_dim > 0 + ) or active is not None + if not adapter_capable: + if configured is not None: + raise ValueError( + f"Checkpoint requests target language {configured!r}, but its metadata does not " + "expose language-adapter support." + ) + return None + if configured is not None and configured != active: + setter = getattr(tokenizer, "set_target_lang", None) + if not callable(setter): + raise ValueError( + f"Checkpoint requests target language {configured!r}, but its tokenizer " + "cannot select a language." + ) + try: + setter(configured) + except Exception as error: + raise ValueError( + f"Checkpoint tokenizer cannot select requested adapter language {configured!r}." + ) from error + active = getattr(tokenizer, "target_lang", configured) + if active != configured: + raise ValueError( + f"Checkpoint tokenizer did not activate requested adapter language {configured!r}." + ) + if not active: + raise ValueError("The checkpoint supports language adapters but no language is active.") + return str(active) if active is not None else None + + +def _normalize_identity_scalar(value: Any, field: str) -> str | int | float: + if isinstance(value, np.generic): + value = value.item() + if isinstance(value, str): + value = unicodedata.normalize("NFKC", value).strip() + if value: + return str(value) + elif isinstance(value, int) and not isinstance(value, bool): + return int(value) + elif isinstance(value, float) and math.isfinite(value): + return float(value) + raise DatasetValidationError(f"Selected ASR row has no valid {field} provenance.") + + +def _selected_row_provenance(row: dict[str, Any], audio_column: str) -> dict[str, Any]: + source_index = row.get("_winml_source_index") + if isinstance(source_index, np.integer): + source_index = int(source_index) + if not isinstance(source_index, int) or isinstance(source_index, bool) or source_index < 0: + raise DatasetValidationError("Selected ASR row has no valid source index provenance.") + raw_dataset_id = row.get("id") + dataset_id = ( + _normalize_identity_scalar(raw_dataset_id, "dataset ID") + if raw_dataset_id is not None + else None + ) + audio = row.get(audio_column) + audio_path: str | None = None + audio_key: str | int | float | None = None + if isinstance(audio, dict): + path = audio.get("path") + if isinstance(path, str) and path: + audio_path = unicodedata.normalize( + "NFKC", path.replace("\\", "/").rsplit("/", 1)[-1] + ).strip() + key = audio.get("key") + if key not in (None, ""): + audio_key = _normalize_identity_scalar(key, "audio key") + elif isinstance(audio, str) and audio: + audio_path = unicodedata.normalize( + "NFKC", audio.replace("\\", "/").rsplit("/", 1)[-1] + ).strip() + return { + "source_index": source_index, + "dataset_id": dataset_id, + "audio_path": audio_path, + "audio_key": audio_key, + } + + +class WinMLCTCASREvaluator(WinMLEvaluator): + """Evaluate metadata-resolved CTC ASR models with bounded full utterances.""" + + def __init__(self, config: WinMLEvaluationConfig, model: Any) -> None: + _require_soundfile() + mapping = config.dataset.columns_mapping + self._audio_column = mapping.get("input_column", "audio") + self._transcription_column = mapping.get("label_column", "transcription") + model_config = getattr(model, "config", None) + if model_config is None or not _is_ctc_config(model_config): + raise ValueError( + "automatic-speech-recognition evaluation currently supports only " + "metadata-resolved *ForCTC checkpoints." + ) + _validate_ctc_model_inputs(model) + if not config.model_id: + raise ValueError("CTC ASR evaluation requires model_id to load its processor.") + + self.processor = _load_ctc_processor( + config.model_id, + trust_remote_code=config.trust_remote_code, + ) + self.target_lang = _configure_processor_language(self.processor, model_config) + feature_extractor = getattr(self.processor, "feature_extractor", None) + sampling_rate = getattr(feature_extractor, "sampling_rate", None) + if not isinstance(sampling_rate, int) or sampling_rate <= 0: + raise ValueError( + "The checkpoint processor has no valid feature-extractor sampling rate." + ) + self.sampling_rate = sampling_rate + self.blank_token_id = self._resolve_blank_token_id(model_config) + self.tokenizer_vocab_size = self._resolve_tokenizer_vocab_size() + super().__init__(config, model) + + def prepare_pipeline(self) -> Pipeline | None: # type: ignore[override] + """Skip the HF pipeline because CTC decoding drives the model directly.""" + return None + + def prepare_data(self) -> Dataset: + """Load deterministic rows and force explicit decode=False audio values.""" + from datasets import Audio, Dataset, load_dataset, load_from_disk + + ds = self.config.dataset + try: + ds_path = Path(ds.path).expanduser() if ds.path else None + if ds_path and ds_path.is_dir(): + dataset = load_from_disk(str(ds_path)) + else: + dataset = load_dataset( + ds.path, + name=ds.name, + split=ds.split, + streaming=ds.streaming, + revision=ds.revision, + ) + dataset = dataset.cast_column(self._audio_column, Audio(decode=False)) + except Exception as error: + raise DatasetValidationError( + f"Failed to load ASR dataset '{ds.path}' " + f"(name={ds.name!r}, split='{ds.split}'): {error}" + ) from error + + validate_dataset_columns(dataset, "automatic-speech-recognition", ds.columns_mapping) + dataset = dataset.map( + lambda _row, index: {"_winml_source_index": index}, + with_indices=True, + ) + if ds.shuffle: + dataset = dataset.shuffle(seed=ds.seed) + if ds.streaming: + rows = list(dataset.take(ds.samples)) + else: + rows = list(dataset.select(range(min(ds.samples, len(dataset))))) + elif ds.streaming: + rows = list(dataset.take(ds.samples)) + else: + rows = [dataset[index] for index in range(min(ds.samples, len(dataset)))] + return Dataset.from_list(rows) + + def compute(self) -> dict[str, Any]: + """Decode bounded rows and report corpus WER/CER plus accounting.""" + predictions: list[str] = [] + references: list[str] = [] + rejection_reasons: Counter[str] = Counter() + selected_rows = [_selected_row_provenance(row, self._audio_column) for row in self.data] + source_indices = [row["source_index"] for row in selected_rows] + if len(source_indices) != len(set(source_indices)): + raise DatasetValidationError("Selected ASR rows contain duplicate source indices.") + + for row in self.data: + try: + reference = _normalize_transcript(row.get(self._transcription_column)) + if not reference: + raise _RejectedSampleError("normalized transcription is empty") + prediction = self._transcribe(row.get(self._audio_column)) + except _RejectedSampleError as error: + rejection_reasons[str(error)] += 1 + continue + references.append(reference) + predictions.append(prediction) + + rejected = sum(rejection_reasons.values()) + if not predictions: + raise DatasetValidationError( + "No usable ASR samples remain after validation; " + f"processed=0, rejected={rejected}, reasons={dict(rejection_reasons)}" + ) + return { + "wer": _corpus_error_rate(references, predictions, words=True), + "cer": _corpus_error_rate(references, predictions, words=False), + "predictions": predictions, + "references": references, + "requested_samples": self.config.dataset.samples, + "selected_samples": len(selected_rows), + "selected_source_ids": [row["dataset_id"] for row in selected_rows], + "selected_source_indices": source_indices, + "selected_audio_paths": [row["audio_path"] for row in selected_rows], + "selected_audio_keys": [row["audio_key"] for row in selected_rows], + "selected_rows": selected_rows, + "processed_samples": len(predictions), + "skipped_samples": 0, + "rejected_samples": rejected, + "rejection_reasons": dict(rejection_reasons), + "target_lang": self.target_lang, + } + + def _transcribe(self, audio_value: Any) -> str: + waveform, source_rate = _decode_audio(audio_value) + waveform = _resample_audio(waveform, source_rate, self.sampling_rate) + input_names = list((getattr(self.model, "io_config", None) or {}).get("input_names", [])) + if not input_names: + input_names = ["input_values"] + encoded = self.processor( + waveform, + sampling_rate=self.sampling_rate, + return_tensors=("pt" if self.config.runtime == "pytorch" else "np"), + return_attention_mask="attention_mask" in input_names, + ) + model_values = dict(encoded) + if "input_values" not in model_values or model_values["input_values"].ndim != 2: + raise _RejectedSampleError("processor did not produce rank-2 input_values") + + window_size = self._fixed_waveform_size(input_names) + sample_count = model_values["input_values"].shape[1] + window_count = 1 if window_size is None else math.ceil(sample_count / window_size) + if ( + window_size is not None + and sample_count % window_size + and "attention_mask" not in input_names + ): + raise _RejectedSampleError( + "fixed-size input_values require attention_mask for a partial final window" + ) + if window_count > _MAX_WINDOWS_PER_UTTERANCE: + raise _RejectedSampleError( + f"utterance requires {window_count} windows; cap is {_MAX_WINDOWS_PER_UTTERANCE}" + ) + + predicted_ids: list[int] = [] + for window_index in range(window_count): + inputs = self._window_inputs( + model_values, + input_names, + window_index=window_index, + window_size=window_size, + ) + outputs = self.model(**inputs) + logits = self._extract_logits(outputs) + if logits.ndim != 3 or logits.shape[0] != 1: + raise ValueError( + f"CTC logits must have shape [1, frames, vocab], got {logits.shape}." + ) + if logits.shape[-1] != self.tokenizer_vocab_size: + raise ValueError( + f"CTC output vocabulary {logits.shape[-1]} does not match active tokenizer " + f"vocabulary {self.tokenizer_vocab_size}." + ) + if window_size is not None: + valid_samples = min(window_size, sample_count - window_index * window_size) + valid_frames = self._valid_logit_frames( + valid_samples, + window_size, + logits.shape[1], + ) + logits = logits[:, :valid_frames] + if predicted_ids: + predicted_ids.append(self.blank_token_id) + predicted_ids.extend(np.argmax(logits, axis=-1)[0].astype(int).tolist()) + + tokenizer = getattr(self.processor, "tokenizer", self.processor) + decoded = tokenizer.batch_decode([predicted_ids]) + if not isinstance(decoded, list) or len(decoded) != 1: + raise ValueError("CTC tokenizer.batch_decode must return one transcript per utterance.") + return _normalize_transcript(decoded[0]) + + def _fixed_waveform_size(self, input_names: list[str]) -> int | None: + io_config = getattr(self.model, "io_config", None) or {} + shapes = io_config.get("input_shapes", []) + try: + input_index = input_names.index("input_values") + shape = shapes[input_index] + except (ValueError, IndexError): + return None + return shape[1] if len(shape) == 2 and isinstance(shape[1], int) else None + + @staticmethod + def _window_inputs( + arrays: dict[str, Any], + input_names: list[str], + *, + window_index: int, + window_size: int | None, + ) -> dict[str, Any]: + inputs: dict[str, Any] = {} + for name in input_names: + if name not in arrays: + raise ValueError(f"Processor did not produce required ONNX input {name!r}.") + value = arrays[name] + if window_size is None: + inputs[name] = value + continue + start = window_index * window_size + chunk = value[:, start : start + window_size] + if chunk.shape[1] < window_size: + padding = window_size - chunk.shape[1] + if isinstance(chunk, np.ndarray): + chunk = np.pad(chunk, ((0, 0), (0, padding))) + else: + import torch.nn.functional as functional + + chunk = functional.pad(chunk, (0, padding)) + inputs[name] = chunk + return inputs + + def _valid_logit_frames( + self, + valid_samples: int, + padded_samples: int, + total_frames: int, + ) -> int: + """Map valid samples in a padded window to valid CTC output frames.""" + if valid_samples >= padded_samples: + return total_frames + model_config = getattr(self.model, "config", None) + conv_kernels = getattr(model_config, "conv_kernel", None) + conv_strides = getattr(model_config, "conv_stride", None) + if ( + isinstance(conv_kernels, list) + and isinstance(conv_strides, list) + and len(conv_kernels) == len(conv_strides) + ): + output_length = valid_samples + for kernel, stride in zip(conv_kernels, conv_strides, strict=True): + if not isinstance(kernel, int) or not isinstance(stride, int) or stride <= 0: + break + output_length = max(0, (output_length - kernel) // stride + 1) + else: + return min(total_frames, output_length) + return min(total_frames, math.ceil(total_frames * valid_samples / padded_samples)) + + @staticmethod + def _extract_logits(outputs: Any) -> np.ndarray: + logits = ( + outputs.get("logits") if isinstance(outputs, dict) else getattr(outputs, "logits", None) + ) + if logits is None: + raise ValueError("CTC model output does not contain logits.") + if hasattr(logits, "detach"): + logits = logits.detach().cpu().numpy() + return np.asarray(logits) + + def _resolve_blank_token_id(self, model_config: Any) -> int: + tokenizer = getattr(self.processor, "tokenizer", self.processor) + blank_token_id = getattr(tokenizer, "pad_token_id", None) + if blank_token_id is None: + blank_token_id = getattr(model_config, "pad_token_id", None) + if not isinstance(blank_token_id, int): + raise TypeError("The active CTC tokenizer has no integer blank/pad token ID.") + return blank_token_id + + def _resolve_tokenizer_vocab_size(self) -> int: + tokenizer = getattr(self.processor, "tokenizer", self.processor) + get_vocab = getattr(tokenizer, "get_vocab", None) + vocabulary = get_vocab() if callable(get_vocab) else getattr(tokenizer, "vocab", None) + if not isinstance(vocabulary, dict) or not vocabulary: + raise TypeError("The active CTC tokenizer has no vocabulary metadata.") + return len(vocabulary) + + +__all__ = ["WinMLCTCASREvaluator"] diff --git a/src/winml/modelkit/eval/evaluate.py b/src/winml/modelkit/eval/evaluate.py index 42a9ab8a5..b6d65ba05 100644 --- a/src/winml/modelkit/eval/evaluate.py +++ b/src/winml/modelkit/eval/evaluate.py @@ -100,6 +100,8 @@ def _select_model_loader(config: WinMLEvaluationConfig) -> _ModelLoaderKind: "winml.modelkit.eval.mask_generation_evaluator:WinMLMaskGenerationEvaluator", "text-generation": "winml.modelkit.eval.text_generation_evaluator:WinMLTextGenerationEvaluator", + "automatic-speech-recognition": + "winml.modelkit.eval.ctc_asr_evaluator:WinMLCTCASREvaluator", } # fmt: on @@ -283,6 +285,18 @@ def _validate_pytorch_runtime_config(config: WinMLEvaluationConfig) -> None: "split": "test", "columns_mapping": {"input_column": "text"}, }, + "automatic-speech-recognition": { + "path": "google/fleurs", + "name": "en_us", + "split": "validation", + "samples": 2, + "shuffle": False, + "revision": "70bb2e84b976b7e960aa89f1c648e09c59f894dd", + "columns_mapping": { + "input_column": "audio", + "label_column": "transcription", + }, + }, } diff --git a/src/winml/modelkit/loader/config.py b/src/winml/modelkit/loader/config.py index b2bb053b5..36dde3baa 100644 --- a/src/winml/modelkit/loader/config.py +++ b/src/winml/modelkit/loader/config.py @@ -47,6 +47,7 @@ class WinMLLoaderConfig: Requires trust_remote_code=True for security. trust_remote_code: Whether to trust remote/custom code. Required when using user_script. + target_lang: Language adapter to load before exporting an MMS CTC model. Example: # Standard usage with auto-detection @@ -73,6 +74,7 @@ class WinMLLoaderConfig: module_path: str | None = None user_script: str | None = None trust_remote_code: bool = False + target_lang: str | None = None def to_dict(self) -> dict[str, Any]: """Serialize to dictionary. @@ -93,6 +95,8 @@ def to_dict(self) -> dict[str, Any]: result["user_script"] = self.user_script if self.trust_remote_code: result["trust_remote_code"] = self.trust_remote_code + if self.target_lang is not None: + result["target_lang"] = self.target_lang return result @classmethod @@ -112,6 +116,7 @@ def from_dict(cls, data: dict[str, Any]) -> WinMLLoaderConfig: module_path=data.get("module_path"), user_script=data.get("user_script"), trust_remote_code=data.get("trust_remote_code", False), + target_lang=data.get("target_lang"), ) @@ -267,6 +272,12 @@ def resolve_loader_config( model_class=resolved_class.__name__, model_type=resolved_model_type, trust_remote_code=trust_remote_code, + target_lang=( + getattr(resolved_hf_config, "target_lang", None) + if resolved_task == "automatic-speech-recognition" + and getattr(resolved_hf_config, "adapter_attn_dim", None) + else None + ), ) return loader_config, resolved_hf_config, resolved_class, resolution diff --git a/src/winml/modelkit/loader/hf.py b/src/winml/modelkit/loader/hf.py index caf8dc615..42b0b95b3 100644 --- a/src/winml/modelkit/loader/hf.py +++ b/src/winml/modelkit/loader/hf.py @@ -148,6 +148,7 @@ def load_hf_model( *, torch_dtype: Any | None = None, attn_implementation: str | None = None, + target_lang: str | None = None, ) -> tuple[nn.Module, PretrainedConfig, str]: """Load, detect task, and prepare HuggingFace model. @@ -179,6 +180,7 @@ def load_hf_model( Pass ``"auto"`` to preserve the checkpoint's stored dtype. attn_implementation: Optional Transformers attention implementation forwarded to ``from_pretrained`` before attention modules are created. + target_lang: Optional MMS language adapter to activate before export. Returns: Tuple of (model, hf_config, task) @@ -296,6 +298,19 @@ def load_hf_model( if attn_implementation is not None: load_kwargs["attn_implementation"] = attn_implementation model = loader_cls.from_pretrained(model_name_or_path, **load_kwargs) + if target_lang is not None: + load_adapter = getattr(model, "load_adapter", None) + if not callable(load_adapter): + raise ValueError( + f"Model {model.__class__.__name__} cannot load requested language adapter " + f"{target_lang!r}." + ) + try: + load_adapter(target_lang) + except Exception as error: + raise ValueError(f"Could not load language adapter {target_lang!r}.") from error + model.config.target_lang = target_lang + hf_config.target_lang = target_lang # [5] Export Preparation model.eval() diff --git a/src/winml/modelkit/models/winml/__init__.py b/src/winml/modelkit/models/winml/__init__.py index 251c40dad..dc54eb2f3 100644 --- a/src/winml/modelkit/models/winml/__init__.py +++ b/src/winml/modelkit/models/winml/__init__.py @@ -50,6 +50,7 @@ "feature-extraction": "WinMLModelForFeatureExtraction", "sentence-similarity": "WinMLModelForFeatureExtraction", "image-feature-extraction": "WinMLModelForFeatureExtraction", + "automatic-speech-recognition": "WinMLModelForGenericTask", } # Level 2: (model_type, task) -> Specialized class (exceptions only) diff --git a/src/winml/modelkit/utils/eval_utils.py b/src/winml/modelkit/utils/eval_utils.py index f85c672b4..d340ccc7e 100644 --- a/src/winml/modelkit/utils/eval_utils.py +++ b/src/winml/modelkit/utils/eval_utils.py @@ -448,6 +448,23 @@ class TaskSchema: ), ) +_AUTOMATIC_SPEECH_RECOGNITION_SCHEMA = TaskSchema( + columns=( + SchemaItem( + "input_column", + "audio bytes or path decoded by SoundFile", + default="audio", + remap_hint="", + ), + SchemaItem( + "label_column", + "reference transcription", + default="transcription", + remap_hint="", + ), + ), +) + TASK_SCHEMAS: dict[str, TaskSchema] = { "image-classification": _IMAGE_CLASSIFICATION_SCHEMA, "text-classification": _TEXT_CLASSIFICATION_SCHEMA, @@ -468,6 +485,7 @@ class TaskSchema: "keypoint-detection": _KEYPOINT_DETECTION_SCHEMA, "mask-generation": _MASK_GENERATION_SCHEMA, "text-generation": _TEXT_GENERATION_SCHEMA, + "automatic-speech-recognition": _AUTOMATIC_SPEECH_RECOGNITION_SCHEMA, } diff --git a/tests/unit/commands/test_build.py b/tests/unit/commands/test_build.py index 37f6bde79..d894c489c 100644 --- a/tests/unit/commands/test_build.py +++ b/tests/unit/commands/test_build.py @@ -2020,6 +2020,58 @@ def test_quant_model_type_carried_into_quantize_stage( assert config.quant.model_type == "qwen3_transformer_only" mock_quantize.assert_called_once() + @patch("winml.modelkit.commands.build._run_compile_stage") + @patch("winml.modelkit.commands.build._run_quantize_stage") + @patch("winml.modelkit.commands.build._run_optimize_stage") + @patch("winml.modelkit.commands.build._show_io") + @patch("winml.modelkit.utils.console.StageLive") + @patch("winml.modelkit.export.export_onnx") + @patch("winml.modelkit.build.hf._load_model") + def test_skip_optimize_reaches_hf_optimize_stage( + self, + mock_load_model: MagicMock, + mock_export_onnx: MagicMock, + mock_stage_live: MagicMock, + mock_show_io: MagicMock, + mock_optimize: MagicMock, + mock_quantize: MagicMock, + mock_compile: MagicMock, + tmp_path: Path, + ) -> None: + from winml.modelkit.commands.build import _build_hf_pipeline + + mock_stage_live.return_value.__enter__ = MagicMock(return_value=MagicMock()) + mock_stage_live.return_value.__exit__ = MagicMock(return_value=False) + mock_load_model.return_value = MagicMock() + exported = tmp_path / "export.onnx" + mock_optimize.return_value = (exported, None) + mock_quantize.return_value = exported + mock_compile.side_effect = RuntimeError("stop-after-stages") + + config = MagicMock() + config.loader.model_type = "wav2vec2" + config.loader.task = "automatic-speech-recognition" + config.loader.model_class = "AutoModelForCTC" + config.export = MagicMock() + config.quant = None + config.skip_optimize = False + config.to_dict.return_value = {} + + with pytest.raises(RuntimeError, match="stop-after-stages"): + _build_hf_pipeline( + config=config, + model_id="org/ctc-model", + output_dir=tmp_path / "out", + rebuild=True, + cache_key=None, + ep=None, + device="cpu", + extra_kwargs={"skip_optimize": True}, + preloaded_hf_config=None, + ) + + assert mock_optimize.call_args.kwargs["skip_optimize"] is True + class TestBuildEpResolution: """--ep forwarding into config generation + the compile EP-availability gate.""" diff --git a/tests/unit/eval/test_ctc_asr_evaluator.py b/tests/unit/eval/test_ctc_asr_evaluator.py new file mode 100644 index 000000000..7b9bd476d --- /dev/null +++ b/tests/unit/eval/test_ctc_asr_evaluator.py @@ -0,0 +1,975 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- + +from __future__ import annotations + +from io import BytesIO +from types import SimpleNamespace +from typing import ClassVar +from unittest.mock import MagicMock, patch + +import numpy as np +import pytest +import soundfile as sf +import torch + +from winml.modelkit.eval.ctc_asr_evaluator import ( + WinMLCTCASREvaluator, + _configure_processor_language, + _corpus_error_rate, + _decode_audio, + _is_ctc_config, + _load_ctc_processor, + _normalize_transcript, + _RejectedSampleError, + _require_soundfile, + _resample_audio, + _validate_ctc_model_inputs, +) +from winml.modelkit.utils.eval_utils import TASK_SCHEMAS, DatasetValidationError + + +class _Tokenizer: + pad_token_id = 0 + + def __init__(self, target_lang: str | None = None, vocab_size: int = 3) -> None: + self.target_lang = target_lang + self._vocab = {str(index): index for index in range(vocab_size)} + + def get_vocab(self) -> dict[str, int]: + return self._vocab + + def set_target_lang(self, target_lang: str) -> None: + self.target_lang = target_lang + + def batch_decode(self, sequences): + decoded_ids = list(sequences[0]) + collapsed: list[int] = [] + previous = None + for token_id in decoded_ids: + if token_id != 0 and token_id != previous: + collapsed.append(token_id) + previous = token_id + return [" ".join(str(token_id) for token_id in collapsed)] + + +class _Processor: + def __init__( + self, + *, + target_lang: str | None = None, + vocab_size: int = 3, + return_attention_mask: bool = False, + ) -> None: + self.tokenizer = _Tokenizer(target_lang, vocab_size) + self.feature_extractor = SimpleNamespace(sampling_rate=16_000) + self.return_attention_mask = return_attention_mask + + def __call__(self, waveform, **kwargs): + values = np.asarray(waveform, dtype=np.float32)[None, :] + if kwargs.get("return_tensors") == "pt": + encoded = {"input_values": torch.from_numpy(values)} + if self.return_attention_mask: + encoded["attention_mask"] = torch.ones_like( + encoded["input_values"], dtype=torch.int64 + ) + return encoded + encoded = {"input_values": values} + if self.return_attention_mask: + encoded["attention_mask"] = np.ones_like(values, dtype=np.int64) + return encoded + + def batch_decode(self, _sequences): + raise AssertionError("Greedy decoding must use the tokenizer, not the LM processor") + + +def _wav_bytes(samples: np.ndarray, sampling_rate: int = 16_000) -> bytes: + buffer = BytesIO() + sf.write(buffer, samples, sampling_rate, format="WAV", subtype="FLOAT") + return buffer.getvalue() + + +def _evaluator( + *, + input_shape: list[object], + processor: _Processor | None = None, + with_attention_mask: bool = False, +): + from winml.modelkit.eval import DatasetConfig, WinMLEvaluationConfig + + evaluator = WinMLCTCASREvaluator.__new__(WinMLCTCASREvaluator) + evaluator.processor = processor or _Processor() + evaluator.sampling_rate = 16_000 + evaluator.blank_token_id = 0 + evaluator.tokenizer_vocab_size = 3 + evaluator.target_lang = evaluator.processor.tokenizer.target_lang + evaluator._audio_column = "audio" + evaluator._transcription_column = "transcription" + evaluator.config = WinMLEvaluationConfig( + task="automatic-speech-recognition", + dataset=DatasetConfig(path="dataset", samples=2, shuffle=False), + ) + evaluator.model = MagicMock() + input_names = ["input_values", "attention_mask"] if with_attention_mask else ["input_values"] + evaluator.model.io_config = { + "input_names": input_names, + "input_shapes": [input_shape for _ in input_names], + } + return evaluator + + +def test_asr_registry_schema_and_generic_model_route() -> None: + from winml.modelkit.eval import WinMLEvaluationConfig, get_evaluator_class + from winml.modelkit.inference import TASK_REGISTRY + from winml.modelkit.models.winml import get_winml_class + from winml.modelkit.models.winml.base import WinMLModelForGenericTask + + config = WinMLEvaluationConfig(task="automatic-speech-recognition") + assert get_evaluator_class(config) is WinMLCTCASREvaluator + assert "automatic-speech-recognition" in TASK_SCHEMAS + assert get_winml_class("wav2vec2", "automatic-speech-recognition") is WinMLModelForGenericTask + assert TASK_REGISTRY["automatic-speech-recognition"].user_inputs[0].type == "audio" + + +def test_asr_default_dataset_is_pinned_and_bounded() -> None: + from winml.modelkit.eval.evaluate import _DEFAULT_DATASETS + + default = _DEFAULT_DATASETS["automatic-speech-recognition"] + assert default["path"] == "google/fleurs" + assert default["name"] == "en_us" + assert default["revision"] == "70bb2e84b976b7e960aa89f1c648e09c59f894dd" + assert default["samples"] == 2 + assert default["shuffle"] is False + + +@pytest.mark.parametrize("architecture", ["Wav2Vec2ForCTC", "HubertForCTC", "AutoModelForCTC"]) +def test_ctc_architectures_are_metadata_driven(architecture: str) -> None: + assert _is_ctc_config(SimpleNamespace(architectures=[architecture])) + + +def test_non_ctc_asr_fails_closed() -> None: + assert not _is_ctc_config(SimpleNamespace(architectures=["WhisperForConditionalGeneration"])) + + from winml.modelkit.eval import DatasetConfig, WinMLEvaluationConfig + + model = MagicMock() + model.config = SimpleNamespace(architectures=["WhisperForConditionalGeneration"]) + config = WinMLEvaluationConfig( + model_id="org/whisper-model", + task="automatic-speech-recognition", + dataset=DatasetConfig(path="dataset"), + ) + with pytest.raises(ValueError, match=r"only metadata-resolved \*ForCTC"): + WinMLCTCASREvaluator(config, model) + + +@pytest.mark.parametrize( + ("input_names", "input_shapes", "message"), + [ + (["input_features"], [[1, 80, 3000]], "input contract"), + (["input_values"], [[1, 16000, 1]], "rank-2 input_values"), + ], +) +def test_ctc_model_input_contract_fails_closed( + input_names: list[str], + input_shapes: list[list[int]], + message: str, +) -> None: + model = SimpleNamespace(io_config={"input_names": input_names, "input_shapes": input_shapes}) + + with pytest.raises(ValueError, match=message): + _validate_ctc_model_inputs(model) + + +def test_ctc_model_input_contract_accepts_optional_attention_mask() -> None: + model = SimpleNamespace( + io_config={ + "input_names": ["input_values", "attention_mask"], + "input_shapes": [[1, 16000], [1, 16000]], + } + ) + + _validate_ctc_model_inputs(model) + + +def test_constructor_initializes_base_evaluator_state() -> None: + from winml.modelkit.eval import DatasetConfig, WinMLEvaluationConfig + + processor = _Processor(target_lang="eng") + model = MagicMock() + model.config = SimpleNamespace(architectures=["Wav2Vec2ForCTC"], pad_token_id=0) + config = WinMLEvaluationConfig( + model_id="org/ctc-model", + task="automatic-speech-recognition", + dataset=DatasetConfig(path="dataset"), + ) + prepared_data = object() + + with ( + patch("transformers.AutoProcessor.from_pretrained", return_value=processor), + patch.object(WinMLCTCASREvaluator, "prepare_data", return_value=prepared_data), + ): + evaluator = WinMLCTCASREvaluator(config, model) + + assert evaluator.model is model + assert evaluator.config is config + assert evaluator.data is prepared_data + assert evaluator.pipe is None + + +def test_missing_soundfile_fails_before_base_initialization() -> None: + import builtins + + real_import = builtins.__import__ + + def reject_soundfile(name, *args, **kwargs): + if name == "soundfile": + raise ImportError("missing soundfile") + return real_import(name, *args, **kwargs) + + with ( + patch("builtins.__import__", side_effect=reject_soundfile), + pytest.raises(ImportError, match=r"winml-cli\[audio\]"), + ): + _require_soundfile() + + +@pytest.mark.parametrize("streaming", [False, True]) +def test_prepare_data_preserves_source_order_with_duplicate_ids(streaming: bool) -> None: + from winml.modelkit.eval import DatasetConfig, WinMLEvaluationConfig + + rows = [ + {"id": 1525, "audio": {"path": "first.wav"}, "transcription": "first"}, + {"id": 1657, "audio": {"path": "second.wav"}, "transcription": "second"}, + {"id": 1525, "audio": {"path": "third.wav"}, "transcription": "third"}, + ] + indexed_rows = [{**row, "_winml_source_index": index} for index, row in enumerate(rows)] + dataset = MagicMock(column_names=["id", "audio", "transcription"]) + dataset.cast_column.return_value = dataset + dataset.__len__.return_value = len(rows) + + def add_source_indices(*_args, **_kwargs): + dataset.__getitem__.side_effect = indexed_rows.__getitem__ + dataset.take.return_value = indexed_rows + return dataset + + dataset.map.side_effect = add_source_indices + evaluator = WinMLCTCASREvaluator.__new__(WinMLCTCASREvaluator) + evaluator._audio_column = "audio" + evaluator.config = WinMLEvaluationConfig( + task="automatic-speech-recognition", + dataset=DatasetConfig(path="dataset", samples=3, shuffle=False, streaming=streaming), + ) + + with patch("datasets.load_dataset", return_value=dataset): + selected = evaluator.prepare_data() + + assert selected["id"] == [1525, 1657, 1525] + assert selected["_winml_source_index"] == [0, 1, 2] + assert selected["transcription"] == ["first", "second", "third"] + dataset.sort.assert_not_called() + + +@pytest.mark.parametrize("streaming", [False, True]) +def test_prepare_data_applies_seeded_shuffle_before_bounded_selection(streaming: bool) -> None: + from datasets import Audio, Dataset + + from winml.modelkit.eval import DatasetConfig, WinMLEvaluationConfig + + rows = [ + { + "id": index, + "audio": f"{index}.wav", + "transcription": str(index), + } + for index in range(10) + ] + dataset = Dataset.from_list(rows) + if streaming: + dataset = dataset.cast_column("audio", Audio(decode=False)).to_iterable_dataset() + evaluator = WinMLCTCASREvaluator.__new__(WinMLCTCASREvaluator) + evaluator._audio_column = "audio" + evaluator.config = WinMLEvaluationConfig( + task="automatic-speech-recognition", + dataset=DatasetConfig( + path="dataset", + samples=4, + shuffle=True, + seed=123, + streaming=streaming, + ), + ) + + with patch("datasets.load_dataset", return_value=dataset): + selected = evaluator.prepare_data() + + assert selected["id"] == [7, 4, 0, 2] + assert selected["_winml_source_index"] == [7, 4, 0, 2] + + +@pytest.mark.parametrize("streaming", [False, True]) +def test_prepare_data_shuffle_is_repeatable_and_seed_discriminating(streaming: bool) -> None: + from datasets import Dataset, Features, IterableDataset, Value + + from winml.modelkit.eval import DatasetConfig, WinMLEvaluationConfig + + rows = [ + {"id": index, "audio": f"{index}.wav", "transcription": str(index)} for index in range(10) + ] + features = Features( + { + "id": Value("int64"), + "audio": Value("string"), + "transcription": Value("string"), + } + ) + + def select(seed: int) -> list[int]: + dataset = ( + IterableDataset.from_generator(lambda: iter(rows), features=features) + if streaming + else Dataset.from_list(rows) + ) + evaluator = WinMLCTCASREvaluator.__new__(WinMLCTCASREvaluator) + evaluator._audio_column = "audio" + evaluator.config = WinMLEvaluationConfig( + task="automatic-speech-recognition", + dataset=DatasetConfig( + path="dataset", + samples=4, + shuffle=True, + seed=seed, + streaming=streaming, + ), + ) + with ( + patch.object(dataset, "cast_column", return_value=dataset), + patch("datasets.load_dataset", return_value=dataset), + ): + return list(evaluator.prepare_data()["_winml_source_index"]) + + assert select(123) == select(123) == [7, 4, 0, 2] + assert select(456) != select(123) + + +def test_prepare_data_map_shuffle_is_exact_beyond_stream_buffer() -> None: + from datasets import Dataset + + from winml.modelkit.eval import DatasetConfig, WinMLEvaluationConfig + + rows = [ + {"id": index, "audio": f"{index}.wav", "transcription": str(index)} + for index in range(1_200) + ] + dataset = Dataset.from_list(rows) + + def select() -> list[int]: + evaluator = WinMLCTCASREvaluator.__new__(WinMLCTCASREvaluator) + evaluator._audio_column = "audio" + evaluator.config = WinMLEvaluationConfig( + task="automatic-speech-recognition", + dataset=DatasetConfig( + path="dataset", + samples=4, + shuffle=True, + seed=123, + streaming=False, + ), + ) + with ( + patch.object(dataset, "cast_column", return_value=dataset), + patch("datasets.load_dataset", return_value=dataset), + ): + return list(evaluator.prepare_data()["_winml_source_index"]) + + indexed = dataset.map( + lambda _row, index: {"_winml_source_index": index}, + with_indices=True, + ) + expected = list(indexed.shuffle(seed=123).select(range(4))["_winml_source_index"]) + assert select() == expected + + +def test_prepare_data_streaming_shuffle_is_bounded() -> None: + from datasets import Features, IterableDataset, Value + + from winml.modelkit.eval import DatasetConfig, WinMLEvaluationConfig + + yielded = 0 + + def rows(): + nonlocal yielded + for index in range(10_000): + yielded += 1 + if yielded > 1_100: + raise AssertionError("streaming shuffle consumed an unbounded source") + yield {"id": index, "audio": f"{index}.wav", "transcription": str(index)} + + dataset = IterableDataset.from_generator( + rows, + features=Features( + { + "id": Value("int64"), + "audio": Value("string"), + "transcription": Value("string"), + } + ), + ) + evaluator = WinMLCTCASREvaluator.__new__(WinMLCTCASREvaluator) + evaluator._audio_column = "audio" + evaluator.config = WinMLEvaluationConfig( + task="automatic-speech-recognition", + dataset=DatasetConfig( + path="dataset", + samples=4, + shuffle=True, + seed=123, + streaming=True, + ), + ) + + with ( + patch.object(dataset, "cast_column", return_value=dataset), + patch("datasets.load_dataset", return_value=dataset), + ): + selected = evaluator.prepare_data() + + assert len(selected) == 4 + assert yielded <= 1_100 + assert all(0 <= index < yielded for index in selected["_winml_source_index"]) + + +def test_compute_rejects_invalid_present_id_selected_by_shuffle_before_inference() -> None: + from datasets import Dataset + + from winml.modelkit.eval import DatasetConfig, WinMLEvaluationConfig + + rows = [ + { + "id": " " if index == 7 else str(index), + "audio": f"{index}.wav", + "transcription": str(index), + } + for index in range(10) + ] + evaluator = WinMLCTCASREvaluator.__new__(WinMLCTCASREvaluator) + evaluator._audio_column = "audio" + evaluator._transcription_column = "transcription" + evaluator.config = WinMLEvaluationConfig( + task="automatic-speech-recognition", + dataset=DatasetConfig(path="dataset", samples=1, shuffle=True, seed=123), + ) + with patch("datasets.load_dataset", return_value=Dataset.from_list(rows)): + evaluator.data = evaluator.prepare_data() + evaluator._transcribe = MagicMock() + + assert evaluator.data["_winml_source_index"] == [7] + with pytest.raises(DatasetValidationError, match="dataset ID"): + evaluator.compute() + evaluator._transcribe.assert_not_called() + + +def test_soundfile_decode_mixes_stereo_to_mono_float32() -> None: + stereo = np.column_stack( + [np.linspace(-1, 1, 32, dtype=np.float32), np.linspace(1, -1, 32, dtype=np.float32)] + ) + waveform, sampling_rate = _decode_audio({"bytes": _wav_bytes(stereo), "path": None}) + assert sampling_rate == 16_000 + assert waveform.shape == (32,) + assert waveform.dtype == np.float32 + np.testing.assert_allclose(waveform, 0.0, atol=1e-6) + + +def test_soundfile_decode_rejects_invalid_schema() -> None: + with pytest.raises(_RejectedSampleError, match="neither bytes nor path"): + _decode_audio({"bytes": None, "path": None}) + + +def test_resample_uses_processor_rate_only_when_needed() -> None: + waveform = np.arange(80, dtype=np.float32) + assert _resample_audio(waveform, 16_000, 16_000) is waveform + assert _resample_audio(waveform, 8_000, 16_000).shape == (160,) + + +def test_optional_lm_dependency_falls_back_to_plain_wav2vec2_processor() -> None: + feature_extractor = object() + tokenizer = object() + processor = object() + missing_lm = ImportError( + "Wav2Vec2ProcessorWithLM requires the pyctcdecode library but it was not found" + ) + + with ( + patch("transformers.AutoProcessor.from_pretrained", side_effect=missing_lm), + patch("transformers.AutoFeatureExtractor.from_pretrained", return_value=feature_extractor), + patch("transformers.AutoTokenizer.from_pretrained", return_value=tokenizer), + patch("transformers.Wav2Vec2Processor", return_value=processor) as processor_class, + ): + assert _load_ctc_processor("org/ctc-model", trust_remote_code=False) is processor + + processor_class.assert_called_once_with( + feature_extractor=feature_extractor, + tokenizer=tokenizer, + ) + + +def test_available_lm_processor_is_preserved() -> None: + processor = object() + with patch("transformers.AutoProcessor.from_pretrained", return_value=processor): + assert _load_ctc_processor("org/lm-model", trust_remote_code=False) is processor + + +def test_unrelated_processor_import_error_fails_closed() -> None: + with ( + patch( + "transformers.AutoProcessor.from_pretrained", + side_effect=ImportError("unsupported custom processor dependency"), + ), + pytest.raises(ImportError, match="unsupported custom processor dependency"), + ): + _load_ctc_processor("org/unsupported-model", trust_remote_code=False) + + +def test_incompatible_greedy_processor_components_fail_clearly() -> None: + missing_lm = ImportError( + "Wav2Vec2ProcessorWithLM requires the pyctcdecode library but it was not found" + ) + with ( + patch("transformers.AutoProcessor.from_pretrained", side_effect=missing_lm), + patch("transformers.AutoFeatureExtractor.from_pretrained", return_value=object()), + patch("transformers.AutoTokenizer.from_pretrained", return_value=object()), + patch("transformers.Wav2Vec2Processor", side_effect=TypeError("incompatible components")), + pytest.raises(ValueError, match="cannot form a greedy CTC processor"), + ): + _load_ctc_processor("org/unsupported-model", trust_remote_code=False) + + +def test_ordinary_wav2vec2_tokenizer_accepts_null_target_language() -> None: + processor = _Processor(target_lang=None) + assert _configure_processor_language(processor, SimpleNamespace(adapter_attn_dim=None)) is None + + +def test_mms_language_selection_uses_adapter_metadata() -> None: + processor = _Processor(target_lang="eng") + active = _configure_processor_language( + processor, + SimpleNamespace(target_lang="deu", adapter_attn_dim=16), + ) + assert active == "deu" + assert processor.tokenizer.target_lang == "deu" + + +def test_mms_published_language_is_preserved_without_config_override() -> None: + processor = _Processor(target_lang="eng") + assert _configure_processor_language(processor, SimpleNamespace(adapter_attn_dim=16)) == "eng" + + +def test_invalid_mms_adapter_fails_closed() -> None: + processor = _Processor(target_lang="eng") + processor.tokenizer.set_target_lang = MagicMock(side_effect=ValueError("invalid adapter")) + with pytest.raises(ValueError, match="cannot select requested adapter language 'invalid'"): + _configure_processor_language( + processor, + SimpleNamespace(target_lang="invalid", adapter_attn_dim=16), + ) + + +def test_mms_adapter_metadata_requires_an_active_language() -> None: + processor = _Processor(target_lang=None) + with pytest.raises(ValueError, match="supports language adapters but no language is active"): + _configure_processor_language(processor, SimpleNamespace(adapter_attn_dim=16)) + + +def test_static_audio_windows_pad_and_insert_blank_for_processor_ctc_decode() -> None: + processor = _Processor(return_attention_mask=True) + evaluator = _evaluator( + input_shape=[1, 4], + processor=processor, + with_attention_mask=True, + ) + first_logits = np.array([[[0, 5, 0], [0, 5, 0]]], dtype=np.float32) + final_logits = np.array([[[0, 5, 0], [0, 0, 5]]], dtype=np.float32) + evaluator.model.side_effect = [ + {"logits": first_logits}, + {"logits": final_logits}, + ] + audio = {"bytes": _wav_bytes(np.arange(6, dtype=np.float32)), "path": None} + + assert evaluator._transcribe(audio) == "1 1" + assert len(evaluator.model.call_args_list) == 2 + final_input = evaluator.model.call_args_list[1].kwargs["input_values"] + assert final_input.shape == (1, 4) + np.testing.assert_array_equal(final_input[0, 2:], np.zeros(2)) + final_mask = evaluator.model.call_args_list[1].kwargs["attention_mask"] + np.testing.assert_array_equal(final_mask, [[1, 1, 0, 0]]) + + +def test_static_audio_rejects_unmasked_partial_final_window() -> None: + evaluator = _evaluator(input_shape=[1, 4]) + audio = {"bytes": _wav_bytes(np.arange(6, dtype=np.float32)), "path": None} + + with pytest.raises(_RejectedSampleError, match="require attention_mask"): + evaluator._transcribe(audio) + + evaluator.model.assert_not_called() + + +def test_pytorch_runtime_preserves_processor_tensors_for_native_model() -> None: + evaluator = _evaluator(input_shape=[1, "samples"]) + evaluator.config.runtime = "pytorch" + + class NativeCTCModel: + io_config: ClassVar[dict[str, list[object]]] = {} + + def __call__(self, *, input_values): + assert isinstance(input_values, torch.Tensor) + return SimpleNamespace(logits=torch.tensor([[[0.0, 5.0, 0.0]]])) + + evaluator.model = NativeCTCModel() + audio = {"bytes": _wav_bytes(np.ones(4, dtype=np.float32)), "path": None} + + assert evaluator._transcribe(audio) == "1" + + +def test_dynamic_audio_runs_full_utterance_once() -> None: + evaluator = _evaluator(input_shape=[1, "samples"]) + evaluator.model.return_value = { + "logits": np.array([[[0, 5, 0], [5, 0, 0], [0, 0, 5]]], dtype=np.float32) + } + audio = {"bytes": _wav_bytes(np.arange(9, dtype=np.float32)), "path": None} + + assert evaluator._transcribe(audio) == "1 2" + assert evaluator.model.call_count == 1 + assert evaluator.model.call_args.kwargs["input_values"].shape == (1, 9) + + +def test_static_audio_rejects_more_than_64_windows_before_inference() -> None: + evaluator = _evaluator(input_shape=[1, 1]) + audio = {"bytes": _wav_bytes(np.arange(65, dtype=np.float32)), "path": None} + with pytest.raises(_RejectedSampleError, match="cap is 64"): + evaluator._transcribe(audio) + evaluator.model.assert_not_called() + + +def test_vocab_mismatch_fails_closed() -> None: + evaluator = _evaluator(input_shape=[1, "samples"]) + evaluator.model.return_value = {"logits": np.zeros((1, 2, 4), dtype=np.float32)} + audio = {"bytes": _wav_bytes(np.ones(4, dtype=np.float32)), "path": None} + with pytest.raises(ValueError, match="does not match active tokenizer"): + evaluator._transcribe(audio) + + +def test_transcript_normalization_and_exact_wer_cer() -> None: + assert _normalize_transcript(" hello\t world ") == "hello world" + assert _corpus_error_rate(["hello world"], ["hello there"], words=True) == 0.5 + assert _corpus_error_rate(["abc"], ["adc"], words=False) == pytest.approx(1 / 3) + + +def test_compute_preserves_accounting_and_rejection_reasons() -> None: + evaluator = _evaluator(input_shape=[1, "samples"]) + evaluator.data = [ + { + "_winml_source_index": 0, + "id": 1525, + "audio": {"path": r"C:\scratch\10018492969996036091.wav", "key": "first"}, + "transcription": " hello world ", + }, + { + "_winml_source_index": 1, + "id": 1657, + "audio": {"path": r"C:\scratch\10288018704489549018.wav", "key": "second"}, + "transcription": "ignored", + }, + ] + + def transcribe(value: dict[str, str]) -> str: + if value["key"] == "second": + raise _RejectedSampleError("bad audio") + return "hello world" + + evaluator._transcribe = transcribe + result = evaluator.compute() + assert result["wer"] == 0.0 + assert result["cer"] == 0.0 + assert result["requested_samples"] == 2 + assert result["selected_samples"] == 2 + assert result["selected_source_ids"] == [1525, 1657] + assert result["selected_source_indices"] == [0, 1] + assert result["selected_audio_paths"] == [ + "10018492969996036091.wav", + "10288018704489549018.wav", + ] + assert result["selected_audio_keys"] == ["first", "second"] + assert result["selected_rows"] == [ + { + "source_index": 0, + "dataset_id": 1525, + "audio_path": "10018492969996036091.wav", + "audio_key": "first", + }, + { + "source_index": 1, + "dataset_id": 1657, + "audio_path": "10288018704489549018.wav", + "audio_key": "second", + }, + ] + assert result["processed_samples"] == 1 + assert result["skipped_samples"] == 0 + assert result["rejected_samples"] == 1 + assert result["rejection_reasons"] == {"bad audio": 1} + assert ( + result["processed_samples"] + result["rejected_samples"] + result["skipped_samples"] + == result["selected_samples"] + ) + + +def test_compute_rejects_duplicate_selected_source_indices_before_inference() -> None: + evaluator = _evaluator(input_shape=[1, "samples"]) + evaluator.data = [ + {"_winml_source_index": 0, "id": 1, "audio": "first", "transcription": "first"}, + {"_winml_source_index": 0, "id": 2, "audio": "second", "transcription": "second"}, + ] + evaluator._transcribe = MagicMock() + + with pytest.raises(DatasetValidationError, match="duplicate source indices"): + evaluator.compute() + + evaluator._transcribe.assert_not_called() + + +def test_compute_rejects_missing_selected_source_index_before_inference() -> None: + evaluator = _evaluator(input_shape=[1, "samples"]) + evaluator.data = [{"audio": "value", "transcription": "reference"}] + evaluator._transcribe = MagicMock() + + with pytest.raises(DatasetValidationError, match="no valid source index provenance"): + evaluator.compute() + + evaluator._transcribe.assert_not_called() + + +@pytest.mark.parametrize( + ("row", "message"), + [ + ( + { + "_winml_source_index": 0, + "id": True, + "audio": {"path": "sample.wav"}, + "transcription": "reference", + }, + "dataset ID", + ), + ( + { + "_winml_source_index": 0, + "id": " ", + "audio": {"path": "sample.wav"}, + "transcription": "reference", + }, + "dataset ID", + ), + ], +) +def test_compute_rejects_malformed_selected_identity_before_inference( + row: dict[str, object], + message: str, +) -> None: + evaluator = _evaluator(input_shape=[1, "samples"]) + evaluator.data = [row] + evaluator._transcribe = MagicMock() + + with pytest.raises(DatasetValidationError, match=message): + evaluator.compute() + + evaluator._transcribe.assert_not_called() + + +def test_compute_accepts_missing_dataset_id_and_duplicate_audio_basename() -> None: + evaluator = _evaluator(input_shape=[1, "samples"]) + evaluator.data = [ + { + "_winml_source_index": 0, + "audio": {"path": r"C:\first\same.wav", "key": "first"}, + "transcription": "first", + }, + { + "_winml_source_index": 1, + "audio": {"path": r"D:\second\same.wav", "key": "second"}, + "transcription": "second", + }, + ] + evaluator._transcribe = MagicMock(side_effect=["first", "second"]) + + result = evaluator.compute() + + assert result["selected_source_ids"] == [None, None] + assert result["selected_audio_paths"] == ["same.wav", "same.wav"] + assert evaluator._transcribe.call_count == 2 + + +def test_compute_accepts_duplicate_dataset_ids_with_distinct_row_and_audio_identity() -> None: + evaluator = _evaluator(input_shape=[1, "samples"]) + evaluator.data = [ + { + "_winml_source_index": 0, + "id": 1525, + "audio": {"path": "first.wav"}, + "transcription": "same sentence", + }, + { + "_winml_source_index": 1, + "id": 1525, + "audio": {"path": "second.wav"}, + "transcription": "same sentence", + }, + ] + evaluator._transcribe = MagicMock(return_value="same sentence") + + result = evaluator.compute() + + assert result["selected_source_ids"] == [1525, 1525] + assert result["selected_source_indices"] == [0, 1] + assert result["selected_audio_paths"] == ["first.wav", "second.wav"] + assert evaluator._transcribe.call_count == 2 + + +def test_compute_accepts_and_emits_pinned_fleurs_row_identities() -> None: + evaluator = _evaluator(input_shape=[1, "samples"]) + evaluator.data = [ + { + "_winml_source_index": 0, + "id": 1525, + "audio": {"path": r"C:\cache\10018492969996036091.wav"}, + "transcription": "first", + }, + { + "_winml_source_index": 1, + "id": 1657, + "audio": {"path": r"C:\cache\10288018704489549018.wav"}, + "transcription": "second", + }, + ] + evaluator._transcribe = MagicMock(side_effect=["first", "second"]) + + result = evaluator.compute() + + assert result["selected_source_ids"] == [1525, 1657] + assert result["selected_source_indices"] == [0, 1] + assert result["selected_audio_paths"] == [ + "10018492969996036091.wav", + "10288018704489549018.wav", + ] + assert evaluator._transcribe.call_count == 2 + + +def test_compute_scores_successful_empty_decode_as_deletions() -> None: + evaluator = _evaluator(input_shape=[1, "samples"]) + evaluator.data = [ + {"_winml_source_index": 0, "id": 1, "audio": "empty", "transcription": "hello world"}, + {"_winml_source_index": 1, "id": 2, "audio": "decoded", "transcription": "good day"}, + ] + evaluator._transcribe = lambda value: "" if value == "empty" else "good day" + + result = evaluator.compute() + + assert result["wer"] == 0.5 + assert result["cer"] == pytest.approx(11 / 19) + assert result["predictions"] == ["", "good day"] + assert result["processed_samples"] == 2 + assert result["skipped_samples"] == 0 + assert result["rejected_samples"] == 0 + assert result["rejection_reasons"] == {} + + +def test_compute_scores_all_empty_hypotheses_as_deletions() -> None: + evaluator = _evaluator(input_shape=[1, "samples"]) + evaluator.data = [ + {"_winml_source_index": 0, "id": 1, "audio": "first", "transcription": "hello world"}, + {"_winml_source_index": 1, "id": 2, "audio": "second", "transcription": "good day"}, + ] + evaluator._transcribe = lambda _value: "" + + result = evaluator.compute() + + assert result["wer"] == 1.0 + assert result["cer"] == 1.0 + assert result["predictions"] == ["", ""] + assert result["processed_samples"] == 2 + assert result["skipped_samples"] == 0 + assert result["rejected_samples"] == 0 + assert result["rejection_reasons"] == {} + + +def test_compute_rejects_empty_reference_but_keeps_empty_hypothesis() -> None: + evaluator = _evaluator(input_shape=[1, "samples"]) + evaluator.data = [ + { + "_winml_source_index": 0, + "id": 1, + "audio": "invalid-reference", + "transcription": " \t ", + }, + { + "_winml_source_index": 1, + "id": 2, + "audio": "empty-hypothesis", + "transcription": "valid", + }, + ] + evaluator._transcribe = lambda _value: "" + + result = evaluator.compute() + + assert result["wer"] == 1.0 + assert result["cer"] == 1.0 + assert result["predictions"] == [""] + assert result["references"] == ["valid"] + assert result["processed_samples"] == 1 + assert result["skipped_samples"] == 0 + assert result["rejected_samples"] == 1 + assert result["rejection_reasons"] == {"normalized transcription is empty": 1} + + +def test_compute_distinguishes_decode_failure_from_empty_decode() -> None: + evaluator = _evaluator(input_shape=[1, "samples"]) + evaluator.data = [ + {"_winml_source_index": 0, "id": 1, "audio": "failure", "transcription": "rejected"}, + {"_winml_source_index": 1, "id": 2, "audio": "empty", "transcription": "kept"}, + ] + + def transcribe(value: str) -> str: + if value == "failure": + raise _RejectedSampleError("decode failed") + return "" + + evaluator._transcribe = transcribe + result = evaluator.compute() + + assert result["wer"] == 1.0 + assert result["cer"] == 1.0 + assert result["predictions"] == [""] + assert result["processed_samples"] == 1 + assert result["skipped_samples"] == 0 + assert result["rejected_samples"] == 1 + assert result["rejection_reasons"] == {"decode failed": 1} + + +def test_compute_fails_closed_when_all_rows_are_rejected() -> None: + evaluator = _evaluator(input_shape=[1, "samples"]) + evaluator.data = [{"_winml_source_index": 0, "id": 1, "audio": "invalid", "transcription": ""}] + with pytest.raises(DatasetValidationError, match="processed=0, rejected=1"): + evaluator.compute() + + +def test_unexpected_inference_error_propagates() -> None: + evaluator = _evaluator(input_shape=[1, "samples"]) + evaluator.data = [ + {"_winml_source_index": 0, "id": 1, "audio": "value", "transcription": "reference"} + ] + evaluator._transcribe = MagicMock(side_effect=RuntimeError("session failed")) + with pytest.raises(RuntimeError, match="session failed"): + evaluator.compute() diff --git a/tests/unit/loader/test_load_hf_model.py b/tests/unit/loader/test_load_hf_model.py index 01b46e22b..bd560ba44 100644 --- a/tests/unit/loader/test_load_hf_model.py +++ b/tests/unit/loader/test_load_hf_model.py @@ -271,6 +271,68 @@ def test_attention_implementation_is_forwarded_before_model_construction(self, m attn_implementation="eager", ) + def test_target_language_adapter_is_loaded_before_export(self, monkeypatch): + from types import SimpleNamespace + from unittest.mock import MagicMock + + import winml.modelkit.loader.resolution as resolution_module + + model_class = MagicMock() + model_class.__name__ = "Wav2Vec2ForCTC" + model_class.config_class = None + model = MagicMock() + model.config = SimpleNamespace() + model.parameters.return_value = [] + model_class.from_pretrained.return_value = model + config = SimpleNamespace(model_type="wav2vec2", architectures=["Wav2Vec2ForCTC"]) + monkeypatch.setattr( + resolution_module, + "resolve_task", + lambda *_args, **_kwargs: SimpleNamespace( + task="automatic-speech-recognition", + model_class=model_class, + ), + ) + + loaded, loaded_config, _ = load_hf_model( + "facebook/mms-1b-all", + hf_config=config, + target_lang="deu", + ) + + assert loaded is model + model.load_adapter.assert_called_once_with("deu") + assert loaded.config.target_lang == "deu" + assert loaded_config.target_lang == "deu" + + def test_target_language_rejects_model_without_adapter_support(self, monkeypatch): + from types import SimpleNamespace + from unittest.mock import MagicMock + + import winml.modelkit.loader.resolution as resolution_module + + model_class = MagicMock() + model_class.__name__ = "ModelWithoutAdapters" + model_class.config_class = None + model = SimpleNamespace( + config=SimpleNamespace(), + eval=lambda: None, + parameters=list, + ) + model_class.from_pretrained.return_value = model + config = SimpleNamespace(model_type="unit", architectures=["ModelWithoutAdapters"]) + monkeypatch.setattr( + resolution_module, + "resolve_task", + lambda *_args, **_kwargs: SimpleNamespace( + task="automatic-speech-recognition", + model_class=model_class, + ), + ) + + with pytest.raises(ValueError, match="cannot load requested language adapter"): + load_hf_model("org/model", hf_config=config, target_lang="deu") + def test_bert_tiny_uses_model_specific_default_task(self, monkeypatch): """bert-tiny should use model-specific default task when task is omitted.""" from unittest.mock import MagicMock diff --git a/tests/unit/loader/test_loader_config.py b/tests/unit/loader/test_loader_config.py index 3e46293f9..c0315aec4 100644 --- a/tests/unit/loader/test_loader_config.py +++ b/tests/unit/loader/test_loader_config.py @@ -47,6 +47,13 @@ def test_to_dict_with_model_type(self): d = config.to_dict() assert d["model_type"] == "bert" + def test_target_lang_roundtrip(self): + config = WinMLLoaderConfig(target_lang="deu") + + restored = WinMLLoaderConfig.from_dict(config.to_dict()) + + assert restored.target_lang == "deu" + def test_to_dict_empty_when_defaults(self): """Test config serialization returns empty dict for defaults.""" config = WinMLLoaderConfig() diff --git a/tests/unit/loader/test_resolve_loader_config.py b/tests/unit/loader/test_resolve_loader_config.py index 08928be87..ea4cb61b4 100644 --- a/tests/unit/loader/test_resolve_loader_config.py +++ b/tests/unit/loader/test_resolve_loader_config.py @@ -260,6 +260,26 @@ def test_trust_remote_code_propagated(self) -> None: assert loader_config.trust_remote_code is True + def test_mms_target_language_is_persisted_in_loader_config(self) -> None: + mock_config = MagicMock() + mock_config.model_type = "wav2vec2" + mock_config.target_lang = "deu" + mock_config.adapter_attn_dim = 16 + mock_class = MagicMock(spec=[]) + mock_class.__name__ = "Wav2Vec2ForCTC" + mock_class.config_class = None + + with patch( + "winml.modelkit.loader.resolution.resolve_task", + return_value=_make_resolution("automatic-speech-recognition", mock_class), + ): + loader_config, _, _, _ = resolve_loader_config( + "facebook/mms-1b-all", + hf_config=mock_config, + ) + + assert loader_config.target_lang == "deu" + def test_explicit_task_passed_through(self) -> None: """Explicit task is forwarded to resolve_task.""" mock_config = MagicMock()