feat(eval): add translation quality evaluation

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",

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":

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

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'}=<language key>",
+                )
+            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

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):

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="<your_source_column>",
+        ),
+        SchemaItem(
+            "reference_column",
+            "reference text(s), or a translation dict containing the target language key",
+            default="translation",
+            remap_hint="<your_reference_column>",
+        ),
+    ),
+    params=(
+        SchemaItem(
+            "source_lang",
+            "source-language key when source_column contains translation dicts",
+            remap_hint="<source_language_key>",
+        ),
+        SchemaItem(
+            "target_lang",
+            "target-language key when reference_column contains translation dicts",
+            remap_hint="<target_language_key>",
+        ),
+        SchemaItem(
+            "tokenizer_source_lang",
+            "source-language identifier for multilingual tokenizers; defaults to source_lang",
+            remap_hint="<tokenizer_source_language>",
+        ),
+        SchemaItem(
+            "tokenizer_target_lang",
+            "target-language identifier for multilingual tokenizers; defaults to target_lang",
+            remap_hint="<tokenizer_target_language>",
+        ),
+        SchemaItem(
+            "source_prefix",
+            "optional task prefix prepended before tokenization (for example, a T5 prompt)",
+            remap_hint="<translation_prefix>",
+        ),
+    ),
+    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,