diff --git a/src/winml/modelkit/datasets/__init__.py b/src/winml/modelkit/datasets/__init__.py index 66bee45de..1b198c483 100644 --- a/src/winml/modelkit/datasets/__init__.py +++ b/src/winml/modelkit/datasets/__init__.py @@ -24,6 +24,7 @@ from .mask_generation import MaskGenerationDataset from .object_detection import DEFAULT_OBJECT_DETECTION_SIZE, ObjectDetectionDataset from .processor_utils import get_image_processor_config +from .prompt_dataset import PromptDataset, PromptRecord from .random_dataset import RandomDataset from .text import TextDataset diff --git a/src/winml/modelkit/datasets/base.py b/src/winml/modelkit/datasets/base.py index ea73e7bf6..b8d1fe410 100644 --- a/src/winml/modelkit/datasets/base.py +++ b/src/winml/modelkit/datasets/base.py @@ -25,7 +25,10 @@ class BaseTaskDataset(ABC): properties are readonly to ensure dataset immutability and thread safety. Attributes: - model_name: HuggingFace model identifier or local model path + model_name: HuggingFace model identifier or local model path. + Optional — task-agnostic data sources (e.g. prompt corpora used + by generative-model evaluators) may leave this ``None``. Most + calibration and task-oriented subclasses require it. dataset_name: Dataset identifier (HF dataset or local path) data_split: Dataset split to use (e.g., 'train', 'validation', 'test') """ @@ -35,7 +38,7 @@ class BaseTaskDataset(ABC): def __init__( self, - model_name: str, + model_name: str | None = None, dataset_name: str | None = None, max_samples: int | None = None, data_split: str | None = None, @@ -44,7 +47,10 @@ def __init__( """Initialize dataset with readonly properties. Args: - model_name: HuggingFace model identifier or path + model_name: HuggingFace model identifier or path. Optional — + task-agnostic subclasses may pass ``None``. Task-oriented + subclasses that need a tokenizer/processor should validate + its presence themselves. dataset_name: Dataset name (uses DEFAULT_DATASET if None) max_samples: Maximum number of samples (None = use all) data_split: Dataset split (None = let subclass decide) @@ -78,8 +84,12 @@ def _initialize(self) -> None: # Readonly properties @property - def model_name(self) -> str: - """Get the model name (readonly).""" + def model_name(self) -> str | None: + """Get the model name (readonly). + + ``None`` for task-agnostic subclasses that are not bound to a + specific model. + """ return self._model_name @property @@ -92,7 +102,6 @@ def data_split(self) -> str | None: """Get the dataset split (readonly).""" return self._data_split - def __len__(self) -> int: """Return the number of samples in the dataset.""" if self._dataset is None: diff --git a/src/winml/modelkit/datasets/prompt_dataset.py b/src/winml/modelkit/datasets/prompt_dataset.py new file mode 100644 index 000000000..91a68dfaf --- /dev/null +++ b/src/winml/modelkit/datasets/prompt_dataset.py @@ -0,0 +1,466 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- +"""Prompt-corpus dataset for evaluators that iterate over text prompts. + +``PromptDataset`` is a task-agnostic data source: unlike the calibration- +oriented ``TextDataset`` (which yields tokenized tensors bound to a +specific model) or the image-primary ``MaskGenerationDataset`` (which +carries a prompt column alongside images), ``PromptDataset``'s primary +yield is a bare prompt record. It is intended for evaluators that consume +prompts to produce something and score the output against a reference +(e.g. zero-shot classification with textual labels, VLM caption/QA +evaluation, retrieval, future generative-image workflows). + +Records +------- + +Each item yielded by ``PromptDataset`` is a plain ``dict`` with these keys:: + + { + "prompt": str, # required + "negative_prompt": str | None, # optional + "reference_text": str | None, # optional + "reference_image": str | None, # optional (path or URL) + "metadata": dict[str, Any], # optional, arbitrary + } + +The dict shape (rather than a dataclass) matches the convention used by +sibling datasets in this package. A parallel :class:`PromptRecord` +dataclass is provided for callers that prefer a typed representation +during construction. + +Loading +------- + +Three constructors are supported:: + + PromptDataset.from_list([{"prompt": "..."}, ...]) + PromptDataset.from_jsonl("path/to/prompts.jsonl") + PromptDataset.from_hf("dataset/repo", split="test", prompt_col="prompt") + +``from_hf`` also accepts a pre-loaded ``datasets.Dataset`` object in +place of the name, which makes it straightforward to test without hitting +the network. +""" + +from __future__ import annotations + +import json +import logging +from dataclasses import dataclass, field +from pathlib import Path +from random import Random +from typing import TYPE_CHECKING, Any + +from .base import BaseTaskDataset + + +if TYPE_CHECKING: + from collections.abc import Callable, Iterable, Iterator + +logger = logging.getLogger(__name__) + + +# --------------------------------------------------------------------------- +# Record type +# --------------------------------------------------------------------------- + + +@dataclass(frozen=True) +class PromptRecord: + """Typed record for a single prompt-corpus entry. + + Callers may construct :class:`PromptDataset` from a list of + :class:`PromptRecord` instances (via :meth:`PromptDataset.from_list`) + or from plain dicts \u2014 they are coerced to the same internal shape. + """ + + prompt: str + negative_prompt: str | None = None + reference_text: str | None = None + reference_image: str | None = None + metadata: dict[str, Any] = field(default_factory=dict) + + def to_dict(self) -> dict[str, Any]: + """Return the record as a dict with all keys present.""" + return { + "prompt": self.prompt, + "negative_prompt": self.negative_prompt, + "reference_text": self.reference_text, + "reference_image": self.reference_image, + "metadata": dict(self.metadata), + } + + +_ALLOWED_KEYS = frozenset( + {"prompt", "negative_prompt", "reference_text", "reference_image", "metadata"} +) + + +def _coerce_record(raw: Any) -> dict[str, Any]: + """Coerce a raw input (dict or PromptRecord) into the canonical dict shape. + + Raises: + ValueError: If ``prompt`` is missing / empty, or if unknown keys + are present at the top level. + TypeError: If ``raw`` is not a dict or PromptRecord. + """ + if isinstance(raw, PromptRecord): + return raw.to_dict() + if not isinstance(raw, dict): + raise TypeError( + f"PromptDataset record must be a dict or PromptRecord, got {type(raw).__name__}" + ) + + unknown = set(raw) - _ALLOWED_KEYS + if unknown: + raise ValueError( + f"PromptDataset record has unknown keys {sorted(unknown)}; " + f"allowed keys are {sorted(_ALLOWED_KEYS)}", + ) + + prompt = raw.get("prompt") + if not isinstance(prompt, str) or not prompt: + raise ValueError( + "PromptDataset record requires a non-empty string 'prompt' field", + ) + + metadata = raw.get("metadata", {}) or {} + if not isinstance(metadata, dict): + raise TypeError( + f"PromptDataset record 'metadata' must be a dict, got {type(metadata).__name__}" + ) + + return { + "prompt": prompt, + "negative_prompt": raw.get("negative_prompt"), + "reference_text": raw.get("reference_text"), + "reference_image": raw.get("reference_image"), + "metadata": dict(metadata), + } + + +# --------------------------------------------------------------------------- +# Dataset +# --------------------------------------------------------------------------- + + +class PromptDataset(BaseTaskDataset): + """Prompt-corpus dataset for evaluators. + + Not intended as a calibration dataset \u2014 use :class:`TextDataset` for + that. Not registered in ``TASK_DATASET_MAPPING``. + + Typical usage:: + + ds = PromptDataset.from_jsonl("prompts/drawbench.jsonl") + for sample in ds: + image = pipeline(sample["prompt"]) + score = clip_score(image, sample["reference_text"]) + """ + + def __init__( + self, + records: Iterable[PromptRecord | dict[str, Any]], + *, + model_name: str | None = None, + dataset_name: str | None = None, + max_samples: int | None = None, + data_split: str | None = None, + **kwargs: Any, + ) -> None: + """Initialize from an iterable of records. + + Args: + records: Iterable of :class:`PromptRecord` or plain dicts. + Each record is validated and coerced to the canonical shape. + model_name: Optional. PromptDataset is task-agnostic and does + not require a model. Accepted for API consistency with the + base class. + dataset_name: Optional label describing the prompt corpus + (e.g. ``"drawbench"``); surfaced by ``dataset_name``. + max_samples: Truncate the record list to at most this many + entries. Applied after validation. + data_split: Optional split label, purely informational. + **kwargs: Forwarded to the base class. + + Raises: + ValueError: If any record fails validation or if ``records`` + is empty. + """ + # Coerce and validate before the base class initialises so that + # _initialize can just assign the list. + self._raw_records = list(records) + + super().__init__( + model_name=model_name, + dataset_name=dataset_name, + max_samples=max_samples, + data_split=data_split, + **kwargs, + ) + + def _initialize(self) -> None: + coerced = [_coerce_record(r) for r in self._raw_records] + if not coerced: + raise ValueError("PromptDataset requires at least one record") + if self._max_samples is not None and self._max_samples < len(coerced): + coerced = coerced[: self._max_samples] + self._dataset = coerced + self._metadata = { + "num_records": len(coerced), + "source": self._dataset_name, + } + # Release the pre-init buffer. + del self._raw_records + + # --------------------------------------------------------------------- + # Loaders + # --------------------------------------------------------------------- + + @classmethod + def from_list( + cls, + records: Iterable[PromptRecord | dict[str, Any]], + **kwargs: Any, + ) -> PromptDataset: + """Construct from an in-memory iterable of records.""" + return cls(records, **kwargs) + + @classmethod + def from_jsonl( + cls, + path: str | Path, + **kwargs: Any, + ) -> PromptDataset: + """Construct from a JSONL file (one JSON object per line).""" + path = Path(path) + if not path.exists(): + raise FileNotFoundError(f"PromptDataset JSONL not found: {path}") + logger.info("Loading prompt corpus from JSONL: %s", path) + records: list[dict[str, Any]] = [] + with path.open("r", encoding="utf-8") as fh: + for line_no, line in enumerate(fh, start=1): + line = line.strip() + if not line: + continue + try: + records.append(json.loads(line)) + except json.JSONDecodeError as e: + raise ValueError( + f"PromptDataset JSONL parse error at {path}:{line_no}: {e.msg}", + ) from e + kwargs.setdefault("dataset_name", str(path)) + return cls(records, **kwargs) + + @classmethod + def from_hf( + cls, + dataset_or_name: Any, + *, + split: str = "train", + prompt_col: str = "prompt", + negative_prompt_col: str | None = None, + reference_text_col: str | None = None, + reference_image_col: str | None = None, + metadata_cols: list[str] | None = None, + **kwargs: Any, + ) -> PromptDataset: + """Construct from a HuggingFace dataset repo name or a pre-loaded dataset. + + Accepts either a HF dataset repo id (calls ``load_dataset``) or an + already-loaded ``datasets.Dataset`` object. The object path lets + callers pre-load or mock without hitting the network (useful in tests). + + Args: + dataset_or_name: Either a HF dataset repo id (``str``) or a + pre-loaded ``datasets.Dataset``. Accepting the object + directly lets callers pre-load or mock without invoking + ``load_dataset`` (useful in tests). + split: Split name (used only when passed a name string). + prompt_col: Source column mapped to ``prompt``. + negative_prompt_col, reference_text_col, reference_image_col: + Optional source columns for the corresponding record fields. + metadata_cols: Extra source columns to collect into ``metadata``. + **kwargs: Forwarded to :class:`PromptDataset`. + + Raises: + KeyError: If ``prompt_col`` (or any mapped column) is missing + from the dataset schema. + """ + if isinstance(dataset_or_name, str): + from datasets import load_dataset + + logger.info( + "Loading prompt corpus from HF dataset: %s (split=%s, prompt_col=%s)", + dataset_or_name, + split, + prompt_col, + ) + hf_dataset = load_dataset(dataset_or_name, split=split) + kwargs.setdefault("dataset_name", dataset_or_name) + kwargs.setdefault("data_split", split) + else: + hf_dataset = dataset_or_name + + col_names = set(getattr(hf_dataset, "column_names", []) or []) + required_cols = [prompt_col] + optional_col_map = { + "negative_prompt": negative_prompt_col, + "reference_text": reference_text_col, + "reference_image": reference_image_col, + } + required_cols.extend(src for src in optional_col_map.values() if src is not None) + required_cols.extend(metadata_cols or []) + + if col_names: + missing = [c for c in required_cols if c not in col_names] + if missing: + raise KeyError( + f"PromptDataset.from_hf: missing columns {missing} " + f"in dataset with columns {sorted(col_names)}", + ) + + records: list[dict[str, Any]] = [] + for row in hf_dataset: + record: dict[str, Any] = {"prompt": row[prompt_col]} + for target, src in optional_col_map.items(): + if src is not None: + value = row.get(src) if hasattr(row, "get") else row[src] + record[target] = value + if metadata_cols: + record["metadata"] = {c: row[c] for c in metadata_cols} + records.append(record) + + return cls(records, **kwargs) + + # --------------------------------------------------------------------- + # BaseTaskDataset interface + # --------------------------------------------------------------------- + + def __getitem__(self, idx: int) -> dict[str, Any]: + """Return the record at ``idx`` as a dict.""" + return dict(self._dataset[idx]) + + @property + def label_col(self) -> str: + """PromptDataset is task-agnostic and has no fixed label column. + + Returns an empty string to satisfy the abstract-property contract. + Callers should inspect ``reference_text`` / ``reference_image`` + on each record when a reference is available. + """ + return "" + + # --------------------------------------------------------------------- + # Convenience + # --------------------------------------------------------------------- + + def __iter__(self) -> Iterator[dict[str, Any]]: + """Iterate over the coerced record dicts.""" + return iter(self._dataset) + + # --------------------------------------------------------------------- + # Derived-dataset operations (filter / sample) + # + # Design notes for future contributors + # ------------------------------------ + # * These operations follow a common shape: transform the record list, + # then build a new dataset instance from the derived records. The + # :meth:`_derive` hook is the single extension point -- override it in + # a subclass to change the resulting type, inject extra metadata, or + # customise how kwargs cascade. ``filter`` and ``sample`` delegate to + # ``_derive`` so any new derived-dataset operation you add stays + # consistent with those two. + # * All operations return a *new* dataset (never mutate ``self``) -- + # this keeps ``PromptDataset`` safely reusable across evaluations and + # supports chaining (``ds.filter(...).sample(...)``). + # * Adding a new derived operation (e.g. ``deduplicate``, ``group_by``) + # is a matter of one method that transforms ``self._dataset`` into a + # new record list and calls ``self._derive(new_records)``. + # --------------------------------------------------------------------- + + def filter( + self, + predicate: Callable[[dict[str, Any]], bool], + **kwargs: Any, + ) -> PromptDataset: + """Return a new dataset keeping only records where ``predicate`` is true. + + Args: + predicate: Callable receiving a record ``dict`` (a shallow copy; + see note) and returning ``True`` to keep the record. + **kwargs: Forwarded to the derived dataset's constructor. By + default the derived dataset inherits ``model_name``, + ``dataset_name`` and ``data_split`` from ``self``; pass + them explicitly here to override. + + Raises: + ValueError: If the predicate rejects every record + (:class:`PromptDataset` requires at least one). + + Note: + The predicate receives a shallow copy of each record (matching + ``__getitem__`` semantics), so mutating top-level fields is + safe. Mutating nested containers such as ``metadata`` still + affects the source -- predicates should be read-only. + """ + kept = [record for record in self._dataset if predicate(dict(record))] + return self._derive(kept, **kwargs) + + def sample( + self, + n: int, + *, + seed: int | None = None, + **kwargs: Any, + ) -> PromptDataset: + """Return a new dataset containing ``n`` randomly-drawn records. + + Args: + n: Number of records to draw. Must satisfy + ``1 <= n <= len(self)`` -- oversampling is rejected because + ``PromptDataset`` records are unique instances; use + ``max_samples`` at construction time to enforce a size + cap. + seed: Optional integer for a reproducible draw. ``None`` uses + a fresh, non-deterministic RNG (matches + :func:`random.Random` semantics). + **kwargs: Forwarded to the derived dataset's constructor. + + Raises: + ValueError: If ``n`` is outside ``[1, len(self)]``. + """ + if n < 1: + raise ValueError(f"n must be >= 1, got {n}") + if n > len(self): + raise ValueError( + f"Cannot sample {n} records from a dataset of size {len(self)}; " + "use max_samples at construction to cap size.", + ) + rng = Random(seed) + sampled = rng.sample(list(self._dataset), n) + return self._derive(sampled, **kwargs) + + def _derive( + self, + records: list[dict[str, Any]], + **kwargs: Any, + ) -> PromptDataset: + """Construct a derived dataset from ``records`` inheriting metadata. + + Subclasses may override this to return their own type, inject + additional metadata, or customise the cascade of default kwargs. + The default preserves ``model_name`` / ``dataset_name`` / + ``data_split`` from ``self`` unless the caller has passed an + explicit override in ``kwargs``. + + Uses ``type(self)`` (not ``PromptDataset``) as the constructor so + subclasses receive an instance of themselves. + """ + kwargs.setdefault("model_name", self._model_name) + kwargs.setdefault("dataset_name", self._dataset_name) + kwargs.setdefault("data_split", self._data_split) + return type(self)(records, **kwargs) diff --git a/src/winml/modelkit/eval/image_feature_extraction_evaluator.py b/src/winml/modelkit/eval/image_feature_extraction_evaluator.py index fe250c8c6..5a57be60a 100644 --- a/src/winml/modelkit/eval/image_feature_extraction_evaluator.py +++ b/src/winml/modelkit/eval/image_feature_extraction_evaluator.py @@ -8,7 +8,10 @@ Evaluates image embedding models (e.g. DINOv2, DINO, ViT-in21k) by: 1. Extracting the CLS token embedding for each image via the pipeline. 2. Running a leave-one-out k-Nearest Neighbor classifier. - 3. Reporting kNN top-1 and top-5 accuracy. + 3. Reporting kNN top-1 and top-5 accuracy alongside standard retrieval + metrics (Recall@K and MRR) computed on the same cosine ranking -- + the numbers the SSL / embedding-quality literature (DINO, DINOv2, + MoCo, MAE) actually reports. Pipeline output contract (HF image-feature-extraction): pipe(image) -> [[[float, ...]]] shape: [1, num_tokens, hidden_dim] @@ -70,7 +73,17 @@ def align_labels(self, dataset: Dataset, ds_config: DatasetConfig) -> Dataset: return dataset def compute(self) -> dict[str, Any]: - """Run kNN evaluation and return accuracy metrics.""" + """Run kNN evaluation and return accuracy + retrieval metrics. + + Returns: + ``knn_top1_accuracy`` and ``knn_top5_accuracy`` (classification + accuracy via distance-weighted kNN majority vote), plus + ``recall_at_1`` / ``recall_at_5`` / ``recall_at_10`` (fraction + of queries whose top-K cosine neighbours contain a same-class + item) and ``mrr`` (mean reciprocal rank of the first same-class + neighbour). All accuracy figures are percentages in + ``[0, 100]``; recall and MRR are in ``[0, 1]``. + """ from .metrics.knn_accuracy import KNNAccuracyMetric embeddings: list[np.ndarray] = [] @@ -95,8 +108,62 @@ def compute(self) -> dict[str, Any]: embeddings_array = np.array(embeddings) labels_array = np.array(labels) - metric = KNNAccuracyMetric(k=10) - return metric.compute(embeddings_array, labels_array) + knn_result = KNNAccuracyMetric(k=10).compute(embeddings_array, labels_array) + retrieval_result = self._compute_retrieval_metrics(embeddings_array, labels_array) + return {**knn_result, **retrieval_result} + + @staticmethod + def _compute_retrieval_metrics( + embeddings: np.ndarray, + labels: np.ndarray, + ) -> dict[str, Any]: + """Compute Recall@{1, 5, 10} + MRR on the leave-one-out cosine ranking. + + Uses the same L2-normalisation and self-exclusion as + :class:`~winml.modelkit.eval.metrics.KNNAccuracyMetric` so the + rankings driving the two report families are consistent -- the + retrieval numbers describe *the same neighbour ordering* the kNN + classifier voted on. + + For every query, a same-class neighbour is treated as the single + relevant match (hit@K / rank-of-first-hit). This matches the + classification-as-retrieval convention used across DINO, DINOv2, + MoCo and MAE evaluations. + """ + from .metrics.mean_reciprocal_rank import MeanReciprocalRankMetric + from .metrics.recall_at_k import RecallAtKMetric + + # L2-normalise with an eps floor to guard degenerate embeddings. + # Matches KNNAccuracyMetric so the ranking is bit-identical. + norms = np.linalg.norm(embeddings, axis=1, keepdims=True) + norms = np.maximum(norms, 1e-9) + normalized = embeddings / norms + + similarity = normalized @ normalized.T + np.fill_diagonal(similarity, -np.inf) # exclude self + + # Full descending sort so MRR can find hits at arbitrary rank. + # Self naturally sits at the last position (its -inf became +inf + # under negation) so slicing off the tail drops the self entry. + ranked_indices = np.argsort(-similarity, axis=1)[:, :-1] + ranked_labels = labels[ranked_indices] + + recall = RecallAtKMetric(k_values=(1, 5, 10)) + mrr = MeanReciprocalRankMetric() + for i in range(len(labels)): + query_label = int(labels[i]) + recall.update(ranked_labels[i], query_label) + mrr.update(ranked_labels[i], query_label) + + # Merge into a flat dict; drop the duplicate ``n_samples`` keys + # (both metrics report the same count -- the outer evaluator + # already reports it via KNNAccuracyMetric-adjacent bookkeeping). + result = recall.compute() + result.pop("n_samples", None) + mrr_result = mrr.compute() + mrr_result.pop("n_samples", None) + result.update(mrr_result) + return result @staticmethod def _extract_image_embedding(raw: Any) -> np.ndarray: diff --git a/src/winml/modelkit/eval/metrics/__init__.py b/src/winml/modelkit/eval/metrics/__init__.py index 0488db527..872661b70 100644 --- a/src/winml/modelkit/eval/metrics/__init__.py +++ b/src/winml/modelkit/eval/metrics/__init__.py @@ -14,12 +14,15 @@ if TYPE_CHECKING: from .binary_segmentation import BinarySegmentationMetric from .classification import ClassificationMetric + from .clip_score import CLIPScoreMetric from .depth import DepthMetric from .keypoint import KeypointAPMetric from .knn_accuracy import KNNAccuracyMetric from .mean_average_precision import MAPMetric from .mean_iou import IGNORE_INDEX, MeanIoUMetric + from .mean_reciprocal_rank import MeanReciprocalRankMetric from .pseudo_perplexity import PseudoPerplexityMetric + from .recall_at_k import RecallAtKMetric from .spearman_correlation import SpearmanCorrelationMetric from .top_k_accuracy import TopKAccuracyMetric @@ -29,6 +32,7 @@ # that do not actually use the metric in question. _LAZY_ATTRS: dict[str, str] = { "BinarySegmentationMetric": ".binary_segmentation:BinarySegmentationMetric", + "CLIPScoreMetric": ".clip_score:CLIPScoreMetric", "ClassificationMetric": ".classification:ClassificationMetric", "DepthMetric": ".depth:DepthMetric", "KeypointAPMetric": ".keypoint:KeypointAPMetric", @@ -36,7 +40,9 @@ "KNNAccuracyMetric": ".knn_accuracy:KNNAccuracyMetric", "MAPMetric": ".mean_average_precision:MAPMetric", "MeanIoUMetric": ".mean_iou:MeanIoUMetric", + "MeanReciprocalRankMetric": ".mean_reciprocal_rank:MeanReciprocalRankMetric", "PseudoPerplexityMetric": ".pseudo_perplexity:PseudoPerplexityMetric", + "RecallAtKMetric": ".recall_at_k:RecallAtKMetric", "SpearmanCorrelationMetric": ".spearman_correlation:SpearmanCorrelationMetric", "TopKAccuracyMetric": ".top_k_accuracy:TopKAccuracyMetric", } @@ -61,13 +67,16 @@ def __dir__() -> list[str]: __all__ = [ "IGNORE_INDEX", "BinarySegmentationMetric", + "CLIPScoreMetric", "ClassificationMetric", "DepthMetric", "KNNAccuracyMetric", "KeypointAPMetric", "MAPMetric", "MeanIoUMetric", + "MeanReciprocalRankMetric", "PseudoPerplexityMetric", + "RecallAtKMetric", "SpearmanCorrelationMetric", "TopKAccuracyMetric", ] diff --git a/src/winml/modelkit/eval/metrics/clip_score.py b/src/winml/modelkit/eval/metrics/clip_score.py new file mode 100644 index 000000000..9378d8f9c --- /dev/null +++ b/src/winml/modelkit/eval/metrics/clip_score.py @@ -0,0 +1,140 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- + +"""CLIPScore metric for text-image alignment. + +Standard evaluation metric for text-to-image and image captioning workflows +(Hessel et al., "CLIPScore: A Reference-free Evaluation Metric for Image +Captioning", EMNLP 2021). For a text-image pair:: + + clip_score(text, image) = weight * max(0, cos(t_emb, i_emb)) + +where ``t_emb`` and ``i_emb`` are CLIP text and image embeddings, and +``weight`` is a fixed multiplier (2.5 in the original paper, keeping the +reported score in a ``[0, 2.5]`` range for typical positive cosines around +``[0, 1]``). + +This metric handles the scoring math only. Callers are responsible for +running whichever CLIP variant they want to obtain the embeddings -- the +metric stays model-agnostic and works for any embedding pair (image-image, +text-text, or cross-modal). +""" + +from __future__ import annotations + +import math +from typing import Any + +import numpy as np + + +# Hessel et al. (2021) report scores in a ``[0, 2.5]`` range for typical +# positive cosines in ``[0, 1]``; the ``2.5`` multiplier keeps our output +# comparable to numbers in the CLIPScore literature. +_DEFAULT_WEIGHT = 2.5 + + +class CLIPScoreMetric: + """Cosine-based text-image alignment score with CLIPScore semantics. + + Typical usage:: + + metric = CLIPScoreMetric() + for text_emb, image_emb in embedding_pairs: + metric.update(text_emb, image_emb) + result = metric.compute() + # {"clip_score_mean": 0.75, "clip_score_std": 0.1, ..., "n_samples": 100} + + Attributes: + weight: Scaling factor applied to each positive cosine (default 2.5, + matching Hessel et al. 2021). Set to ``1.0`` to report raw + positive-cosine values in ``[0, 1]``. + """ + + def __init__(self, weight: float = _DEFAULT_WEIGHT) -> None: + """Initialize the metric with a scaling weight. + + Args: + weight: Non-negative, finite multiplier applied to each positive + cosine similarity. Defaults to ``2.5`` (Hessel et al. + convention). + + Raises: + ValueError: If ``weight`` is negative, ``NaN``, or ``±inf``. + """ + if not math.isfinite(weight) or weight < 0: + raise ValueError( + f"weight must be non-negative and finite, got {weight!r}", + ) + self._weight = float(weight) + self._scores: list[float] = [] + + def update( + self, + text_embedding: np.ndarray, + image_embedding: np.ndarray, + ) -> None: + """Record one text-image pair's alignment score. + + Computes ``weight * max(0, cos(text, image))`` and accumulates it. + Zero-norm inputs (dead embeddings) score 0 -- a dead vector against + anything else has an undefined angle, so we treat it as no alignment. + + Args: + text_embedding: 1-D CLIP text embedding (or any shape that + flattens to 1-D; typically ``(D,)`` or ``(1, D)``). + image_embedding: 1-D CLIP image embedding of the same total size + as ``text_embedding``. + + Raises: + ValueError: If the two embeddings do not share a total size. + """ + text = np.asarray(text_embedding, dtype=np.float64).ravel() + image = np.asarray(image_embedding, dtype=np.float64).ravel() + if text.shape != image.shape: + raise ValueError( + f"text/image embedding size mismatch: {text.shape} vs {image.shape}", + ) + if text.size == 0: + raise ValueError("embeddings cannot be empty") + + norm_t = float(np.linalg.norm(text)) + norm_i = float(np.linalg.norm(image)) + if norm_t == 0.0 or norm_i == 0.0: + score = 0.0 + else: + cos_sim = float(np.dot(text, image) / (norm_t * norm_i)) + score = max(0.0, cos_sim) * self._weight + self._scores.append(score) + + def compute(self) -> dict[str, Any]: + """Return aggregate statistics over all recorded pairs. + + Returns: + Dictionary with ``clip_score_mean``, ``clip_score_std``, + ``clip_score_min``, ``clip_score_max`` (each rounded to 4 + decimals) and ``n_samples``. Every stat is ``None`` when no + samples have been recorded. + """ + if not self._scores: + return { + "clip_score_mean": None, + "clip_score_std": None, + "clip_score_min": None, + "clip_score_max": None, + "n_samples": 0, + } + arr = np.asarray(self._scores, dtype=np.float64) + return { + "clip_score_mean": round(float(arr.mean()), 4), + "clip_score_std": round(float(arr.std()), 4), + "clip_score_min": round(float(arr.min()), 4), + "clip_score_max": round(float(arr.max()), 4), + "n_samples": len(self._scores), + } + + def reset(self) -> None: + """Clear all accumulated scores.""" + self._scores = [] diff --git a/src/winml/modelkit/eval/metrics/mean_reciprocal_rank.py b/src/winml/modelkit/eval/metrics/mean_reciprocal_rank.py new file mode 100644 index 000000000..2211218f0 --- /dev/null +++ b/src/winml/modelkit/eval/metrics/mean_reciprocal_rank.py @@ -0,0 +1,109 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- + +"""Mean Reciprocal Rank (MRR) metric for retrieval and ranked-neighbor evaluation. + +MRR reports the mean of the reciprocal of the rank at which the first +relevant item appears in each query's ranked list:: + + MRR = (1/N) * Σᵢ 1 / rank_first_relevant(qᵢ) + +Queries whose ranked list contains no relevant item contribute ``0`` to +the mean (equivalent to treating the rank of the first relevant item as +infinity). + +Compared with :class:`~winml.modelkit.eval.metrics.RecallAtKMetric` -- +which reports whether a relevant item is anywhere in the top K -- MRR is +sensitive to the *position* of the first hit and rewards ranking a +correct answer higher. The two metrics are complementary and are +usually reported together in retrieval evaluations. +""" + +from __future__ import annotations + +from typing import Any + +import numpy as np + + +class MeanReciprocalRankMetric: + """Streaming Mean Reciprocal Rank over ranked prediction lists. + + Typical usage:: + + metric = MeanReciprocalRankMetric() + for ranked, gt in per_query: + metric.update(ranked, gt) + result = metric.compute() + # {"mrr": 0.623, "n_samples": 100} + + ``update`` accepts the same input shape as + :class:`~winml.modelkit.eval.metrics.RecallAtKMetric` so that a single + ``(ranked_predictions, ground_truth)`` stream can drive both metrics. + """ + + def __init__(self) -> None: + self._rr_sum = 0.0 + self._count = 0 + + def update( + self, + ranked_predictions: np.ndarray, + ground_truth: int | np.integer | np.ndarray | list[int] | tuple[int, ...] | set[int], + ) -> None: + """Record one query's ranked prediction list. + + Args: + ranked_predictions: 1-D array of predicted item IDs (or labels), + sorted in descending score order. + ground_truth: A single relevant ID (``int``) or a collection of + relevant IDs. Empty collections are rejected -- a query + with zero relevant items has undefined MRR. + + Raises: + ValueError: If ``ranked_predictions`` is empty or the + ground-truth collection is empty. + """ + ranked = np.asarray(ranked_predictions).ravel() + if ranked.size == 0: + raise ValueError("ranked_predictions cannot be empty") + + if isinstance(ground_truth, (int, np.integer)): + relevant: set[int] = {int(ground_truth)} + else: + relevant = {int(x) for x in np.asarray(list(ground_truth)).ravel()} + if not relevant: + raise ValueError("ground_truth must contain at least one relevant ID") + + # First-hit rank is 1-indexed. A query with no hit contributes 0 + # (equivalent to 1/infinity) and does not raise -- callers reason + # about MRR = 0 vs None distinctly. + rr = 0.0 + for position, item in enumerate(ranked, start=1): + if int(item) in relevant: + rr = 1.0 / position + break + self._rr_sum += rr + self._count += 1 + + def compute(self) -> dict[str, Any]: + """Return the mean reciprocal rank and the sample count. + + Returns: + Dictionary with ``mrr`` (rounded to 4 decimals) and + ``n_samples``. ``mrr`` is ``None`` when no samples have been + recorded. + """ + if self._count == 0: + return {"mrr": None, "n_samples": 0} + return { + "mrr": round(self._rr_sum / self._count, 4), + "n_samples": self._count, + } + + def reset(self) -> None: + """Clear the accumulated reciprocal-rank sum.""" + self._rr_sum = 0.0 + self._count = 0 diff --git a/src/winml/modelkit/eval/metrics/recall_at_k.py b/src/winml/modelkit/eval/metrics/recall_at_k.py new file mode 100644 index 000000000..b069a2f56 --- /dev/null +++ b/src/winml/modelkit/eval/metrics/recall_at_k.py @@ -0,0 +1,132 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- + +"""Recall@K metric for retrieval and ranked-neighbor evaluations. + +For each query, given a ranked list of predicted item IDs (or labels) +sorted by descending score and one or more ground-truth relevant IDs, +Recall@K reports the fraction of relevant items retrieved in the top K:: + + recall@k(query) = |relevant ∩ ranked[:k]| / |relevant| + +Two input shapes are supported: + +* **Single-relevant** -- pass an ``int`` ground truth. ``recall@k`` is + ``1.0`` if the ground-truth ID appears in ``ranked[:k]`` else ``0.0``. + This matches the classification-as-retrieval convention used in the SSL + embedding literature (DINO, DINOv2, MoCo, MAE): for each query, count a + hit if *any* neighbor with the correct label falls in the top K. +* **Multi-relevant** -- pass a 1-D array/tuple/set of relevant IDs. + ``recall@k`` becomes the classical retrieval recall. + +The metric is fully model-agnostic: it consumes ranked ID lists produced +by any similarity computation. +""" + +from __future__ import annotations + +from typing import Any + +import numpy as np + + +_DEFAULT_K_VALUES: tuple[int, ...] = (1, 5, 10) + + +class RecallAtKMetric: + """Streaming Recall@K over ranked prediction lists. + + Typical usage:: + + metric = RecallAtKMetric(k_values=(1, 5, 10)) + for ranked, gt in per_query: + metric.update(ranked, gt) + result = metric.compute() + # {"recall_at_1": 0.42, "recall_at_5": 0.71, "recall_at_10": 0.83, + # "n_samples": 100} + """ + + def __init__(self, k_values: tuple[int, ...] = _DEFAULT_K_VALUES) -> None: + """Initialize the metric with the K values to report. + + Args: + k_values: Non-empty iterable of positive integers. Duplicates are + collapsed and values are sorted ascending. Defaults to + ``(1, 5, 10)`` matching the SSL embedding-evaluation convention. + + Raises: + ValueError: If ``k_values`` is empty or contains a non-positive + value. + """ + ks = tuple(sorted({int(k) for k in k_values})) + if not ks: + raise ValueError("k_values must be a non-empty iterable of positive ints") + if any(k < 1 for k in ks): + raise ValueError(f"k_values must all be >= 1, got {sorted(k_values)}") + self._k_values: tuple[int, ...] = ks + # Sum of per-query recall values, one running sum per K. + self._recall_sums: dict[int, float] = dict.fromkeys(ks, 0.0) + self._count = 0 + + def update( + self, + ranked_predictions: np.ndarray, + ground_truth: int | np.integer | np.ndarray | list[int] | tuple[int, ...] | set[int], + ) -> None: + """Record one query's ranked prediction list. + + Args: + ranked_predictions: 1-D array of predicted item IDs (or labels), + sorted in descending score order. Anything that + ``np.asarray`` accepts as a 1-D integer array works. + ground_truth: A single relevant ID (``int``) or a collection of + relevant IDs. Passing an empty collection is rejected -- + a query with zero relevant items has undefined Recall@K. + + Raises: + ValueError: If ``ranked_predictions`` is not 1-D or the + ground-truth collection is empty. + """ + ranked = np.asarray(ranked_predictions).ravel() + if ranked.size == 0: + raise ValueError("ranked_predictions cannot be empty") + + # Normalize ground truth to a set of ints for uniform handling. + if isinstance(ground_truth, (int, np.integer)): + relevant: set[int] = {int(ground_truth)} + else: + relevant = {int(x) for x in np.asarray(list(ground_truth)).ravel()} + if not relevant: + raise ValueError("ground_truth must contain at least one relevant ID") + + total_relevant = len(relevant) + for k in self._k_values: + top_k = ranked[:k] + hits = sum(1 for item in top_k if int(item) in relevant) + self._recall_sums[k] += hits / total_relevant + self._count += 1 + + def compute(self) -> dict[str, Any]: + """Return mean Recall@K for every configured K, plus ``n_samples``. + + Returns: + Dictionary with keys ``recall_at_{k}`` (rounded to 4 decimals) for + every ``k`` in ``k_values``, plus ``n_samples``. Every recall + value is ``None`` when no samples have been recorded. + """ + if self._count == 0: + result: dict[str, Any] = {f"recall_at_{k}": None for k in self._k_values} + result["n_samples"] = 0 + return result + result = { + f"recall_at_{k}": round(self._recall_sums[k] / self._count, 4) for k in self._k_values + } + result["n_samples"] = self._count + return result + + def reset(self) -> None: + """Clear all accumulated recall sums.""" + self._recall_sums = dict.fromkeys(self._k_values, 0.0) + self._count = 0 diff --git a/tests/unit/datasets/test_prompt_dataset.py b/tests/unit/datasets/test_prompt_dataset.py new file mode 100644 index 000000000..0a49ef527 --- /dev/null +++ b/tests/unit/datasets/test_prompt_dataset.py @@ -0,0 +1,450 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- +"""Unit tests for :class:`PromptDataset` and :class:`PromptRecord`.""" + +from __future__ import annotations + +import json +from typing import TYPE_CHECKING, Any + +import pytest + +from winml.modelkit.datasets import PromptDataset, PromptRecord + + +if TYPE_CHECKING: + from pathlib import Path + + +# ============================================================================= +# PromptRecord dataclass +# ============================================================================= + + +class TestPromptRecord: + def test_defaults(self) -> None: + r = PromptRecord(prompt="a photo of a cat") + assert r.prompt == "a photo of a cat" + assert r.negative_prompt is None + assert r.reference_text is None + assert r.reference_image is None + assert r.metadata == {} + + def test_to_dict_populates_all_keys(self) -> None: + r = PromptRecord( + prompt="hi", + negative_prompt="lo", + reference_text="ref", + reference_image="ref.png", + metadata={"src": "unit"}, + ) + d = r.to_dict() + assert set(d) == { + "prompt", + "negative_prompt", + "reference_text", + "reference_image", + "metadata", + } + assert d["metadata"] == {"src": "unit"} + + def test_metadata_dict_is_copied(self) -> None: + meta = {"k": 1} + r = PromptRecord(prompt="hi", metadata=meta) + d = r.to_dict() + d["metadata"]["k"] = 99 + # Original metadata is untouched (to_dict copies). + assert meta == {"k": 1} + + +# ============================================================================= +# Construction & validation +# ============================================================================= + + +class TestConstruction: + def test_from_list_dicts(self) -> None: + ds = PromptDataset.from_list([{"prompt": "hello"}, {"prompt": "world"}]) + assert len(ds) == 2 + assert ds[0]["prompt"] == "hello" + assert ds[1]["prompt"] == "world" + + def test_from_list_records(self) -> None: + records = [PromptRecord(prompt="a"), PromptRecord(prompt="b")] + ds = PromptDataset.from_list(records) + assert len(ds) == 2 + assert ds[0]["prompt"] == "a" + + def test_from_list_mixed(self) -> None: + records = [PromptRecord(prompt="a"), {"prompt": "b"}] + ds = PromptDataset.from_list(records) + assert len(ds) == 2 + + def test_optional_fields_survive(self) -> None: + ds = PromptDataset.from_list( + [ + { + "prompt": "hi", + "negative_prompt": "lo", + "reference_text": "ref", + "reference_image": "img.png", + "metadata": {"src": "unit"}, + } + ] + ) + sample = ds[0] + assert sample["negative_prompt"] == "lo" + assert sample["reference_text"] == "ref" + assert sample["reference_image"] == "img.png" + assert sample["metadata"] == {"src": "unit"} + + def test_missing_optional_fields_default_to_none(self) -> None: + ds = PromptDataset.from_list([{"prompt": "hi"}]) + sample = ds[0] + assert sample["negative_prompt"] is None + assert sample["reference_text"] is None + assert sample["reference_image"] is None + assert sample["metadata"] == {} + + def test_model_name_optional(self) -> None: + # Passing no model_name works (task-agnostic dataset). + ds = PromptDataset.from_list([{"prompt": "hi"}]) + assert ds.model_name is None + + def test_model_name_accepted(self) -> None: + ds = PromptDataset.from_list([{"prompt": "hi"}], model_name="anything") + assert ds.model_name == "anything" + + def test_max_samples_truncates(self) -> None: + ds = PromptDataset.from_list([{"prompt": f"p{i}"} for i in range(10)], max_samples=3) + assert len(ds) == 3 + assert [ds[i]["prompt"] for i in range(3)] == ["p0", "p1", "p2"] + + def test_max_samples_beyond_size_is_noop(self) -> None: + ds = PromptDataset.from_list([{"prompt": "a"}, {"prompt": "b"}], max_samples=100) + assert len(ds) == 2 + + def test_dataset_name_stored(self) -> None: + ds = PromptDataset.from_list([{"prompt": "hi"}], dataset_name="my_corpus") + assert ds.dataset_name == "my_corpus" + + def test_data_split_stored(self) -> None: + ds = PromptDataset.from_list([{"prompt": "hi"}], data_split="test") + assert ds.data_split == "test" + + def test_iteration(self) -> None: + records = [{"prompt": f"p{i}"} for i in range(4)] + ds = PromptDataset.from_list(records) + prompts = [sample["prompt"] for sample in ds] + assert prompts == ["p0", "p1", "p2", "p3"] + + +class TestValidation: + def test_empty_records_rejected(self) -> None: + with pytest.raises(ValueError, match="at least one record"): + PromptDataset.from_list([]) + + def test_missing_prompt_rejected(self) -> None: + with pytest.raises(ValueError, match="non-empty string 'prompt'"): + PromptDataset.from_list([{}]) + + def test_empty_prompt_rejected(self) -> None: + with pytest.raises(ValueError, match="non-empty string 'prompt'"): + PromptDataset.from_list([{"prompt": ""}]) + + def test_non_string_prompt_rejected(self) -> None: + with pytest.raises(ValueError, match="non-empty string 'prompt'"): + PromptDataset.from_list([{"prompt": 42}]) + + def test_unknown_keys_rejected(self) -> None: + with pytest.raises(ValueError, match="unknown keys"): + PromptDataset.from_list([{"prompt": "hi", "surprise": 1}]) + + def test_non_dict_record_rejected(self) -> None: + with pytest.raises(TypeError, match="dict or PromptRecord"): + PromptDataset.from_list(["just a string"]) # type: ignore[list-item] + + def test_non_dict_metadata_rejected(self) -> None: + with pytest.raises(TypeError, match="metadata"): + PromptDataset.from_list([{"prompt": "hi", "metadata": "nope"}]) + + +# ============================================================================= +# from_jsonl +# ============================================================================= + + +class TestFromJsonl: + def test_basic_load(self, tmp_path: Path) -> None: + path = tmp_path / "prompts.jsonl" + lines = [ + {"prompt": "a"}, + {"prompt": "b", "negative_prompt": "c"}, + {"prompt": "d", "metadata": {"k": 1}}, + ] + path.write_text("\n".join(json.dumps(line) for line in lines), encoding="utf-8") + + ds = PromptDataset.from_jsonl(path) + assert len(ds) == 3 + assert ds[0]["prompt"] == "a" + assert ds[1]["negative_prompt"] == "c" + assert ds[2]["metadata"] == {"k": 1} + + def test_blank_lines_skipped(self, tmp_path: Path) -> None: + path = tmp_path / "prompts.jsonl" + path.write_text('{"prompt": "a"}\n\n{"prompt": "b"}\n', encoding="utf-8") + ds = PromptDataset.from_jsonl(path) + assert len(ds) == 2 + + def test_invalid_json_raises_with_line_number(self, tmp_path: Path) -> None: + path = tmp_path / "bad.jsonl" + path.write_text('{"prompt": "a"}\nnot json\n', encoding="utf-8") + with pytest.raises(ValueError, match=r"line 2|:2:"): + PromptDataset.from_jsonl(path) + + def test_missing_file_raises(self, tmp_path: Path) -> None: + with pytest.raises(FileNotFoundError, match="not found"): + PromptDataset.from_jsonl(tmp_path / "does_not_exist.jsonl") + + def test_dataset_name_defaults_to_path(self, tmp_path: Path) -> None: + path = tmp_path / "prompts.jsonl" + path.write_text('{"prompt": "a"}\n', encoding="utf-8") + ds = PromptDataset.from_jsonl(path) + assert ds.dataset_name == str(path) + + def test_str_path_accepted(self, tmp_path: Path) -> None: + path = tmp_path / "prompts.jsonl" + path.write_text('{"prompt": "a"}\n', encoding="utf-8") + ds = PromptDataset.from_jsonl(str(path)) + assert len(ds) == 1 + + +# ============================================================================= +# from_hf (network-free via pre-loaded fake dataset object) +# ============================================================================= + + +class _FakeHFDataset: + """Minimal stand-in for ``datasets.Dataset`` used to test from_hf + without invoking ``load_dataset`` (and without the ``datasets`` + library exposing a network-free construction path in tests).""" + + def __init__(self, rows: list[dict[str, Any]]) -> None: + self._rows = rows + self.column_names = sorted({k for row in rows for k in row}) + + def __iter__(self): + return iter(self._rows) + + +class TestFromHf: + def test_basic_mapping(self) -> None: + fake = _FakeHFDataset( + [ + {"question": "q1", "answer": "a1"}, + {"question": "q2", "answer": "a2"}, + ] + ) + ds = PromptDataset.from_hf( + fake, + prompt_col="question", + reference_text_col="answer", + ) + assert len(ds) == 2 + assert ds[0]["prompt"] == "q1" + assert ds[0]["reference_text"] == "a1" + # Column not mapped -> field stays None + assert ds[0]["negative_prompt"] is None + + def test_metadata_columns_collected(self) -> None: + fake = _FakeHFDataset([{"p": "hi", "src": "wiki", "difficulty": "easy"}]) + ds = PromptDataset.from_hf(fake, prompt_col="p", metadata_cols=["src", "difficulty"]) + assert ds[0]["metadata"] == {"src": "wiki", "difficulty": "easy"} + + def test_missing_prompt_col_raises(self) -> None: + fake = _FakeHFDataset([{"other": "value"}]) + with pytest.raises(KeyError, match="prompt"): + PromptDataset.from_hf(fake, prompt_col="prompt") + + def test_missing_optional_col_raises(self) -> None: + fake = _FakeHFDataset([{"prompt": "hi"}]) + with pytest.raises(KeyError, match="reference"): + PromptDataset.from_hf(fake, prompt_col="prompt", reference_text_col="reference") + + +# ============================================================================= +# BaseTaskDataset contract +# ============================================================================= + + +class TestBaseContract: + def test_len_matches_records(self) -> None: + ds = PromptDataset.from_list([{"prompt": f"p{i}"} for i in range(7)]) + assert len(ds) == 7 + + def test_getitem_returns_dict(self) -> None: + ds = PromptDataset.from_list([{"prompt": "hi"}]) + sample = ds[0] + assert isinstance(sample, dict) + # Returned dict is a copy \u2014 mutating it does not affect the dataset. + sample["prompt"] = "mutated" + assert ds[0]["prompt"] == "hi" + + def test_label_col_is_empty_string(self) -> None: + # PromptDataset is task-agnostic \u2014 no fixed label column. + ds = PromptDataset.from_list([{"prompt": "hi"}]) + assert ds.label_col == "" + + +# ============================================================================= +# Derived-dataset operations: filter / sample / _derive extension hook +# ============================================================================= + + +class TestFilter: + def test_keeps_matching_records(self) -> None: + ds = PromptDataset.from_list( + [ + {"prompt": "short"}, + {"prompt": "a much longer prompt here"}, + {"prompt": "med"}, + ] + ) + filtered = ds.filter(lambda r: len(r["prompt"]) < 10) + assert len(filtered) == 2 + assert [r["prompt"] for r in filtered] == ["short", "med"] + + def test_returns_new_instance_source_untouched(self) -> None: + ds = PromptDataset.from_list([{"prompt": "keep"}, {"prompt": "drop"}]) + filtered = ds.filter(lambda r: r["prompt"] == "keep") + assert filtered is not ds + assert len(ds) == 2 # source untouched + assert len(filtered) == 1 + + def test_filter_returns_prompt_dataset(self) -> None: + ds = PromptDataset.from_list([{"prompt": "keep"}, {"prompt": "drop"}]) + filtered = ds.filter(lambda r: r["prompt"] == "keep") + assert isinstance(filtered, PromptDataset) + + def test_metadata_inherited_by_default(self) -> None: + ds = PromptDataset.from_list( + [{"prompt": "a"}, {"prompt": "b"}], + dataset_name="my_corpus", + data_split="test", + ) + filtered = ds.filter(lambda r: True) + assert filtered.dataset_name == "my_corpus" + assert filtered.data_split == "test" + + def test_metadata_override_via_kwargs(self) -> None: + ds = PromptDataset.from_list([{"prompt": "a"}], dataset_name="orig") + filtered = ds.filter(lambda r: True, dataset_name="filtered_view") + assert filtered.dataset_name == "filtered_view" + + def test_predicate_receives_dict_copy_semantics(self) -> None: + # The predicate sees the record; mutating what it sees must not + # affect the source dataset. + ds = PromptDataset.from_list([{"prompt": "hi", "metadata": {"k": 1}}]) + + def mutating(record: dict) -> bool: + record["prompt"] = "MUTATED" + return True + + _ = ds.filter(mutating) + assert ds[0]["prompt"] == "hi" + + def test_rejects_all_raises(self) -> None: + ds = PromptDataset.from_list([{"prompt": "a"}, {"prompt": "b"}]) + # PromptDataset requires >= 1 record; filtering everything out + # surfaces the base validation error. + with pytest.raises(ValueError, match="at least one"): + ds.filter(lambda r: False) + + def test_chainable_with_sample(self) -> None: + ds = PromptDataset.from_list([{"prompt": f"p{i}"} for i in range(10)]) + result = ds.filter(lambda r: int(r["prompt"][1:]) % 2 == 0).sample(2, seed=1) + assert len(result) == 2 + for record in result: + assert int(record["prompt"][1:]) % 2 == 0 + + +class TestSample: + def test_size_matches_n(self) -> None: + ds = PromptDataset.from_list([{"prompt": f"p{i}"} for i in range(10)]) + sampled = ds.sample(3, seed=0) + assert len(sampled) == 3 + + def test_seed_reproducible(self) -> None: + ds = PromptDataset.from_list([{"prompt": f"p{i}"} for i in range(20)]) + a = list(ds.sample(5, seed=42)) + b = list(ds.sample(5, seed=42)) + assert a == b + + def test_different_seeds_differ(self) -> None: + ds = PromptDataset.from_list([{"prompt": f"p{i}"} for i in range(50)]) + a = list(ds.sample(10, seed=1)) + b = list(ds.sample(10, seed=2)) + # Astronomically unlikely to match with 50-choose-10. + assert a != b + + def test_full_size_returns_permutation(self) -> None: + ds = PromptDataset.from_list([{"prompt": f"p{i}"} for i in range(5)]) + sampled = ds.sample(5, seed=0) + assert len(sampled) == 5 + assert {r["prompt"] for r in sampled} == {"p0", "p1", "p2", "p3", "p4"} + + def test_returns_new_instance_source_untouched(self) -> None: + ds = PromptDataset.from_list([{"prompt": f"p{i}"} for i in range(4)]) + _ = ds.sample(2, seed=0) + assert len(ds) == 4 + + def test_oversample_rejected(self) -> None: + ds = PromptDataset.from_list([{"prompt": "a"}, {"prompt": "b"}]) + with pytest.raises(ValueError, match="Cannot sample 5"): + ds.sample(5, seed=0) + + def test_zero_or_negative_n_rejected(self) -> None: + ds = PromptDataset.from_list([{"prompt": "a"}]) + with pytest.raises(ValueError, match=">= 1"): + ds.sample(0) + with pytest.raises(ValueError, match=">= 1"): + ds.sample(-1) + + def test_metadata_inherited_by_default(self) -> None: + ds = PromptDataset.from_list( + [{"prompt": f"p{i}"} for i in range(5)], + dataset_name="orig", + data_split="val", + ) + sampled = ds.sample(2, seed=0) + assert sampled.dataset_name == "orig" + assert sampled.data_split == "val" + + +class TestDeriveExtensionHook: + """``_derive`` is the extension point for subclasses adding new derived- + dataset operations (dedup, group_by, etc.).""" + + def test_subclass_receives_own_type(self) -> None: + # Contract: derived datasets are instances of ``type(self)``. + class MyPromptDataset(PromptDataset): + pass + + ds = MyPromptDataset.from_list([{"prompt": "a"}, {"prompt": "b"}]) + filtered = ds.filter(lambda r: True) + sampled = ds.sample(1, seed=0) + assert isinstance(filtered, MyPromptDataset) + assert isinstance(sampled, MyPromptDataset) + + def test_subclass_can_override_derive_for_extra_metadata(self) -> None: + class TaggedPromptDataset(PromptDataset): + def _derive(self, records, **kwargs): + # Contract: subclasses can inject extra fields safely. + kwargs.setdefault("dataset_name", f"tagged:{self._dataset_name}") + return super()._derive(records, **kwargs) + + ds = TaggedPromptDataset.from_list([{"prompt": "a"}, {"prompt": "b"}], dataset_name="base") + filtered = ds.filter(lambda r: True) + assert filtered.dataset_name == "tagged:base" diff --git a/tests/unit/eval/test_clip_score_metric.py b/tests/unit/eval/test_clip_score_metric.py new file mode 100644 index 000000000..913ba918f --- /dev/null +++ b/tests/unit/eval/test_clip_score_metric.py @@ -0,0 +1,221 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- + +"""Unit tests for :class:`~winml.modelkit.eval.metrics.CLIPScoreMetric`.""" + +from __future__ import annotations + +import math + +import numpy as np +import pytest + +from winml.modelkit.eval.metrics import CLIPScoreMetric + + +# ============================================================================= +# Construction & validation +# ============================================================================= + + +class TestConstruction: + def test_default_weight_is_2_5(self) -> None: + metric = CLIPScoreMetric() + v = np.array([1.0, 0.0]) + metric.update(v, v) + # identical -> cos=1 -> score = 1 * 2.5 + assert metric.compute()["clip_score_mean"] == 2.5 + + def test_custom_weight(self) -> None: + metric = CLIPScoreMetric(weight=1.0) + v = np.array([1.0, 0.0]) + metric.update(v, v) + # identical -> cos=1 -> score = 1 * 1.0 + assert metric.compute()["clip_score_mean"] == 1.0 + + def test_zero_weight_valid(self) -> None: + metric = CLIPScoreMetric(weight=0.0) + v = np.array([1.0, 0.0]) + metric.update(v, v) + assert metric.compute()["clip_score_mean"] == 0.0 + + def test_negative_weight_rejected(self) -> None: + with pytest.raises(ValueError, match="non-negative"): + CLIPScoreMetric(weight=-1.0) + + def test_nan_weight_rejected(self) -> None: + with pytest.raises(ValueError, match="finite"): + CLIPScoreMetric(weight=float("nan")) + + def test_inf_weight_rejected(self) -> None: + with pytest.raises(ValueError, match="finite"): + CLIPScoreMetric(weight=float("inf")) + + +# ============================================================================= +# Cosine semantics — the actual scoring math +# ============================================================================= + + +class TestCosineSemantics: + def test_identical_embeddings_max_score(self) -> None: + metric = CLIPScoreMetric(weight=1.0) + v = np.array([1.0, 2.0, 3.0]) + metric.update(v, v) + assert metric.compute()["clip_score_mean"] == 1.0 + + def test_orthogonal_embeddings_zero(self) -> None: + metric = CLIPScoreMetric(weight=1.0) + metric.update(np.array([1.0, 0.0]), np.array([0.0, 1.0])) + assert metric.compute()["clip_score_mean"] == 0.0 + + def test_antiparallel_clipped_to_zero(self) -> None: + # cos = -1 -> max(0, -1) = 0 + metric = CLIPScoreMetric(weight=1.0) + v = np.array([1.0, 2.0, 3.0]) + metric.update(v, -v) + assert metric.compute()["clip_score_mean"] == 0.0 + + def test_known_cosine_value(self) -> None: + # cos([1,0], [1,1]) = 1/sqrt(2), scaled x1.0 + metric = CLIPScoreMetric(weight=1.0) + metric.update(np.array([1.0, 0.0]), np.array([1.0, 1.0])) + expected = 1.0 / math.sqrt(2) + assert metric.compute()["clip_score_mean"] == round(expected, 4) + + def test_weight_scales_positive_cosines(self) -> None: + # cos = 1/sqrt(2), weight = 2.5 -> 2.5/sqrt(2) + metric = CLIPScoreMetric(weight=2.5) + metric.update(np.array([1.0, 0.0]), np.array([1.0, 1.0])) + expected = 2.5 / math.sqrt(2) + assert metric.compute()["clip_score_mean"] == round(expected, 4) + + +# ============================================================================= +# Zero-vector handling +# ============================================================================= + + +class TestZeroVectorHandling: + def test_zero_text_scores_zero(self) -> None: + metric = CLIPScoreMetric(weight=1.0) + metric.update(np.zeros(3), np.array([1.0, 2.0, 3.0])) + assert metric.compute()["clip_score_mean"] == 0.0 + + def test_zero_image_scores_zero(self) -> None: + metric = CLIPScoreMetric(weight=1.0) + metric.update(np.array([1.0, 2.0, 3.0]), np.zeros(3)) + assert metric.compute()["clip_score_mean"] == 0.0 + + def test_both_zero_scores_zero(self) -> None: + # Both-zero is an undefined angle; treated as no alignment. + metric = CLIPScoreMetric(weight=1.0) + metric.update(np.zeros(3), np.zeros(3)) + assert metric.compute()["clip_score_mean"] == 0.0 + + +# ============================================================================= +# Shape handling +# ============================================================================= + + +class TestShapeHandling: + def test_2d_input_flattened(self) -> None: + # (1, D) and (D,) with the same D flatten to the same 1-D vector. + metric = CLIPScoreMetric(weight=1.0) + metric.update(np.array([[1.0, 2.0]]), np.array([1.0, 2.0])) + assert metric.compute()["clip_score_mean"] == 1.0 + + def test_size_mismatch_raises(self) -> None: + metric = CLIPScoreMetric() + with pytest.raises(ValueError, match="size mismatch"): + metric.update(np.array([1.0, 2.0]), np.array([1.0, 2.0, 3.0])) + + def test_empty_embedding_rejected(self) -> None: + metric = CLIPScoreMetric() + with pytest.raises(ValueError, match="empty"): + metric.update(np.array([]), np.array([])) + + def test_integer_arrays_accepted(self) -> None: + # ``update`` coerces to float64 -- ints should work. + metric = CLIPScoreMetric(weight=1.0) + metric.update(np.array([1, 0]), np.array([1, 0])) + assert metric.compute()["clip_score_mean"] == 1.0 + + +# ============================================================================= +# Aggregation over multiple samples +# ============================================================================= + + +class TestAggregation: + def test_empty_state_returns_nones(self) -> None: + metric = CLIPScoreMetric() + result = metric.compute() + assert result["clip_score_mean"] is None + assert result["clip_score_std"] is None + assert result["clip_score_min"] is None + assert result["clip_score_max"] is None + assert result["n_samples"] == 0 + + def test_batch_statistics(self) -> None: + # Two orthogonal pairs (score 0) + one identical pair (score 2.5). + metric = CLIPScoreMetric(weight=2.5) + metric.update(np.array([1.0, 0.0]), np.array([0.0, 1.0])) + metric.update(np.array([1.0, 0.0]), np.array([0.0, 1.0])) + metric.update(np.array([1.0, 0.0]), np.array([1.0, 0.0])) + result = metric.compute() + assert result["n_samples"] == 3 + # scores = [0, 0, 2.5] -> mean = 2.5 / 3 + assert result["clip_score_mean"] == round(2.5 / 3, 4) + assert result["clip_score_min"] == 0.0 + assert result["clip_score_max"] == 2.5 + + def test_reset_clears_state(self) -> None: + metric = CLIPScoreMetric() + metric.update(np.array([1.0, 0.0]), np.array([1.0, 0.0])) + assert metric.compute()["n_samples"] == 1 + metric.reset() + assert metric.compute()["n_samples"] == 0 + assert metric.compute()["clip_score_mean"] is None + + +# ============================================================================= +# Real-shaped embeddings (CLIP typically emits 512-D vectors) +# ============================================================================= + + +class TestRealisticShapes: + def test_512_dim_deterministic(self) -> None: + rng = np.random.default_rng(0) + metric = CLIPScoreMetric(weight=2.5) + + text_embeddings = rng.standard_normal(size=(10, 512)) + image_embeddings = rng.standard_normal(size=(10, 512)) + for t, i in zip(text_embeddings, image_embeddings, strict=True): + metric.update(t, i) + + result = metric.compute() + assert result["n_samples"] == 10 + # Random Gaussian pairs in 512-D are near-orthogonal on average -- + # the mean score should be modest and clipped positive. + assert 0.0 <= result["clip_score_mean"] <= 2.5 + + def test_matched_pairs_score_higher_than_random(self) -> None: + # Sanity: identical pairs beat random pairs on average. + rng = np.random.default_rng(42) + + matched = CLIPScoreMetric(weight=1.0) + random_pairs = CLIPScoreMetric(weight=1.0) + for _ in range(20): + v = rng.standard_normal(size=(512,)) + matched.update(v, v) + random_pairs.update(v, rng.standard_normal(size=(512,))) + + matched_mean = matched.compute()["clip_score_mean"] + random_mean = random_pairs.compute()["clip_score_mean"] + assert matched_mean is not None + assert random_mean is not None + assert matched_mean > random_mean diff --git a/tests/unit/eval/test_image_feature_extraction_evaluator.py b/tests/unit/eval/test_image_feature_extraction_evaluator.py index 69de16829..8da8b3767 100644 --- a/tests/unit/eval/test_image_feature_extraction_evaluator.py +++ b/tests/unit/eval/test_image_feature_extraction_evaluator.py @@ -19,6 +19,7 @@ # Helpers # --------------------------------------------------------------------------- + def make_evaluator(columns_mapping=None): """Instantiate evaluator by patching external dependencies.""" from winml.modelkit.eval import DatasetConfig, WinMLEvaluationConfig @@ -47,8 +48,10 @@ def make_evaluator(columns_mapping=None): dataset=DatasetConfig(path="timm/mini-imagenet", columns_mapping=mapping), ) - with patch("datasets.load_dataset", return_value=mock_ds), \ - patch("transformers.pipeline", return_value=mock_pipe): + with ( + patch("datasets.load_dataset", return_value=mock_ds), + patch("transformers.pipeline", return_value=mock_pipe), + ): return WinMLImageFeatureExtractionEvaluator(config, model) @@ -56,17 +59,20 @@ def make_evaluator(columns_mapping=None): # KNNAccuracyMetric # --------------------------------------------------------------------------- + class TestKNNAccuracyMetric: def test_perfect_clusters(self): """Embeddings from the same class are identical -> 100% accuracy.""" metric = KNNAccuracyMetric(k=3) # 4 samples, 2 classes. Class 0 at origin-ish, class 1 far away. - embeddings = np.array([ - [1.0, 0.0, 0.0], - [0.99, 0.01, 0.0], - [0.0, 0.0, 1.0], - [0.01, 0.0, 0.99], - ]) + embeddings = np.array( + [ + [1.0, 0.0, 0.0], + [0.99, 0.01, 0.0], + [0.0, 0.0, 1.0], + [0.01, 0.0, 0.99], + ] + ) labels = np.array([0, 0, 1, 1]) result = metric.compute(embeddings, labels) assert result["knn_top1_accuracy"] == 100.0 @@ -85,11 +91,13 @@ def test_random_embeddings_returns_valid_range(self): def test_k_capped_to_n_minus_1(self): """k should be capped when larger than N-1.""" metric = KNNAccuracyMetric(k=100) - embeddings = np.array([ - [1.0, 0.0], - [0.9, 0.1], - [0.0, 1.0], - ]) + embeddings = np.array( + [ + [1.0, 0.0], + [0.9, 0.1], + [0.0, 1.0], + ] + ) labels = np.array([0, 0, 1]) # Should not raise, k capped to 2 result = metric.compute(embeddings, labels) @@ -127,6 +135,7 @@ def test_two_samples_minimal(self): # WinMLImageFeatureExtractionEvaluator # --------------------------------------------------------------------------- + class TestImageFeatureExtractionEvaluatorInit: def test_default_label_column(self): evaluator = make_evaluator() @@ -205,6 +214,7 @@ def test_default_dataset_registered(self): # WinMLImageFeatureExtractionEvaluator.compute # --------------------------------------------------------------------------- + class TestCompute: """End-to-end: pipeline output -> CLS extraction -> kNN metric.""" @@ -229,12 +239,14 @@ def test_end_to_end_flow(self): {"image": "img3", "label": 1}, {"image": "img4", "label": 1}, ] - outputs = iter([ - self._token_sequence(cluster_a), - self._token_sequence([0.99, 0.01, 0.0]), - self._token_sequence(cluster_b), - self._token_sequence([0.01, 0.99, 0.0]), - ]) + outputs = iter( + [ + self._token_sequence(cluster_a), + self._token_sequence([0.99, 0.01, 0.0]), + self._token_sequence(cluster_b), + self._token_sequence([0.01, 0.99, 0.0]), + ] + ) ev.pipe = MagicMock(side_effect=lambda _img: next(outputs)) result = ev.compute() @@ -244,20 +256,112 @@ def test_end_to_end_flow(self): # Perfectly separable clusters -> 100% top-1. assert result["knn_top1_accuracy"] == 100.0 + def test_retrieval_metrics_reported_alongside_knn(self): + """compute() emits Recall@K and MRR keys next to KNN accuracies.""" + ev = make_evaluator() + + cluster_a = [1.0, 0.0, 0.0] + cluster_b = [0.0, 1.0, 0.0] + ev.data = [ + {"image": "img1", "label": 0}, + {"image": "img2", "label": 0}, + {"image": "img3", "label": 1}, + {"image": "img4", "label": 1}, + ] + outputs = iter( + [ + self._token_sequence(cluster_a), + self._token_sequence([0.99, 0.01, 0.0]), + self._token_sequence(cluster_b), + self._token_sequence([0.01, 0.99, 0.0]), + ] + ) + ev.pipe = MagicMock(side_effect=lambda _img: next(outputs)) + + result = ev.compute() + + # Backward-compat: existing keys unchanged. + assert "knn_top1_accuracy" in result + assert "knn_top5_accuracy" in result + # New retrieval metrics reported. + assert set(result) >= { + "knn_top1_accuracy", + "knn_top5_accuracy", + "recall_at_1", + "recall_at_5", + "recall_at_10", + "mrr", + } + + def test_perfect_clusters_score_max_retrieval(self): + """Well-separated same-class pairs -> recall@1 = mrr = 1.0.""" + ev = make_evaluator() + + cluster_a = [1.0, 0.0, 0.0] + cluster_b = [0.0, 1.0, 0.0] + ev.data = [ + {"image": "img1", "label": 0}, + {"image": "img2", "label": 0}, + {"image": "img3", "label": 1}, + {"image": "img4", "label": 1}, + ] + outputs = iter( + [ + self._token_sequence(cluster_a), + self._token_sequence([0.99, 0.01, 0.0]), + self._token_sequence(cluster_b), + self._token_sequence([0.01, 0.99, 0.0]), + ] + ) + ev.pipe = MagicMock(side_effect=lambda _img: next(outputs)) + + result = ev.compute() + assert result["recall_at_1"] == 1.0 + assert result["mrr"] == 1.0 + + def test_compute_retrieval_metrics_static_helper(self): + """Static helper returns retrieval-only dict without n_samples noise.""" + # Two two-sample clusters -> each query's nearest neighbour is same-class. + embeddings = np.array( + [ + [1.0, 0.0, 0.0], + [0.9, 0.1, 0.0], + [0.0, 0.0, 1.0], + [0.0, 0.1, 0.9], + ] + ) + labels = np.array([0, 0, 1, 1]) + result = WinMLImageFeatureExtractionEvaluator._compute_retrieval_metrics(embeddings, labels) + assert set(result) == {"recall_at_1", "recall_at_5", "recall_at_10", "mrr"} + assert result["recall_at_1"] == 1.0 + assert result["mrr"] == 1.0 + + def test_no_same_class_neighbours_scores_zero(self): + """When every sample is a singleton class, recall@1 and mrr are 0.""" + # Each sample is its own class -> no same-class neighbour exists. + embeddings = np.eye(4, dtype=np.float64) + labels = np.array([0, 1, 2, 3]) + result = WinMLImageFeatureExtractionEvaluator._compute_retrieval_metrics(embeddings, labels) + assert result["recall_at_1"] == 0.0 + assert result["recall_at_5"] == 0.0 + assert result["mrr"] == 0.0 + def test_skips_samples_with_none_image_or_label(self): """Samples missing image or label are dropped before embedding.""" ev = make_evaluator() ev.data = [ {"image": "img1", "label": 0}, - {"image": None, "label": 0}, # skipped - {"image": "img2", "label": None}, # skipped + {"image": None, "label": 0}, # skipped + {"image": "img2", "label": None}, # skipped {"image": "img3", "label": 1}, ] - outputs = iter([ - self._token_sequence([1.0, 0.0]), - self._token_sequence([0.0, 1.0]), - ]) + outputs = iter( + [ + self._token_sequence([1.0, 0.0]), + self._token_sequence([0.0, 1.0]), + ] + ) ev.pipe = MagicMock(side_effect=lambda _img: next(outputs)) result = ev.compute() @@ -274,10 +378,7 @@ def test_raises_when_fewer_than_two_valid_samples(self): {"image": "img1", "label": 0}, {"image": None, "label": 0}, ] - ev.pipe = MagicMock( - return_value=self._token_sequence([1.0, 0.0]) - ) + ev.pipe = MagicMock(return_value=self._token_sequence([1.0, 0.0])) with pytest.raises(ValueError, match="at least 2 valid samples"): ev.compute() - diff --git a/tests/unit/eval/test_mean_reciprocal_rank_metric.py b/tests/unit/eval/test_mean_reciprocal_rank_metric.py new file mode 100644 index 000000000..4e30e9724 --- /dev/null +++ b/tests/unit/eval/test_mean_reciprocal_rank_metric.py @@ -0,0 +1,129 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- + +"""Unit tests for :class:`~winml.modelkit.eval.metrics.MeanReciprocalRankMetric`.""" + +from __future__ import annotations + +import numpy as np +import pytest + +from winml.modelkit.eval.metrics import MeanReciprocalRankMetric + + +# ============================================================================= +# Single-relevant semantics +# ============================================================================= + + +class TestSingleRelevant: + def test_hit_at_rank_1(self) -> None: + metric = MeanReciprocalRankMetric() + metric.update(np.array([42, 1, 2, 3]), 42) + assert metric.compute()["mrr"] == 1.0 + + def test_hit_at_rank_2(self) -> None: + metric = MeanReciprocalRankMetric() + metric.update(np.array([1, 42, 2, 3]), 42) + assert metric.compute()["mrr"] == 0.5 + + def test_hit_at_rank_3(self) -> None: + metric = MeanReciprocalRankMetric() + metric.update(np.array([1, 2, 42, 3]), 42) + # 1/3 rounded to 4 dp + assert metric.compute()["mrr"] == round(1 / 3, 4) + + def test_no_hit_scores_zero(self) -> None: + metric = MeanReciprocalRankMetric() + metric.update(np.array([1, 2, 3, 4]), 999) + assert metric.compute()["mrr"] == 0.0 + + def test_np_integer_ground_truth_accepted(self) -> None: + metric = MeanReciprocalRankMetric() + metric.update(np.array([1, 2, 3]), np.int64(1)) + assert metric.compute()["mrr"] == 1.0 + + +# ============================================================================= +# Multi-relevant semantics (uses first hit) +# ============================================================================= + + +class TestMultiRelevant: + def test_first_hit_wins(self) -> None: + # Relevant {2, 3}; ranked [1, 2, 3] -> first hit at rank 2 -> RR = 1/2. + metric = MeanReciprocalRankMetric() + metric.update(np.array([1, 2, 3]), [2, 3]) + assert metric.compute()["mrr"] == 0.5 + + def test_earliest_position_used(self) -> None: + # Relevant {5, 2}; ranked [1, 2, 3, 4, 5] -> first hit at rank 2. + metric = MeanReciprocalRankMetric() + metric.update(np.array([1, 2, 3, 4, 5]), [5, 2]) + assert metric.compute()["mrr"] == 0.5 + + def test_no_relevant_in_ranked_scores_zero(self) -> None: + metric = MeanReciprocalRankMetric() + metric.update(np.array([10, 20, 30]), [1, 2, 3]) + assert metric.compute()["mrr"] == 0.0 + + def test_set_input(self) -> None: + metric = MeanReciprocalRankMetric() + metric.update(np.array([1, 2, 3]), {3}) + # Hit at rank 3 -> 1/3 + assert metric.compute()["mrr"] == round(1 / 3, 4) + + def test_empty_ground_truth_rejected(self) -> None: + metric = MeanReciprocalRankMetric() + with pytest.raises(ValueError, match="at least one"): + metric.update(np.array([1, 2, 3]), []) + + +# ============================================================================= +# Aggregation +# ============================================================================= + + +class TestAggregation: + def test_empty_state_returns_none(self) -> None: + metric = MeanReciprocalRankMetric() + result = metric.compute() + assert result["mrr"] is None + assert result["n_samples"] == 0 + + def test_batch_mean(self) -> None: + # RR values: [1.0, 0.5, 0.0] -> mean 0.5 + metric = MeanReciprocalRankMetric() + metric.update(np.array([1, 2, 3]), 1) # rank 1 -> 1.0 + metric.update(np.array([1, 2, 3]), 2) # rank 2 -> 0.5 + metric.update(np.array([1, 2, 3]), 999) # miss -> 0.0 + result = metric.compute() + assert result["mrr"] == 0.5 + assert result["n_samples"] == 3 + + def test_reset_clears_state(self) -> None: + metric = MeanReciprocalRankMetric() + metric.update(np.array([1, 2, 3]), 1) + assert metric.compute()["n_samples"] == 1 + metric.reset() + assert metric.compute()["n_samples"] == 0 + assert metric.compute()["mrr"] is None + + +# ============================================================================= +# Shape handling +# ============================================================================= + + +class TestShapeHandling: + def test_2d_input_flattened(self) -> None: + metric = MeanReciprocalRankMetric() + metric.update(np.array([[1, 42, 2]]), 42) + assert metric.compute()["mrr"] == 0.5 + + def test_empty_ranked_predictions_rejected(self) -> None: + metric = MeanReciprocalRankMetric() + with pytest.raises(ValueError, match="cannot be empty"): + metric.update(np.array([]), 1) diff --git a/tests/unit/eval/test_recall_at_k_metric.py b/tests/unit/eval/test_recall_at_k_metric.py new file mode 100644 index 000000000..f5dbdc915 --- /dev/null +++ b/tests/unit/eval/test_recall_at_k_metric.py @@ -0,0 +1,182 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- + +"""Unit tests for :class:`~winml.modelkit.eval.metrics.RecallAtKMetric`.""" + +from __future__ import annotations + +import numpy as np +import pytest + +from winml.modelkit.eval.metrics import RecallAtKMetric + + +# ============================================================================= +# Construction & validation +# ============================================================================= + + +class TestConstruction: + def test_default_k_values(self) -> None: + # Sanity: default (1, 5, 10) surfaces in compute() output keys. + metric = RecallAtKMetric() + metric.update(np.array([1, 2, 3, 4, 5, 6, 7, 8, 9, 10]), 1) + result = metric.compute() + assert set(result) == {"recall_at_1", "recall_at_5", "recall_at_10", "n_samples"} + + def test_custom_k_values(self) -> None: + metric = RecallAtKMetric(k_values=(2, 7)) + metric.update(np.array([10, 20, 30]), 20) + result = metric.compute() + assert set(result) == {"recall_at_2", "recall_at_7", "n_samples"} + + def test_k_values_sorted_and_deduplicated(self) -> None: + # Order shouldn't matter; duplicates collapse. + metric = RecallAtKMetric(k_values=(5, 1, 5, 10)) + metric.update(np.array([1, 2, 3, 4, 5]), 1) + result = metric.compute() + # Keys sorted ascending by K. + assert list(result)[:-1] == ["recall_at_1", "recall_at_5", "recall_at_10"] + + def test_empty_k_values_rejected(self) -> None: + with pytest.raises(ValueError, match="non-empty"): + RecallAtKMetric(k_values=()) + + def test_non_positive_k_rejected(self) -> None: + with pytest.raises(ValueError, match=">= 1"): + RecallAtKMetric(k_values=(0, 5)) + with pytest.raises(ValueError, match=">= 1"): + RecallAtKMetric(k_values=(-1,)) + + +# ============================================================================= +# Single-relevant semantics (hit@k) +# ============================================================================= + + +class TestSingleRelevant: + def test_hit_at_top(self) -> None: + # gt is rank 1 -> hits all K's. + metric = RecallAtKMetric(k_values=(1, 5, 10)) + metric.update(np.array([42, 1, 2, 3, 4, 5, 6, 7, 8, 9]), 42) + result = metric.compute() + assert result["recall_at_1"] == 1.0 + assert result["recall_at_5"] == 1.0 + assert result["recall_at_10"] == 1.0 + + def test_hit_at_5_but_not_1(self) -> None: + # gt at rank 3 -> hits @5 and @10 but not @1. + metric = RecallAtKMetric(k_values=(1, 5, 10)) + metric.update(np.array([1, 2, 42, 3, 4, 5, 6, 7, 8, 9]), 42) + result = metric.compute() + assert result["recall_at_1"] == 0.0 + assert result["recall_at_5"] == 1.0 + assert result["recall_at_10"] == 1.0 + + def test_miss_all(self) -> None: + # gt not in ranked list -> zeros. + metric = RecallAtKMetric(k_values=(1, 5, 10)) + metric.update(np.array([1, 2, 3, 4, 5, 6, 7, 8, 9, 10]), 999) + result = metric.compute() + assert result["recall_at_1"] == 0.0 + assert result["recall_at_5"] == 0.0 + assert result["recall_at_10"] == 0.0 + + def test_np_integer_ground_truth_accepted(self) -> None: + # numpy integer scalar (from argmax etc.) is a common gt source. + metric = RecallAtKMetric(k_values=(1,)) + metric.update(np.array([1, 2, 3]), np.int64(1)) + assert metric.compute()["recall_at_1"] == 1.0 + + +# ============================================================================= +# Multi-relevant semantics (classical retrieval recall) +# ============================================================================= + + +class TestMultiRelevant: + def test_all_relevant_retrieved(self) -> None: + # 3 relevant items, all in top 3 -> recall@3 = 3/3 = 1.0 + metric = RecallAtKMetric(k_values=(3,)) + metric.update(np.array([1, 2, 3, 4, 5]), [1, 2, 3]) + assert metric.compute()["recall_at_3"] == 1.0 + + def test_partial_relevant_retrieved(self) -> None: + # 3 relevant items {1, 2, 3}, top-2 = [1, 2] -> recall@2 = 2/3 + metric = RecallAtKMetric(k_values=(2,)) + metric.update(np.array([1, 2, 4, 3, 5]), [1, 2, 3]) + assert metric.compute()["recall_at_2"] == round(2 / 3, 4) + + def test_no_relevant_retrieved(self) -> None: + metric = RecallAtKMetric(k_values=(3,)) + metric.update(np.array([10, 20, 30]), [1, 2, 3]) + assert metric.compute()["recall_at_3"] == 0.0 + + def test_set_ground_truth_accepted(self) -> None: + metric = RecallAtKMetric(k_values=(2,)) + metric.update(np.array([5, 3, 1]), {1, 3, 7}) + # top-2 = [5, 3], relevant overlap = {3}, |relevant| = 3 -> 1/3 + assert metric.compute()["recall_at_2"] == round(1 / 3, 4) + + def test_empty_ground_truth_rejected(self) -> None: + metric = RecallAtKMetric() + with pytest.raises(ValueError, match="at least one"): + metric.update(np.array([1, 2, 3]), []) + + +# ============================================================================= +# Aggregation over multiple queries +# ============================================================================= + + +class TestAggregation: + def test_empty_state_returns_nones(self) -> None: + metric = RecallAtKMetric(k_values=(1, 5)) + result = metric.compute() + assert result["recall_at_1"] is None + assert result["recall_at_5"] is None + assert result["n_samples"] == 0 + + def test_batch_mean(self) -> None: + # Two queries: one hit@1, one miss@1 -> mean = 0.5 + metric = RecallAtKMetric(k_values=(1, 5)) + metric.update(np.array([1, 2, 3, 4, 5]), 1) # hit @1 and @5 + metric.update(np.array([2, 3, 4, 5, 6]), 1) # miss @1, miss @5 + result = metric.compute() + assert result["recall_at_1"] == 0.5 + assert result["recall_at_5"] == 0.5 + assert result["n_samples"] == 2 + + def test_reset_clears_state(self) -> None: + metric = RecallAtKMetric(k_values=(1,)) + metric.update(np.array([1, 2, 3]), 1) + assert metric.compute()["n_samples"] == 1 + metric.reset() + assert metric.compute()["n_samples"] == 0 + assert metric.compute()["recall_at_1"] is None + + +# ============================================================================= +# Shape handling +# ============================================================================= + + +class TestShapeHandling: + def test_2d_input_flattened(self) -> None: + # A (1, K) shape from `.reshape(1, -1)` or slicing flattens. + metric = RecallAtKMetric(k_values=(1,)) + metric.update(np.array([[42, 1, 2]]), 42) + assert metric.compute()["recall_at_1"] == 1.0 + + def test_empty_ranked_predictions_rejected(self) -> None: + metric = RecallAtKMetric() + with pytest.raises(ValueError, match="cannot be empty"): + metric.update(np.array([]), 1) + + def test_k_larger_than_ranked_list(self) -> None: + # K=10 with only 3 predictions -> treats top-3 == whole list. + metric = RecallAtKMetric(k_values=(10,)) + metric.update(np.array([1, 2, 3]), 2) + assert metric.compute()["recall_at_10"] == 1.0