diff --git a/src/winml/modelkit/eval/__init__.py b/src/winml/modelkit/eval/__init__.py index 7aa18048f..9bf41ca9a 100644 --- a/src/winml/modelkit/eval/__init__.py +++ b/src/winml/modelkit/eval/__init__.py @@ -40,6 +40,7 @@ from .tensor_similarity_evaluator import TensorSimilarityEvaluator from .text_classification_evaluator import WinMLTextClassificationEvaluator from .token_classification_evaluator import WinMLTokenClassificationEvaluator + from .translation_evaluator import WinMLTranslationEvaluator from .zero_shot_classification_evaluator import WinMLZeroShotClassificationEvaluator from .zero_shot_image_classification_evaluator import WinMLZeroShotImageClassificationEvaluator @@ -68,6 +69,8 @@ ".text_classification_evaluator:WinMLTextClassificationEvaluator", "WinMLTokenClassificationEvaluator": ".token_classification_evaluator:WinMLTokenClassificationEvaluator", + "WinMLTranslationEvaluator": + ".translation_evaluator:WinMLTranslationEvaluator", "WinMLZeroShotClassificationEvaluator": ".zero_shot_classification_evaluator:WinMLZeroShotClassificationEvaluator", "WinMLZeroShotImageClassificationEvaluator": @@ -142,6 +145,7 @@ def __dir__() -> list[str]: "WinMLQuestionAnsweringEvaluator", "WinMLTextClassificationEvaluator", "WinMLTokenClassificationEvaluator", + "WinMLTranslationEvaluator", "WinMLZeroShotClassificationEvaluator", "WinMLZeroShotImageClassificationEvaluator", "evaluate", diff --git a/src/winml/modelkit/eval/evaluate.py b/src/winml/modelkit/eval/evaluate.py index 74863970f..3b71939d4 100644 --- a/src/winml/modelkit/eval/evaluate.py +++ b/src/winml/modelkit/eval/evaluate.py @@ -55,6 +55,8 @@ "winml.modelkit.eval.image_feature_extraction_evaluator:WinMLImageFeatureExtractionEvaluator", "image-to-text": "winml.modelkit.eval.image_to_text_evaluator:WinMLImageToTextEvaluator", + "translation": + "winml.modelkit.eval.translation_evaluator:WinMLTranslationEvaluator", "fill-mask": "winml.modelkit.eval.fill_mask_evaluator:WinMLFillMaskEvaluator", "zero-shot-classification": diff --git a/src/winml/modelkit/eval/metrics/translation.py b/src/winml/modelkit/eval/metrics/translation.py new file mode 100644 index 000000000..11f63d86f --- /dev/null +++ b/src/winml/modelkit/eval/metrics/translation.py @@ -0,0 +1,48 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- + +"""Corpus-level machine-translation metrics.""" + +from __future__ import annotations + +from typing import Any + + +class TranslationMetric: + """Aggregate corpus SacreBLEU and chrF over translated sentences.""" + + def __init__(self) -> None: + self._predictions: list[str] = [] + self._references: list[list[str]] = [] + + def update(self, prediction: str, references: str | list[str]) -> None: + """Record one prediction and one or more non-empty references.""" + refs = [references] if isinstance(references, str) else references + cleaned = [reference.strip() for reference in refs if reference and reference.strip()] + if not cleaned: + raise ValueError("at least one non-empty translation reference is required") + self._predictions.append((prediction or "").strip()) + self._references.append(cleaned) + + def compute(self) -> dict[str, Any]: + """Return corpus SacreBLEU and chrF scores on the conventional 0-100 scale.""" + if not self._predictions: + return {"sacrebleu": None, "chrf": None, "n_samples": 0} + + from torchmetrics.text import CHRFScore, SacreBLEUScore + + sacrebleu = SacreBLEUScore(tokenize="13a")( + self._predictions, + self._references, + ) + # Standard chrF2: character n-grams only (word_order=0), beta=2. + # TorchMetrics otherwise defaults to word_order=2, which is chrF++. + chrf = CHRFScore(n_word_order=0, beta=2.0)(self._predictions, self._references) + + return { + "sacrebleu": round(float(sacrebleu) * 100, 4), + "chrf": round(float(chrf) * 100, 4), + "n_samples": len(self._predictions), + } diff --git a/src/winml/modelkit/eval/translation_evaluator.py b/src/winml/modelkit/eval/translation_evaluator.py new file mode 100644 index 000000000..93484ac25 --- /dev/null +++ b/src/winml/modelkit/eval/translation_evaluator.py @@ -0,0 +1,195 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- + +"""Evaluator for text-to-text machine translation models.""" + +from __future__ import annotations + +import logging +from collections.abc import Mapping +from typing import TYPE_CHECKING, Any + +from .base_evaluator import WinMLEvaluator + + +if TYPE_CHECKING: + from datasets import Dataset + + from ..models.winml.composite_model import WinMLCompositeModel + from .config import DatasetConfig, WinMLEvaluationConfig + + +logger = logging.getLogger(__name__) + + +class WinMLTranslationEvaluator(WinMLEvaluator): + """Evaluate translations with corpus SacreBLEU and chrF. + + Flat datasets can map separate ``source_column`` and + ``reference_column`` values. WMT-style datasets can leave both columns + mapped to ``translation`` and provide explicit ``source_lang`` and + ``target_lang`` keys. Language direction is never guessed from a model ID. + """ + + def __init__( + self, + config: WinMLEvaluationConfig, + model: WinMLCompositeModel, + ) -> None: + from ..utils.eval_utils import get_default + + mapping = config.dataset.columns_mapping + self._source_col = mapping.get( + "source_column", + get_default("translation", "source_column") or "translation", + ) + self._reference_col = mapping.get( + "reference_column", + get_default("translation", "reference_column") or "translation", + ) + self._source_lang = mapping.get("source_lang") + self._target_lang = mapping.get("target_lang") + self._tokenizer_source_lang = mapping.get("tokenizer_source_lang", self._source_lang) + self._tokenizer_target_lang = mapping.get("tokenizer_target_lang", self._target_lang) + self._source_prefix = mapping.get("source_prefix", "") + super().__init__(config, model) + self._pipeline_kwargs: dict[str, Any] = { + "num_beams": 1, + "truncation": True, + } + if self._tokenizer_source_lang: + self._pipeline_kwargs["src_lang"] = self._tokenizer_source_lang + if self._tokenizer_target_lang: + self._pipeline_kwargs["tgt_lang"] = self._tokenizer_target_lang + + max_encoder_length = getattr(model, "max_encoder_length", None) + if ( + isinstance(max_encoder_length, int) + and max_encoder_length > 0 + and self.pipe.tokenizer is not None + ): + self.pipe.tokenizer.model_max_length = max_encoder_length + + max_decode_length = getattr(model, "max_decode_length", None) + if isinstance(max_decode_length, int) and max_decode_length > 1: + self._pipeline_kwargs["max_new_tokens"] = max_decode_length - 1 + # TranslationPipeline.check_inputs() only reads max_length, while + # generate() should use the explicit max_new_tokens capacity. + # Align the check and remove checkpoint-specific beam/token caps. + self.pipe.generation_config.max_length = max_decode_length + self.pipe.generation_config.max_new_tokens = None + self.pipe.generation_config.num_beams = 1 + + def align_labels(self, dataset: Dataset, ds_config: DatasetConfig) -> Dataset: + """Return free-text references unchanged.""" + return dataset + + @staticmethod + def _extract_text(value: Any, language: str | None, field: str) -> str: + """Extract one text value from a flat string or language-keyed mapping.""" + from ..utils.eval_utils import DatasetValidationError + + if isinstance(value, Mapping): + if not language: + raise DatasetValidationError( + f"{field} contains a translation dict; provide --column " + f"{'source_lang' if field == 'source' else 'target_lang'}=", + ) + if language not in value: + raise DatasetValidationError( + f"{field} translation dict has no language key '{language}'; " + f"available keys: {sorted(str(key) for key in value)}", + ) + value = value[language] + if not isinstance(value, str) or not value.strip(): + raise DatasetValidationError(f"{field} must resolve to a non-empty string") + return value.strip() + + def _extract_references(self, value: Any) -> str | list[str]: + """Extract one or more reference translations.""" + if isinstance(value, list): + if not value: + from ..utils.eval_utils import DatasetValidationError + + raise DatasetValidationError("reference must contain at least one translation") + return [self._extract_text(item, self._target_lang, "reference") for item in value] + return self._extract_text(value, self._target_lang, "reference") + + @staticmethod + def _prediction_text(output: Any) -> str: + """Normalize Hugging Face translation pipeline output shapes.""" + from ..utils.eval_utils import DatasetValidationError + + if isinstance(output, list): + if not output: + raise DatasetValidationError("pipeline returned no translations") + output = output[0] + if isinstance(output, Mapping): + prediction = output.get("translation_text", output.get("generated_text", "")) + if isinstance(prediction, str) and prediction.strip(): + return prediction.strip() + raise DatasetValidationError("pipeline returned an empty translation") + if isinstance(output, str) and output.strip(): + return output.strip() + raise DatasetValidationError("pipeline returned an unsupported translation result") + + def _extract_sample(self, sample: Any) -> tuple[str, str | list[str]]: + """Extract and validate source/reference values from one dataset row.""" + from ..utils.eval_utils import DatasetValidationError + + if not isinstance(sample, Mapping): + raise DatasetValidationError( + f"dataset row must be a mapping, got {type(sample).__name__}" + ) + source = self._extract_text( + sample.get(self._source_col), + self._source_lang, + "source", + ) + return f"{self._source_prefix}{source}", self._extract_references( + sample.get(self._reference_col) + ) + + def compute(self) -> dict[str, Any]: + """Translate each valid row and compute corpus-level metrics.""" + from tqdm.auto import tqdm + + from ..utils.eval_utils import DatasetValidationError + from .metrics.translation import TranslationMetric + + metric = TranslationMetric() + skipped = 0 + first_error: str | None = None + attempted = 0 + + for sample in tqdm(self.data, desc="Evaluating", unit="sample"): + attempted += 1 + try: + source, references = self._extract_sample(sample) + except DatasetValidationError as error: + logger.warning("Translation sample rejected: %s", error) + first_error = first_error or str(error) + skipped += 1 + continue + + # Do not catch inference/tokenizer/runtime failures: a broken model + # must fail the evaluation rather than produce partial metrics. + output = self.pipe(source, **self._pipeline_kwargs) + try: + prediction = self._prediction_text(output) + metric.update(prediction, references) + except (DatasetValidationError, ValueError) as error: + logger.warning("Translation output rejected: %s", error) + first_error = first_error or str(error) + skipped += 1 + + result = metric.compute() + if result["n_samples"] == 0: + detail = f": {first_error}" if first_error else "" + raise DatasetValidationError(f"No valid translation samples were evaluated{detail}") + result["attempted"] = attempted + result["evaluated"] = result["n_samples"] + result["skipped"] = skipped + return result diff --git a/src/winml/modelkit/models/winml/encoder_decoder.py b/src/winml/modelkit/models/winml/encoder_decoder.py index d488ae842..6d5a4457a 100644 --- a/src/winml/modelkit/models/winml/encoder_decoder.py +++ b/src/winml/modelkit/models/winml/encoder_decoder.py @@ -191,6 +191,12 @@ def __init__( enc_expected = dict( zip(enc_io.get("input_names", []), enc_io.get("input_shapes", []), strict=False) ) + input_ids_shape = enc_expected.get("input_ids", []) + self._max_enc = ( + input_ids_shape[1] + if len(input_ids_shape) > 1 and isinstance(input_ids_shape[1], int) + else None + ) # Wrap encoder with auto-padding so all callsites just use self._encoder(...) self._encoder = self._EncoderWithInputPadding(raw_encoder, enc_expected) @@ -200,7 +206,12 @@ def __init__( ) # Max decode length and KV dtype from decoder ONNX metadata - self._max_dec = self._dec_expected["past_0_key"][2] + max_decode_length = self._dec_expected["past_0_key"][2] + if not isinstance(max_decode_length, int) or max_decode_length <= 0: + raise ValueError( + "Decoder input 'past_0_key' must have a positive static cache length" + ) + self._max_dec = max_decode_length self._num_kv_layers = sum( 1 for n in self._dec_expected if n.startswith("past_") and n.endswith("_key") ) @@ -219,6 +230,16 @@ def __init__( _np_dtype = dec_type_map["past_0_key"] self._kv_dtype = torch.from_numpy(np.zeros(1, dtype=_np_dtype)).dtype + @property + def max_encoder_length(self) -> int | None: + """Return the encoder's static token capacity, or ``None`` when dynamic.""" + return self._max_enc + + @property + def max_decode_length(self) -> int: + """Return the decoder's static output/cache capacity.""" + return self._max_dec + # ----- Encoder ----- class _EncoderWithInputPadding(torch.nn.Module): diff --git a/src/winml/modelkit/utils/eval_utils.py b/src/winml/modelkit/utils/eval_utils.py index 935202bf2..5bcf5e0a5 100644 --- a/src/winml/modelkit/utils/eval_utils.py +++ b/src/winml/modelkit/utils/eval_utils.py @@ -243,6 +243,51 @@ class TaskSchema: roles=("encoder", "decoder"), ) +_TRANSLATION_SCHEMA = TaskSchema( + columns=( + SchemaItem( + "source_column", + "source text, or a translation dict containing the source language key", + default="translation", + remap_hint="", + ), + SchemaItem( + "reference_column", + "reference text(s), or a translation dict containing the target language key", + default="translation", + remap_hint="", + ), + ), + params=( + SchemaItem( + "source_lang", + "source-language key when source_column contains translation dicts", + remap_hint="", + ), + SchemaItem( + "target_lang", + "target-language key when reference_column contains translation dicts", + remap_hint="", + ), + SchemaItem( + "tokenizer_source_lang", + "source-language identifier for multilingual tokenizers; defaults to source_lang", + remap_hint="", + ), + SchemaItem( + "tokenizer_target_lang", + "target-language identifier for multilingual tokenizers; defaults to target_lang", + remap_hint="", + ), + SchemaItem( + "source_prefix", + "optional task prefix prepended before tokenization (for example, a T5 prompt)", + remap_hint="", + ), + ), + roles=("encoder", "decoder"), +) + _FILL_MASK_SCHEMA = TaskSchema( columns=( SchemaItem( @@ -436,6 +481,7 @@ class TaskSchema: "sentence-similarity": _FEATURE_EXTRACTION_SCHEMA, "image-feature-extraction": _IMAGE_FEATURE_EXTRACTION_SCHEMA, "image-to-text": _IMAGE_TO_TEXT_SCHEMA, + "translation": _TRANSLATION_SCHEMA, "fill-mask": _FILL_MASK_SCHEMA, "zero-shot-classification": _ZERO_SHOT_CLASSIFICATION_SCHEMA, "zero-shot-image-classification": _ZERO_SHOT_IMAGE_CLASSIFICATION_SCHEMA, diff --git a/tests/unit/eval/test_translation_evaluator.py b/tests/unit/eval/test_translation_evaluator.py new file mode 100644 index 000000000..d92dc9188 --- /dev/null +++ b/tests/unit/eval/test_translation_evaluator.py @@ -0,0 +1,271 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- + +"""Unit tests for translation evaluation.""" + +from __future__ import annotations + +from unittest.mock import MagicMock, patch + +import pytest +from datasets import Dataset + +from winml.modelkit.eval.translation_evaluator import WinMLTranslationEvaluator +from winml.modelkit.utils.eval_utils import DatasetValidationError + + +def make_evaluator(rows, columns_mapping=None, pipeline=None): + from winml.modelkit.eval import DatasetConfig, WinMLEvaluationConfig + + dataset = Dataset.from_list(rows) + model = MagicMock() + model.config.label2id = None + model.max_encoder_length = 512 + model.max_decode_length = 512 + config = WinMLEvaluationConfig( + model_id="Helsinki-NLP/opus-mt-fr-en", + task="translation", + dataset=DatasetConfig( + path="wmt14", + samples=len(rows), + shuffle=False, + columns_mapping=columns_mapping or {}, + ), + ) + with ( + patch("datasets.load_dataset", return_value=dataset), + patch( + "winml.modelkit.eval.base_evaluator.WinMLEvaluator.prepare_pipeline", + return_value=pipeline or MagicMock(), + ), + ): + return WinMLTranslationEvaluator(config, model) + + +class TestTranslationMetric: + def test_perfect_corpus_scores_100(self): + from winml.modelkit.eval.metrics.translation import TranslationMetric + + metric = TranslationMetric() + metric.update("This is a complete test sentence.", "This is a complete test sentence.") + result = metric.compute() + + assert result == {"sacrebleu": 100.0, "chrf": 100.0, "n_samples": 1} + + def test_empty_corpus_has_no_scores(self): + from winml.modelkit.eval.metrics.translation import TranslationMetric + + assert TranslationMetric().compute() == { + "sacrebleu": None, + "chrf": None, + "n_samples": 0, + } + + def test_multiple_references_are_supported(self): + from winml.modelkit.eval.metrics.translation import TranslationMetric + + metric = TranslationMetric() + metric.update( + "A sufficiently long reference sentence is here.", + [ + "Another reference sentence is also supplied.", + "A sufficiently long reference sentence is here.", + ], + ) + assert metric.compute()["sacrebleu"] == 100.0 + + def test_multiple_perfect_samples_have_perfect_corpus_scores(self): + from winml.modelkit.eval.metrics.translation import TranslationMetric + + metric = TranslationMetric() + first = "The first complete sentence is correct." + second = "The second complete sentence is correct." + metric.update(first, first) + metric.update(second, second) + result = metric.compute() + + assert result["sacrebleu"] == 100.0 + assert result["chrf"] == 100.0 + + def test_empty_references_are_rejected(self): + from winml.modelkit.eval.metrics.translation import TranslationMetric + + with pytest.raises(ValueError, match="at least one non-empty"): + TranslationMetric().update("prediction", []) + + +class TestTranslationEvaluator: + def test_static_encoder_capacity_allows_pipeline_without_tokenizer(self): + pipeline = MagicMock() + pipeline.tokenizer = None + + evaluator = make_evaluator( + [{"source": "texte", "reference": "text"}], + {"source_column": "source", "reference_column": "reference"}, + pipeline, + ) + + assert evaluator.pipe.tokenizer is None + + def test_nested_translation_dict_uses_explicit_language_keys(self): + evaluator = make_evaluator( + [{"translation": {"fr": "Bonjour le monde", "en": "Hello world"}}], + {"source_lang": "fr", "target_lang": "en"}, + ) + assert evaluator.pipe.tokenizer.model_max_length == 512 + assert evaluator.pipe.generation_config.max_new_tokens is None + assert evaluator.pipe.generation_config.num_beams == 1 + evaluator.pipe = MagicMock(return_value=[{"translation_text": "Hello world"}]) + + result = evaluator.compute() + + evaluator.pipe.assert_called_once_with( + "Bonjour le monde", + num_beams=1, + truncation=True, + src_lang="fr", + tgt_lang="en", + max_new_tokens=511, + ) + assert result["n_samples"] == 1 + assert result["chrf"] == 100.0 + assert result["attempted"] == 1 + assert result["evaluated"] == 1 + assert result["skipped"] == 0 + + def test_flat_custom_columns_need_no_language_keys(self): + evaluator = make_evaluator( + [{"french": "Une phrase source", "english": "A source sentence"}], + {"source_column": "french", "reference_column": "english"}, + ) + evaluator.pipe = MagicMock(return_value={"generated_text": "A source sentence"}) + + result = evaluator.compute() + + assert result["n_samples"] == 1 + assert result["chrf"] == 100.0 + + def test_missing_language_key_rejects_row_with_accounting(self): + evaluator = make_evaluator( + [ + {"translation": {"fr": "Valide", "en": "Valid"}}, + {"translation": {"fr": "Invalide", "de": "Ungültig"}}, + ], + {"source_lang": "fr", "target_lang": "en"}, + ) + evaluator.pipe = MagicMock(return_value=[{"translation_text": "Valid"}]) + + result = evaluator.compute() + + assert result["n_samples"] == 1 + assert result["attempted"] == 2 + assert result["evaluated"] == 1 + assert result["skipped"] == 1 + evaluator.pipe.assert_called_once() + + def test_nested_translation_requires_explicit_direction(self): + evaluator = make_evaluator( + [{"translation": {"fr": "Bonjour", "en": "Hello"}}], + ) + + with pytest.raises(DatasetValidationError, match="provide --column source_lang"): + evaluator.compute() + + def test_pipeline_runtime_failure_propagates(self): + evaluator = make_evaluator( + [ + {"source": "un", "reference": "one"}, + {"source": "deux", "reference": "two"}, + ], + {"source_column": "source", "reference_column": "reference"}, + ) + evaluator.pipe = MagicMock( + side_effect=[ + [{"translation_text": "one"}], + RuntimeError("decoder failed"), + ] + ) + + with pytest.raises(RuntimeError, match="decoder failed"): + evaluator.compute() + + def test_all_rejected_fails_closed(self): + evaluator = make_evaluator( + [{"source": "texte", "reference": "text"}], + {"source_column": "source", "reference_column": "reference"}, + ) + evaluator.pipe = MagicMock(return_value=[{"translation_text": ""}]) + + with pytest.raises(DatasetValidationError, match="No valid translation samples"): + evaluator.compute() + + @pytest.mark.parametrize("output", [[], [{}], [{"translation_text": ""}]]) + def test_empty_pipeline_output_is_rejected(self, output): + evaluator = make_evaluator( + [{"source": "texte", "reference": "text"}], + {"source_column": "source", "reference_column": "reference"}, + ) + evaluator.pipe = MagicMock(return_value=output) + + with pytest.raises(DatasetValidationError, match="No valid translation samples"): + evaluator.compute() + + def test_empty_reference_list_is_rejected(self): + evaluator = make_evaluator( + [{"source": "texte", "references": []}], + {"source_column": "source", "reference_column": "references"}, + ) + + with pytest.raises(DatasetValidationError, match="No valid translation samples"): + evaluator.compute() + + def test_non_mapping_row_is_rejected_with_accounting(self): + evaluator = make_evaluator( + [{"source": "valide", "reference": "valid"}], + {"source_column": "source", "reference_column": "reference"}, + ) + evaluator.data = [None, {"source": "valide", "reference": "valid"}] + evaluator.pipe = MagicMock(return_value=[{"translation_text": "valid"}]) + + result = evaluator.compute() + + assert result["attempted"] == 2 + assert result["evaluated"] == 1 + assert result["skipped"] == 1 + + def test_tokenizer_language_ids_and_source_prefix_are_independent_of_dataset_keys(self): + evaluator = make_evaluator( + [{"translation": {"fra_Latn": "Bonjour", "eng_Latn": "Hello"}}], + { + "source_lang": "fra_Latn", + "target_lang": "eng_Latn", + "tokenizer_source_lang": "fra_Latn", + "tokenizer_target_lang": "eng_Latn", + "source_prefix": "translate French to English: ", + }, + ) + evaluator.pipe = MagicMock(return_value=[{"translation_text": "Hello"}]) + + evaluator.compute() + + evaluator.pipe.assert_called_once_with( + "translate French to English: Bonjour", + num_beams=1, + truncation=True, + src_lang="fra_Latn", + tgt_lang="eng_Latn", + max_new_tokens=511, + ) + + +class TestTranslationRegistration: + def test_registry_and_schema_are_present(self): + from winml.modelkit.eval import WinMLEvaluationConfig, get_evaluator_class + from winml.modelkit.utils.eval_utils import TASK_SCHEMAS + + assert get_evaluator_class(WinMLEvaluationConfig(task="translation")) is ( + WinMLTranslationEvaluator + ) + assert TASK_SCHEMAS["translation"].roles == ("encoder", "decoder")