diff --git a/scripts/e2e_eval/cache/baseline_cache.json b/scripts/e2e_eval/cache/baseline_cache.json index 8acbbe305..9d8dfc857 100644 --- a/scripts/e2e_eval/cache/baseline_cache.json +++ b/scripts/e2e_eval/cache/baseline_cache.json @@ -659,6 +659,16 @@ "elapsed": 1319.8, "command": "python.exe run_pytorch_baseline.py --model bert-base-uncased --task fill-mask --device cpu --num-samples 100 --dataset wikitext --split test --dataset-config wikitext-2-raw-v1 --columns-mapping {\"input_column\": \"text\"}" }, + "Qwen/Qwen3-0.6B|text-generation|Salesforce/wikitext|wikitext-2-raw-v1|test|100": { + "status": "PASS", + "metric": { + "metric": "perplexity", + "value": 23.645425, + "num_samples": 100 + }, + "elapsed": 0.0, + "command": "python.exe run_pytorch_baseline.py --model Qwen/Qwen3-0.6B --task text-generation --device cpu --num-samples 100 --dataset Salesforce/wikitext --split test --dataset-config wikitext-2-raw-v1 --columns-mapping {\"input_column\": \"text\", \"num_tokens\": \"8192\", \"seqlen\": \"2048\"}" + }, "distilbert/distilbert-base-uncased|fill-mask|Salesforce/wikitext|wikitext-2-raw-v1|test|100": { "status": "PASS", "metric": { diff --git a/scripts/e2e_eval/run_pytorch_baseline.py b/scripts/e2e_eval/run_pytorch_baseline.py index 4abc5bff4..ba1f478a3 100644 --- a/scripts/e2e_eval/run_pytorch_baseline.py +++ b/scripts/e2e_eval/run_pytorch_baseline.py @@ -152,6 +152,21 @@ def _measure_pytorch_latency(task_evaluator: Any, warmup: int, iterations: int) def _load_pytorch_model(model_id: str, task: str, device_str: str): """Load a native PyTorch model with the task-appropriate AutoModel class.""" import torch + + device = torch.device( + device_str if device_str != "cuda" or torch.cuda.is_available() else "cpu" + ) + + # text-generation resolves to a composite (embeddings/ctx/iter/lm_head) that + # is not a runnable HF model. Wrap a plain causal LM in the perplexity + # scoring contract so the evaluator scores it through the same path as a + # genai bundle. + if task == "text-generation": + from winml.modelkit.models.winml import HFCausalLM + + _out(f"Loading causal LM baseline adapter for {model_id} on {device}") + return HFCausalLM(model_id, device) + from transformers import AutoConfig from winml.modelkit.loader import resolve_task @@ -159,9 +174,6 @@ def _load_pytorch_model(model_id: str, task: str, device_str: str): config = AutoConfig.from_pretrained(model_id) cls = resolve_task(config, task=task).model_class _out(f"Loading {cls.__name__} for {model_id} on {device_str}") - device = torch.device( - device_str if device_str != "cuda" or torch.cuda.is_available() else "cpu" - ) return cls.from_pretrained(model_id).to(device).eval() diff --git a/scripts/e2e_eval/testsets/models_with_acc.json b/scripts/e2e_eval/testsets/models_with_acc.json index e21423161..e6ea36ecc 100644 --- a/scripts/e2e_eval/testsets/models_with_acc.json +++ b/scripts/e2e_eval/testsets/models_with_acc.json @@ -1246,6 +1246,26 @@ } } }, + { + "hf_id": "Qwen/Qwen3-0.6B", + "task": "text-generation", + "model_type": "qwen3", + "group": "Top200", + "priority": "P1", + "dataset_config": { + "path": "Salesforce/wikitext", + "name": "wikitext-2-raw-v1", + "split": "test", + "samples": 100, + "metric": "perplexity", + "winml_metric_key": "perplexity", + "columns_mapping": { + "input_column": "text", + "num_tokens": "8192", + "seqlen": "2048" + } + } + }, { "hf_id": "FacebookAI/xlm-roberta-base", "task": "fill-mask", diff --git a/scripts/e2e_eval/utils/accuracy.py b/scripts/e2e_eval/utils/accuracy.py index f3fe92b21..4b1250845 100644 --- a/scripts/e2e_eval/utils/accuracy.py +++ b/scripts/e2e_eval/utils/accuracy.py @@ -25,6 +25,8 @@ # WinML-vs-baseline delta is small — pick a tighter threshold than default. "knn_top1_accuracy": ("delta_relative", 0.02, 0.05, True), "pseudo_perplexity": ("delta_relative", 0.05, 0.10, False), + # Perplexity (text-generation PPL): lower is better. + "perplexity": ("delta_relative", 0.05, 0.10, False), # CER (OCR error rate): lower is better. "cer": ("delta_relative", 0.05, 0.10, False), # AbsRel (depth-estimation error): lower is better. diff --git a/src/winml/modelkit/commands/eval.py b/src/winml/modelkit/commands/eval.py index f97f72294..288c01898 100644 --- a/src/winml/modelkit/commands/eval.py +++ b/src/winml/modelkit/commands/eval.py @@ -533,6 +533,13 @@ def _resolve_model_path( "for preprocessor and config resolution." ) return value, model_id + + # A local directory (e.g. an onnxruntime-genai bundle) is a model *path*, + # not a Hub model id. Route it to model_path so downstream loaders read the + # bundle from disk instead of treating the directory string as a Hub ref. + if Path(value).expanduser().is_dir(): + return str(Path(value).expanduser()), model_id + if model_id is not None and model_id != value: raise click.UsageError( "Cannot pass both `-m ` and `--model-id`. " diff --git a/src/winml/modelkit/eval/__init__.py b/src/winml/modelkit/eval/__init__.py index 7aa18048f..5d7f8be5a 100644 --- a/src/winml/modelkit/eval/__init__.py +++ b/src/winml/modelkit/eval/__init__.py @@ -39,6 +39,7 @@ from .question_answering_evaluator import WinMLQuestionAnsweringEvaluator from .tensor_similarity_evaluator import TensorSimilarityEvaluator from .text_classification_evaluator import WinMLTextClassificationEvaluator + from .text_generation_evaluator import WinMLTextGenerationEvaluator from .token_classification_evaluator import WinMLTokenClassificationEvaluator from .zero_shot_classification_evaluator import WinMLZeroShotClassificationEvaluator from .zero_shot_image_classification_evaluator import WinMLZeroShotImageClassificationEvaluator @@ -66,6 +67,8 @@ ".question_answering_evaluator:WinMLQuestionAnsweringEvaluator", "WinMLTextClassificationEvaluator": ".text_classification_evaluator:WinMLTextClassificationEvaluator", + "WinMLTextGenerationEvaluator": + ".text_generation_evaluator:WinMLTextGenerationEvaluator", "WinMLTokenClassificationEvaluator": ".token_classification_evaluator:WinMLTokenClassificationEvaluator", "WinMLZeroShotClassificationEvaluator": @@ -141,6 +144,7 @@ def __dir__() -> list[str]: "WinMLObjectDetectionEvaluator", "WinMLQuestionAnsweringEvaluator", "WinMLTextClassificationEvaluator", + "WinMLTextGenerationEvaluator", "WinMLTokenClassificationEvaluator", "WinMLZeroShotClassificationEvaluator", "WinMLZeroShotImageClassificationEvaluator", diff --git a/src/winml/modelkit/eval/evaluate.py b/src/winml/modelkit/eval/evaluate.py index 74863970f..10e295673 100644 --- a/src/winml/modelkit/eval/evaluate.py +++ b/src/winml/modelkit/eval/evaluate.py @@ -69,6 +69,8 @@ "winml.modelkit.eval.tensor_similarity_evaluator:TensorSimilarityEvaluator", "mask-generation": "winml.modelkit.eval.mask_generation_evaluator:WinMLMaskGenerationEvaluator", + "text-generation": + "winml.modelkit.eval.text_generation_evaluator:WinMLTextGenerationEvaluator", } # fmt: on @@ -199,6 +201,14 @@ def get_evaluator_class(config: WinMLEvaluationConfig) -> type[WinMLEvaluator]: "path": "mattmdjaga/human_parsing_dataset", "split": "train", }, + "text-generation": { + # Raw wikitext-2 test split scored token-by-token for perplexity; + # ``input_column`` names the text field the corpus is built from. + "path": "wikitext", + "name": "wikitext-2-raw-v1", + "split": "test", + "columns_mapping": {"input_column": "text"}, + }, } @@ -231,6 +241,9 @@ def _load_model(config: WinMLEvaluationConfig) -> WinMLPreTrainedModel | WinMLCo from ..models import WinMLAutoModel from ..utils import cli as cli_utils + if config.task == "text-generation": + return _load_genai_causal_lm(config) + quant_override: Any = None if not config.quant: from ..config import WinMLBuildConfig @@ -282,6 +295,51 @@ def _load_model(config: WinMLEvaluationConfig) -> WinMLPreTrainedModel | WinMLCo ) +def _load_genai_causal_lm(config: WinMLEvaluationConfig) -> Any: + """Load a causal LM from an onnxruntime-genai bundle directory. + + ``-m `` resolves to ``config.model_path`` (a local directory), + so the bundle directory is read from there. ``ep`` / ``device`` pass straight + through: an explicit ``--ep`` forces the whole decoder pipeline onto that EP, + while ``ep=None`` respects the bundle's ``genai_config.json`` routing. The + bundle is loaded as-is (``compile=False``); building or compiling a bundle is + ``winml build``'s job, not eval's. + + Raises: + ValueError: no bundle directory was provided, the path is not a + directory, or it is missing the ``genai_config.json`` / ONNX files + that mark a genai bundle. + """ + from pathlib import Path + + from ..models.winml.genai_causal_lm import WinMLGenaiCausalLM + + bundle_path = config.model_path + if not bundle_path or isinstance(bundle_path, dict): + raise ValueError( + "text-generation evaluation requires a genai bundle *directory* via " + "-m ." + ) + + bundle_dir = Path(bundle_path).expanduser() + if not bundle_dir.is_dir(): + raise ValueError(f"Genai bundle directory not found: {bundle_dir}") + if not (bundle_dir / "genai_config.json").is_file(): + raise ValueError( + f"'{bundle_dir}' is not a genai bundle: no genai_config.json found. " + "Point -m at a bundle built with 'winml build ... --device npu --ep qnn'." + ) + if not any(bundle_dir.glob("*.onnx")): + raise ValueError(f"'{bundle_dir}' contains no .onnx files; not a valid genai bundle.") + + return WinMLGenaiCausalLM( + bundle_dir, + config.ep, + device=config.device, + compile=False, + ) + + def _resolve_task(config: WinMLEvaluationConfig) -> str: """Resolve the eval task and validate it is supported. diff --git a/src/winml/modelkit/eval/text_generation_evaluator.py b/src/winml/modelkit/eval/text_generation_evaluator.py new file mode 100644 index 000000000..b8ccd7320 --- /dev/null +++ b/src/winml/modelkit/eval/text_generation_evaluator.py @@ -0,0 +1,152 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- + +"""Perplexity evaluator for causal-LM (text-generation) models. + +Does *not* go through HF's ``pipeline`` / ``evaluate`` libraries: perplexity is +scored by teacher-forcing raw corpus tokens through the model's ``forward``, so +the evaluator only needs a model honoring the causal-LM contract +(``encode(text) -> list[int]`` and ``forward(ids).logits``). The same code +scores a WinML genai bundle or any object exposing that interface. + +Protocol: the dataset text column is concatenated and tokenized with the +model's own tokenizer, capped at ``num_tokens``, then cut into contiguous +non-overlapping ``seqlen``-token blocks (no detokenizer, no sliding window). +Every token after the first in its block is scored once, giving +``perplexity = exp(sum(NLL) / scored_positions)``. +""" + +from __future__ import annotations + +import logging +from typing import TYPE_CHECKING, Any + +import numpy as np + +from .base_evaluator import WinMLEvaluator + + +if TYPE_CHECKING: + from transformers.pipelines.base import Pipeline + + +logger = logging.getLogger(__name__) + +__all__ = ["WinMLTextGenerationEvaluator"] + + +class WinMLTextGenerationEvaluator(WinMLEvaluator): + """Evaluator computing disjoint fixed-length perplexity for causal LMs. + + Constructor keeps the standard ``(config, model)`` signature so the registry + dispatch in :mod:`~winml.modelkit.eval.evaluate` works unmodified. ``model`` + is a causal-LM inference object (e.g. + :class:`~winml.modelkit.models.winml.genai_causal_lm.WinMLGenaiCausalLM`). + + Two scoring parameters are read from ``dataset.columns_mapping`` so they ride + the existing ``--column key=value`` CLI path (defaults come from the + text-generation schema in :mod:`~winml.modelkit.utils.eval_utils`): + + * ``num_tokens`` -- total corpus tokens to score. + * ``seqlen`` -- non-overlapping block length. + """ + + _TASK = "text-generation" + + def prepare_pipeline(self) -> Pipeline | None: # type: ignore[override] + """No HF pipeline -- the model's ``forward`` is driven directly.""" + return None + + def prepare_data(self) -> list[list[int]]: + """Load, tokenize, and block the corpus into fixed-length token blocks. + + Returns a list of token-ID blocks, each at least 2 tokens long (a block + needs a first token to condition on and at least one token to score). + """ + from ..utils.eval_utils import get_default + + mapping = self.config.dataset.columns_mapping + num_tokens = int(mapping.get("num_tokens", get_default(self._TASK, "num_tokens"))) + seqlen = int(mapping.get("seqlen", get_default(self._TASK, "seqlen"))) + if seqlen < 2: + raise ValueError(f"seqlen must be at least 2; got {seqlen}.") + + ids = self._load_corpus_tokens(num_tokens) + blocks = [ids[i : i + seqlen] for i in range(0, len(ids), seqlen)] + blocks = [b for b in blocks if len(b) >= 2] + if not blocks: + raise ValueError( + f"Corpus produced no scorable blocks (got {len(ids)} tokens, " + f"seqlen={seqlen}). Increase num_tokens or lower seqlen." + ) + self._seqlen = seqlen + logger.info( + "Perplexity corpus: %d tokens -> %d blocks (seqlen=%d)", + len(ids), + len(blocks), + seqlen, + ) + return blocks + + def compute(self) -> dict[str, Any]: + """Score every block and return perplexity plus corpus statistics.""" + model = self.model + total_nll = 0.0 + scored = 0 + for block in self.data: + logits = model.forward(block).logits[0] + targets = np.asarray(block[1:], dtype=np.int64) + total_nll += _block_nll(logits, targets) + scored += len(targets) + + if scored == 0: + raise RuntimeError("Perplexity evaluation scored 0 positions.") + + return { + "perplexity": float(np.exp(total_nll / scored)), + "num_scored_positions": scored, + "num_blocks": len(self.data), + "seqlen": self._seqlen, + } + + # ------------------------------------------------------------------ + # Internals + # ------------------------------------------------------------------ + + def _load_corpus_tokens(self, num_tokens: int) -> list[int]: + """Concatenate the dataset text column and tokenize with the model. + + Uses the model's own tokenizer (``model.encode``) so the token stream + matches the model under test exactly. + """ + from datasets import load_dataset + + from ..utils.eval_utils import get_default + + ds_config = self.config.dataset + column = ds_config.columns_mapping.get( + "input_column", get_default(self._TASK, "input_column") + ) + dataset = load_dataset( + ds_config.path, + name=ds_config.name, + split=ds_config.split, + revision=ds_config.revision, + ) + if column not in dataset.column_names: + raise ValueError( + f"Dataset '{ds_config.path}' has no column '{column}'; " + f"available columns: {sorted(dataset.column_names)}. " + "Set it via --column input_column=." + ) + text = "\n\n".join(row for row in dataset[column] if row and row.strip()) + return self.model.encode(text)[:num_tokens] + + +def _block_nll(logits: np.ndarray, targets: np.ndarray) -> float: + """Sum of ``-log P(target)`` over positions, from raw (unnormalized) logits.""" + x = logits.astype(np.float64) + logsumexp = x.max(axis=-1) + np.log(np.exp(x - x.max(axis=-1, keepdims=True)).sum(axis=-1)) + return float((logsumexp - x[np.arange(len(targets)), targets]).sum()) diff --git a/src/winml/modelkit/models/winml/__init__.py b/src/winml/modelkit/models/winml/__init__.py index bd1ac303b..864517dbd 100644 --- a/src/winml/modelkit/models/winml/__init__.py +++ b/src/winml/modelkit/models/winml/__init__.py @@ -198,6 +198,7 @@ def register_specialization(model_type: str, task: str, class_name: str) -> None register_genai_bundle, resolve_genai_bundle, ) +from .genai_causal_lm import CausalLMOutput, HFCausalLM, WinMLGenaiCausalLM from .image_classification import WinMLModelForImageClassification from .image_segmentation import ( ImageSegmentationOutput, @@ -219,15 +220,18 @@ def register_specialization(model_type: str, task: str, class_name: str) -> None "GENAI_BUNDLE_REGISTRY", "TASK_TO_WINML_CLASS", "WINML_MODEL_CLASS_MAPPING", + "CausalLMOutput", "GenaiBundleRecipe", "GenaiCompanionSpec", "GenaiTarget", "GenaiTransformerSpec", + "HFCausalLM", "ImageSegmentationOutput", "WinMLCache", "WinMLCompositeModel", "WinMLDecoderOnlyModel", "WinMLEncoderDecoderModel", + "WinMLGenaiCausalLM", "WinMLModelForDepthEstimation", "WinMLModelForFeatureExtraction", "WinMLModelForGenericTask", diff --git a/src/winml/modelkit/models/winml/genai_causal_lm.py b/src/winml/modelkit/models/winml/genai_causal_lm.py new file mode 100644 index 000000000..42c99a0c4 --- /dev/null +++ b/src/winml/modelkit/models/winml/genai_causal_lm.py @@ -0,0 +1,231 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- +"""Causal-LM inference over an onnxruntime-genai bundle.""" + +from __future__ import annotations + +import logging +from dataclasses import dataclass +from typing import TYPE_CHECKING + +import numpy as np + +from ...session import GenaiSession, GenaiSessionError + + +if TYPE_CHECKING: + from pathlib import Path + + from ...session import GenerationConfig + from ...utils.constants import EPNameOrAlias + +logger = logging.getLogger(__name__) + +# One-hot magnitudes used to steer the next input to the given token, applied +# only after the position's logits have been read. +_FORCE_HIGH = 1e9 +_FORCE_LOW = -1e9 + + +@dataclass +class CausalLMOutput: + """Output of :meth:`WinMLGenaiCausalLM.forward`, mirroring HF's ``.logits``.""" + + logits: np.ndarray + + +class WinMLGenaiCausalLM: + """Causal-LM inference over a genai bundle via :class:`GenaiSession`. + + Constructed with an already-resolved ``ep`` / ``device`` that pass straight + to the session, so inference runs on the same runtime and EP as + ``winml perf --runtime winml-genai``. + + Args: + bundle_dir: Path to the genai bundle directory. + ep: EP override, or ``None`` to respect the bundle's ``genai_config.json``. + device: Concrete device (``npu`` / ``gpu`` / ``cpu``) the forced EP + targets; only meaningful when *ep* is set. + context_length: Override for the static KV-cache length; ``None`` uses + the bundle's configured length. + compile: Compile to EPContext ONNX before loading, reusing a cached + ``_compiled/`` when present. Use ``False`` to load an already-compiled + bundle as-is. + compile_timeout: Seconds allowed for compilation. + verbose: Forward verbose logging to the session. + """ + + def __init__( + self, + bundle_dir: str | Path, + ep: EPNameOrAlias | None = None, + *, + device: str | None = None, + context_length: int | None = None, + compile: bool = True, + compile_timeout: int = 300, + verbose: bool = False, + ) -> None: + self._session = GenaiSession( + bundle_dir, + ep, + device=device, + context_length=context_length, + compile=compile, + compile_timeout=compile_timeout, + verbose=verbose, + ) + + # ------------------------------------------------------------------ + # Tokenizer + # ------------------------------------------------------------------ + + def encode(self, text: str) -> list[int]: + """Encode *text* to token IDs using the bundle tokenizer.""" + return self._session.encode(text) + + # ------------------------------------------------------------------ + # Inference + # ------------------------------------------------------------------ + + def forward(self, input_ids: list[int]) -> CausalLMOutput: + """Run a forward pass over *input_ids*, matching HF's ``forward``. + + Returns the next-token logits at each input position, so a caller can + score the sequence exactly as with an HF causal LM. + + The genai runtime only exposes last-position logits, so the sequence is + replayed token by token through the decode loop, pinning each next input + to the given token (reusing the KV cache, cost O(len)). The final + position, whose prediction lies beyond the input, is not computed. + + Args: + input_ids: Token IDs; at least 2 tokens and at most the bundle's + :attr:`context_length`. + + Returns: + :class:`CausalLMOutput` whose ``logits`` has shape + ``(1, seq_len - 1, vocab)``; row ``i`` predicts ``input_ids[i + 1]``. + + Raises: + ValueError: *input_ids* is too short or exceeds the context length. + """ + self._session._ensure_loaded() + + n = len(input_ids) + if n < 2: + raise ValueError(f"forward needs at least 2 tokens; got {n}.") + max_length = self._session.context_length + if max_length is not None and n > max_length: + raise ValueError( + f"Sequence of {n} tokens exceeds the bundle context length {max_length}." + ) + + import onnxruntime_genai as og + + # set_logits is only on the low-level generator, not the session's + # generate path, so reach the session's loaded model directly. + model = self._session._model + + params = og.GeneratorParams(model) + # Fixed max_length is required by the compiled pipeline; do_sample=False + # keeps the decode deterministic. + params.set_search_options(max_length=max_length, do_sample=False) + gen = og.Generator(model, params) + gen.append_tokens([input_ids[0]]) + + rows: list[np.ndarray] = [] + for i in range(1, n): + step = np.asarray(gen.get_logits())[0, -1, :].astype(np.float32) + rows.append(step) + # The last read predicts input_ids[n-1]; nothing after it is scored, + # so skip the otherwise-wasted decode step. + if i < n - 1: + target = input_ids[i] + # Pin the next input to the given token so the pass follows + # input_ids regardless of the model's own argmax. + if int(step.argmax()) != target: + forced = np.full((1, 1, step.shape[0]), _FORCE_LOW, dtype=np.float32) + forced[0, 0, target] = _FORCE_HIGH + gen.set_logits(forced) + gen.generate_next_token() + + return CausalLMOutput(logits=np.stack(rows)[None, ...]) + + __call__ = forward + + def generate( + self, + prompt: str | list[int], + config: GenerationConfig | None = None, + *, + apply_chat_template: bool = True, + ) -> str: + """Generate text from *prompt*, delegating to the underlying session. + + A string *prompt* is wrapped with the bundle's chat template by default + (matching ``winml perf``); many chat models degenerate without it. Set + *apply_chat_template* to ``False``, or pass pre-encoded token IDs, to + skip the wrapping. + """ + if apply_chat_template and isinstance(prompt, str): + try: + prompt = self._session.apply_chat_template(prompt) + except GenaiSessionError: + pass # Bundle ships no chat template; use the raw prompt. + return self._session.generate(prompt, config) + + +class HFCausalLM: + """Adapt a HuggingFace causal LM to the :class:`WinMLGenaiCausalLM` contract. + + Exposes the identical ``encode`` / ``forward`` pair, so a scorer (e.g. the + text-generation perplexity evaluator) can run an fp32 PyTorch baseline + through the exact same model-agnostic path as a genai bundle, with no + branching on the concrete model type. + + ``transformers`` and ``torch`` are imported lazily so importing this module + never requires them; only constructing / calling the adapter does. + + Args: + model_id: HuggingFace model ID or local path. + device: A ``torch.device`` (or device string) to place the model on. + """ + + def __init__(self, model_id: str, device) -> None: + from transformers import AutoModelForCausalLM, AutoTokenizer + + self._tokenizer = AutoTokenizer.from_pretrained(model_id) + self._model = AutoModelForCausalLM.from_pretrained(model_id).to(device).eval() + self._device = device + + def encode(self, text: str) -> list[int]: + """Encode *text* to token IDs. + + ``add_special_tokens=False`` keeps the raw stream, matching the genai + bundle tokenizer used by :class:`WinMLGenaiCausalLM`. + """ + return self._tokenizer(text, add_special_tokens=False)["input_ids"] + + def forward(self, input_ids: list[int]) -> CausalLMOutput: + """Run a forward pass over *input_ids*, matching HF's ``forward``. + + Returns: + :class:`CausalLMOutput` whose ``logits`` has shape + ``(1, len(input_ids) - 1, vocab)``; row ``i`` predicts + ``input_ids[i + 1]``. The trailing row (predicting past the input) + is dropped to match the genai contract. + """ + import torch + + with torch.no_grad(): + logits = self._model(input_ids=torch.tensor([input_ids], device=self._device)).logits + trimmed = logits[0, :-1, :].to(torch.float32).cpu().numpy() + return CausalLMOutput(logits=trimmed[None, ...]) + + __call__ = forward + + +__all__ = ["CausalLMOutput", "HFCausalLM", "WinMLGenaiCausalLM"] diff --git a/src/winml/modelkit/utils/eval_utils.py b/src/winml/modelkit/utils/eval_utils.py index 935202bf2..f85c672b4 100644 --- a/src/winml/modelkit/utils/eval_utils.py +++ b/src/winml/modelkit/utils/eval_utils.py @@ -423,6 +423,31 @@ class TaskSchema: roles=("image-encoder", "prompt-decoder"), ) +_TEXT_GENERATION_SCHEMA = TaskSchema( + columns=( + SchemaItem( + "input_column", + "text field the perplexity corpus is concatenated from", + default="text", + remap_hint="", + ), + ), + params=( + SchemaItem( + "num_tokens", + "total corpus tokens to score", + default="8192", + remap_hint="", + ), + SchemaItem( + "seqlen", + "non-overlapping block length (tokens)", + default="2048", + remap_hint="", + ), + ), +) + TASK_SCHEMAS: dict[str, TaskSchema] = { "image-classification": _IMAGE_CLASSIFICATION_SCHEMA, "text-classification": _TEXT_CLASSIFICATION_SCHEMA, @@ -442,6 +467,7 @@ class TaskSchema: "depth-estimation": _DEPTH_ESTIMATION_SCHEMA, "keypoint-detection": _KEYPOINT_DETECTION_SCHEMA, "mask-generation": _MASK_GENERATION_SCHEMA, + "text-generation": _TEXT_GENERATION_SCHEMA, } diff --git a/tests/unit/eval/test_text_generation_evaluator.py b/tests/unit/eval/test_text_generation_evaluator.py new file mode 100644 index 000000000..5193de55c --- /dev/null +++ b/tests/unit/eval/test_text_generation_evaluator.py @@ -0,0 +1,190 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- +"""Unit tests for WinMLTextGenerationEvaluator (disjoint fixed-length perplexity). + +The dataset load and the model's ``encode`` / ``forward`` are stubbed so no +weights or corpora are downloaded. The tests verify the corpus-blocking +protocol (``num_tokens`` / ``seqlen`` from ``columns_mapping``), the NLL/PPL +math (uniform logits over a vocab of ``V`` give perplexity ``V``), and the +module-level ``_block_nll`` numerator. +""" + +from __future__ import annotations + +import math +from unittest.mock import MagicMock, patch + +import numpy as np +import pytest + +from winml.modelkit.eval import WinMLTextGenerationEvaluator +from winml.modelkit.eval.text_generation_evaluator import _block_nll + + +class _FakeDataset: + """Minimal stand-in for a HF ``Dataset`` exposing one text column.""" + + def __init__(self, rows, column="text") -> None: + self.column_names = [column] + self._rows = {column: rows} + + def __getitem__(self, column): + return self._rows[column] + + +def _uniform_forward(vocab): + """Forward returning all-zero (uniform) logits of shape (1, len-1, vocab).""" + + def forward(block): + out = MagicMock() + out.logits = np.zeros((1, len(block) - 1, vocab), dtype=np.float32) + return out + + return forward + + +def _make_evaluator( + *, + corpus_tokens, + columns_mapping=None, + corpus_rows=None, + vocab=50, + forward=None, +): + from winml.modelkit.eval import DatasetConfig, WinMLEvaluationConfig + + model = MagicMock() + model.encode.return_value = corpus_tokens + model.forward.side_effect = forward or _uniform_forward(vocab) + + fake_ds = _FakeDataset(["dummy"] if corpus_rows is None else corpus_rows) + config = WinMLEvaluationConfig( + model_id="test/mock-lm", + task="text-generation", + dataset=DatasetConfig( + path="wikitext", + name="wikitext-2-raw-v1", + split="test", + columns_mapping=columns_mapping or {}, + ), + ) + with patch("datasets.load_dataset", return_value=fake_ds): + return WinMLTextGenerationEvaluator(config, model) + + +class TestPreparePipeline: + def test_returns_none(self) -> None: + evaluator = _make_evaluator(corpus_tokens=list(range(10))) + assert evaluator.pipe is None + + +class TestPrepareData: + def test_defaults_from_schema(self) -> None: + """Empty mapping -> num_tokens=8192, seqlen=2048 (single small block).""" + evaluator = _make_evaluator(corpus_tokens=list(range(20))) + assert evaluator._seqlen == 2048 + assert evaluator.data == [list(range(20))] + + def test_blocks_by_seqlen(self) -> None: + evaluator = _make_evaluator( + corpus_tokens=list(range(10)), + columns_mapping={"num_tokens": "10", "seqlen": "4"}, + ) + assert evaluator.data == [[0, 1, 2, 3], [4, 5, 6, 7], [8, 9]] + + def test_truncates_to_num_tokens(self) -> None: + evaluator = _make_evaluator( + corpus_tokens=list(range(100)), + columns_mapping={"num_tokens": "6", "seqlen": "3"}, + ) + assert evaluator.data == [[0, 1, 2], [3, 4, 5]] + + def test_drops_trailing_single_token_block(self) -> None: + """A trailing block with < 2 tokens can't be scored and is dropped.""" + evaluator = _make_evaluator( + corpus_tokens=list(range(9)), + columns_mapping={"num_tokens": "9", "seqlen": "4"}, + ) + assert evaluator.data == [[0, 1, 2, 3], [4, 5, 6, 7]] + + def test_seqlen_below_two_raises(self) -> None: + with pytest.raises(ValueError, match="seqlen must be at least 2"): + _make_evaluator( + corpus_tokens=list(range(10)), + columns_mapping={"seqlen": "1"}, + ) + + def test_no_scorable_blocks_raises(self) -> None: + with pytest.raises(ValueError, match="no scorable blocks"): + _make_evaluator( + corpus_tokens=[0], + columns_mapping={"num_tokens": "1", "seqlen": "4"}, + ) + + +class TestLoadCorpusTokens: + def test_missing_input_column_raises(self) -> None: + with pytest.raises(ValueError, match="no column 'body'"): + _make_evaluator( + corpus_tokens=list(range(10)), + corpus_rows=["a", "b"], + columns_mapping={"input_column": "body"}, + ) + + def test_joins_nonblank_rows_with_double_newline(self) -> None: + evaluator = _make_evaluator( + corpus_tokens=list(range(10)), + corpus_rows=["a", "", " ", "b"], + ) + evaluator.model.encode.assert_called_once_with("a\n\nb") + + +class TestCompute: + def test_uniform_logits_give_perplexity_equal_vocab(self) -> None: + vocab = 50 + evaluator = _make_evaluator( + corpus_tokens=list(range(10)), + columns_mapping={"num_tokens": "10", "seqlen": "5"}, + vocab=vocab, + ) + result = evaluator.compute() + assert result["perplexity"] == pytest.approx(vocab, rel=1e-6) + assert result["num_scored_positions"] == 8 # (5-1) * 2 blocks + assert result["num_blocks"] == 2 + assert result["seqlen"] == 5 + + def test_one_forward_per_block(self) -> None: + evaluator = _make_evaluator( + corpus_tokens=list(range(12)), + columns_mapping={"num_tokens": "12", "seqlen": "4"}, + ) + evaluator.compute() + assert evaluator.model.forward.call_count == 3 + + def test_zero_scored_positions_raises(self) -> None: + evaluator = _make_evaluator(corpus_tokens=list(range(10))) + evaluator.data = [] + with pytest.raises(RuntimeError, match="scored 0 positions"): + evaluator.compute() + + +class TestBlockNLL: + def test_uniform_logits_give_log_vocab_per_position(self) -> None: + vocab, positions = 20, 3 + logits = np.zeros((positions, vocab), dtype=np.float32) + targets = np.array([0, 1, 2], dtype=np.int64) + assert _block_nll(logits, targets) == pytest.approx(positions * math.log(vocab)) + + def test_matches_reference_cross_entropy(self) -> None: + rng = np.random.default_rng(0) + logits = rng.standard_normal((4, 7)).astype(np.float32) + targets = np.array([1, 3, 0, 6], dtype=np.int64) + + x = logits.astype(np.float64) + logsumexp = np.log(np.exp(x - x.max(axis=-1, keepdims=True)).sum(axis=-1)) + logsumexp += x.max(axis=-1) + expected = float((logsumexp - x[np.arange(len(targets)), targets]).sum()) + + assert _block_nll(logits, targets) == pytest.approx(expected) diff --git a/tests/unit/models/winml/test_genai_causal_lm.py b/tests/unit/models/winml/test_genai_causal_lm.py new file mode 100644 index 000000000..d64f33bdc --- /dev/null +++ b/tests/unit/models/winml/test_genai_causal_lm.py @@ -0,0 +1,99 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- +"""Unit tests for the causal-LM inference classes in ``genai_causal_lm``. + +``HFCausalLM`` is the PyTorch-baseline adapter that honours the same +``encode`` / ``forward`` contract as ``WinMLGenaiCausalLM``. The HF tokenizer +and model are stubbed so no weights are downloaded; the tests verify the +adapter maps onto the contract exactly (``add_special_tokens=False`` encoding, +float32 numpy logits trimmed to ``(1, N - 1, vocab)``). +""" + +from __future__ import annotations + +from unittest.mock import MagicMock, patch + +import numpy as np +import torch + +from winml.modelkit.models.winml import CausalLMOutput, HFCausalLM + + +def _make_adapter(*, logits=None, token_ids=None): + """Build an ``HFCausalLM`` with stubbed tokenizer/model (no download).""" + tokenizer = MagicMock() + tokenizer.return_value = {"input_ids": [5, 6, 7] if token_ids is None else token_ids} + + model = MagicMock() + # from_pretrained(...).to(device).eval() must yield the same stub. + model.to.return_value = model + model.eval.return_value = model + call_out = MagicMock() + call_out.logits = logits + model.return_value = call_out + + with ( + patch("transformers.AutoTokenizer.from_pretrained", return_value=tokenizer), + patch("transformers.AutoModelForCausalLM.from_pretrained", return_value=model), + ): + adapter = HFCausalLM("dummy/model", torch.device("cpu")) + return adapter, tokenizer, model + + +class TestCausalLMOutput: + def test_holds_logits(self) -> None: + arr = np.zeros((1, 3, 4), dtype=np.float32) + assert CausalLMOutput(logits=arr).logits is arr + + +class TestEncode: + def test_returns_token_ids(self) -> None: + adapter, _, _ = _make_adapter(token_ids=[10, 20, 30]) + assert adapter.encode("hello world") == [10, 20, 30] + + def test_disables_special_tokens(self) -> None: + """The genai bundle tokenizer adds no specials; the adapter must match.""" + adapter, tokenizer, _ = _make_adapter() + adapter.encode("some text") + tokenizer.assert_called_once_with("some text", add_special_tokens=False) + + +class TestForward: + def test_output_is_causal_lm_output(self) -> None: + logits = torch.zeros(1, 3, 5) + adapter, _, _ = _make_adapter(logits=logits) + out = adapter.forward([1, 2, 3]) + assert isinstance(out, CausalLMOutput) + + def test_trims_trailing_row(self) -> None: + """Row predicting past the input is dropped: shape becomes (1, N-1, V).""" + vocab = 5 + logits = torch.arange(3 * vocab, dtype=torch.float32).reshape(1, 3, vocab) + adapter, _, _ = _make_adapter(logits=logits) + out = adapter.forward([1, 2, 3]) + assert out.logits.shape == (1, 2, vocab) + + def test_logits_match_raw_model(self) -> None: + vocab = 4 + logits = torch.arange(3 * vocab, dtype=torch.float32).reshape(1, 3, vocab) + adapter, _, _ = _make_adapter(logits=logits) + out = adapter.forward([7, 8, 9]) + np.testing.assert_allclose(out.logits[0], logits[0, :-1, :].numpy()) + + def test_casts_to_float32(self) -> None: + logits = torch.zeros(1, 3, 5, dtype=torch.float16) + adapter, _, _ = _make_adapter(logits=logits) + out = adapter.forward([1, 2, 3]) + assert out.logits.dtype == np.float32 + + def test_feeds_input_ids_as_batched_tensor(self) -> None: + logits = torch.zeros(1, 3, 5) + adapter, _, model = _make_adapter(logits=logits) + adapter.forward([11, 22, 33]) + passed = model.call_args.kwargs["input_ids"] + assert torch.equal(passed, torch.tensor([[11, 22, 33]])) + + def test_call_is_forward(self) -> None: + assert HFCausalLM.__call__ is HFCausalLM.forward