From d726e14f8356906551b8ab866026bd7889291b7f Mon Sep 17 00:00:00 2001 From: Abhilash Chenreddy Date: Mon, 20 Jul 2026 15:10:50 -0400 Subject: [PATCH 1/7] feat(datasets): add PromptDataset primitive for prompt-driven evaluators Existing datasets in this package are either calibration-focused (TextDataset: yields tokenized tensors for feeding into the model) or image-primary (ImageDataset/MaskGenerationDataset/etc.: yield images optionally decorated with text columns). None expose a prompt corpus as their primary yield. This PR introduces PromptDataset \u2014 a task-agnostic data source whose primary yield is a prompt record. Intended for evaluators that consume prompts to produce something and score the output against a reference (zero-shot classification with textual labels, VLM caption/QA eval, retrieval, and future generative-image workflows). Record shape (plain dict, matches sibling-dataset convention): { 'prompt': str, # required 'negative_prompt': str | None, # optional 'reference_text': str | None, # optional 'reference_image': str | None, # optional 'metadata': dict, # optional } A parallel PromptRecord dataclass is provided for typed construction. Loading: PromptDataset.from_list([{'prompt': '...'}, ...]) PromptDataset.from_jsonl('path/to/prompts.jsonl') PromptDataset.from_hf(name_or_dataset, prompt_col='prompt', ...) from_hf accepts either a HF repo id (calls load_dataset) or a pre-loaded datasets.Dataset object \u2014 the object path is what tests use to avoid network access. Base-class change: BaseTaskDataset.model_name is now Optional[str] with default None. This is backward-compatible for all existing subclasses (TextDataset, ImageDataset, ObjectDetectionDataset, KeypointDetection Dataset, MaskGenerationDataset, DepthEstimationDataset, ImageSegmentation Dataset) \u2014 they all continue to receive a model_name from callers, and their internal type expectations are unchanged. Task-agnostic subclasses (PromptDataset) can now inherit cleanly without a fake model_name. Not added to TASK_DATASET_MAPPING \u2014 PromptDataset is an evaluator data source, not a calibration dataset. Tests: tests/unit/datasets/test_prompt_dataset.py \u2014 35 tests covering PromptRecord dataclass, construction (list of dicts, list of PromptRecords, mixed), field validation (missing prompt, empty prompt, non-string prompt, unknown keys, non-dict record, non-dict metadata), from_jsonl (basic, blank lines, invalid JSON, missing file, dataset_name default, str path), from_hf (basic mapping, metadata cols, missing prompt col, missing optional col), and the BaseTaskDataset abstract-method contract. All 97 tests in tests/unit/datasets/ pass (35 new + 62 existing). --- src/winml/modelkit/datasets/__init__.py | 1 + src/winml/modelkit/datasets/base.py | 21 +- src/winml/modelkit/datasets/prompt_dataset.py | 362 ++++++++++++++++++ tests/unit/datasets/test_prompt_dataset.py | 298 ++++++++++++++ 4 files changed, 676 insertions(+), 6 deletions(-) create mode 100644 src/winml/modelkit/datasets/prompt_dataset.py create mode 100644 tests/unit/datasets/test_prompt_dataset.py 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..45512b6ea --- /dev/null +++ b/src/winml/modelkit/datasets/prompt_dataset.py @@ -0,0 +1,362 @@ +# ------------------------------------------------------------------------- +# 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 typing import TYPE_CHECKING, Any + +from .base import BaseTaskDataset + + +if TYPE_CHECKING: + from collections.abc import Iterable + +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): + """Iterate over the coerced record dicts.""" + return iter(self._dataset) diff --git a/tests/unit/datasets/test_prompt_dataset.py b/tests/unit/datasets/test_prompt_dataset.py new file mode 100644 index 000000000..b4ede6bd0 --- /dev/null +++ b/tests/unit/datasets/test_prompt_dataset.py @@ -0,0 +1,298 @@ +# ------------------------------------------------------------------------- +# 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 == "" From d0d9c6a4fbff397d47955377daf43af674b467a3 Mon Sep 17 00:00:00 2001 From: Abhilash Chenreddy Date: Tue, 4 Aug 2026 20:51:30 -0400 Subject: [PATCH 2/7] feat(eval): add CLIPScoreMetric for text-image alignment Adds `CLIPScoreMetric` in `eval/metrics/clip_score.py` following the convention of Hessel et al., 2021 (EMNLP): clip_score(text, image) = weight * max(0, cos(t_emb, i_emb)) The metric is model-agnostic: callers supply pre-computed CLIP text/image embeddings, the metric handles the scoring math (cosine, positive-clip, weight, batch aggregation). Works for any embedding pair (image-image, text-text, or cross-modal), not just CLIP-produced ones. Registered on the lazy loader in `eval/metrics/__init__.py` following the same pattern as every existing metric. Tests: `tests/unit/eval/test_clip_score_metric.py` -- 20 tests covering construction/validation, cosine semantics (identical/orthogonal/anti- parallel/known values), weight scaling, zero-vector handling, shape handling (1-D/2-D/int arrays/size mismatch), aggregation (mean/std/min/ max/reset), and realistic 512-D CLIP-sized inputs. All pass. Full `tests/unit/eval/` regression: 600 pass (with the pre-existing network-dependent test skipped). No graph-optimization or model behaviour is changed. Prerequisite for text-to-image evaluation workflows (VLM alignment scoring, cross-modal retrieval, and generative-image eval). --- src/winml/modelkit/eval/metrics/__init__.py | 3 + src/winml/modelkit/eval/metrics/clip_score.py | 140 +++++++++++ tests/unit/eval/test_clip_score_metric.py | 221 ++++++++++++++++++ 3 files changed, 364 insertions(+) create mode 100644 src/winml/modelkit/eval/metrics/clip_score.py create mode 100644 tests/unit/eval/test_clip_score_metric.py diff --git a/src/winml/modelkit/eval/metrics/__init__.py b/src/winml/modelkit/eval/metrics/__init__.py index 0488db527..4c4a7a12c 100644 --- a/src/winml/modelkit/eval/metrics/__init__.py +++ b/src/winml/modelkit/eval/metrics/__init__.py @@ -14,6 +14,7 @@ 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 @@ -29,6 +30,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", @@ -61,6 +63,7 @@ def __dir__() -> list[str]: __all__ = [ "IGNORE_INDEX", "BinarySegmentationMetric", + "CLIPScoreMetric", "ClassificationMetric", "DepthMetric", "KNNAccuracyMetric", 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/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 From b97b07ba27b8593d07fbb72d4bb9a5a0eea39683 Mon Sep 17 00:00:00 2001 From: Abhilash Chenreddy Date: Wed, 5 Aug 2026 05:55:19 -0400 Subject: [PATCH 3/7] feat(eval): add RecallAtKMetric for retrieval and ranked-neighbor evaluation MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Streaming metric that reports Recall@K over ranked prediction lists, for any similarity-based retrieval or classification-as-retrieval workflow. For each query, given a ranked list of predicted IDs / labels (descending score) and one-or-more ground-truth relevant IDs, 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 (int ground_truth) -- 1.0 if the ID appears in the top K, else 0.0. Matches the classification-as-retrieval convention used across SSL embedding papers (DINO, DINOv2, MoCo, MAE). * Multi-relevant (array / set / tuple ground_truth) -- classical retrieval recall. Default k_values=(1, 5, 10) matches the reporting convention in the embedding-evaluation literature. Reports one recall_at_{k} entry per K, plus n_samples. Registered on the lazy loader in eval/metrics/__init__.py following the same pattern as every existing metric. Model-agnostic. Tests: tests/unit/eval/test_recall_at_k_metric.py -- 20 tests covering construction (defaults / custom / dedup / rejects empty and non-positive), single-relevant semantics (hit at top / mid / miss / np.integer scalar), multi-relevant semantics (full / partial / none / set input / empty gt rejected), aggregation (empty state / batch mean / reset), and shape handling (2-D flatten / empty ranked rejected / K larger than list). All pass. --- src/winml/modelkit/eval/metrics/__init__.py | 3 + .../modelkit/eval/metrics/recall_at_k.py | 132 +++++++++++++ tests/unit/eval/test_recall_at_k_metric.py | 182 ++++++++++++++++++ 3 files changed, 317 insertions(+) create mode 100644 src/winml/modelkit/eval/metrics/recall_at_k.py create mode 100644 tests/unit/eval/test_recall_at_k_metric.py diff --git a/src/winml/modelkit/eval/metrics/__init__.py b/src/winml/modelkit/eval/metrics/__init__.py index 4c4a7a12c..358d33b27 100644 --- a/src/winml/modelkit/eval/metrics/__init__.py +++ b/src/winml/modelkit/eval/metrics/__init__.py @@ -21,6 +21,7 @@ from .mean_average_precision import MAPMetric from .mean_iou import IGNORE_INDEX, MeanIoUMetric from .pseudo_perplexity import PseudoPerplexityMetric + from .recall_at_k import RecallAtKMetric from .spearman_correlation import SpearmanCorrelationMetric from .top_k_accuracy import TopKAccuracyMetric @@ -39,6 +40,7 @@ "MAPMetric": ".mean_average_precision:MAPMetric", "MeanIoUMetric": ".mean_iou:MeanIoUMetric", "PseudoPerplexityMetric": ".pseudo_perplexity:PseudoPerplexityMetric", + "RecallAtKMetric": ".recall_at_k:RecallAtKMetric", "SpearmanCorrelationMetric": ".spearman_correlation:SpearmanCorrelationMetric", "TopKAccuracyMetric": ".top_k_accuracy:TopKAccuracyMetric", } @@ -71,6 +73,7 @@ def __dir__() -> list[str]: "MAPMetric", "MeanIoUMetric", "PseudoPerplexityMetric", + "RecallAtKMetric", "SpearmanCorrelationMetric", "TopKAccuracyMetric", ] 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/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 From 6fba49ffded9d2a29284dedf9ba9b0824d754b8b Mon Sep 17 00:00:00 2001 From: Abhilash Chenreddy Date: Wed, 5 Aug 2026 05:57:50 -0400 Subject: [PATCH 4/7] feat(eval): add MeanReciprocalRankMetric for retrieval evaluation MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Streaming metric that reports the Mean Reciprocal Rank (MRR) over ranked prediction lists. For each query, MRR uses the rank of the FIRST relevant item in the ranked list:: MRR = (1/N) * Σ 1 / rank_first_relevant(query) Queries whose ranked list contains no relevant item contribute 0 (rank treated as infinity). Companion to RecallAtKMetric. Recall@K reports 'is any relevant item in the top K', MRR reports 'how highly is the first relevant item ranked'. Retrieval evaluations typically report both together. Same input shape as RecallAtKMetric so a single (ranked_predictions, ground_truth) stream can drive both metrics from the same evaluator. Registered on the lazy loader. Model-agnostic. Tests: tests/unit/eval/test_mean_reciprocal_rank_metric.py -- 17 tests covering single-relevant (rank 1 / 2 / 3 / miss / np.integer scalar), multi-relevant (first hit wins / earliest position / no hit / set input / empty gt rejected), aggregation (empty state / batch mean / reset), and shape handling (2-D flatten / empty ranked rejected). All pass. --- src/winml/modelkit/eval/metrics/__init__.py | 3 + .../eval/metrics/mean_reciprocal_rank.py | 109 +++++++++++++++ .../eval/test_mean_reciprocal_rank_metric.py | 129 ++++++++++++++++++ 3 files changed, 241 insertions(+) create mode 100644 src/winml/modelkit/eval/metrics/mean_reciprocal_rank.py create mode 100644 tests/unit/eval/test_mean_reciprocal_rank_metric.py diff --git a/src/winml/modelkit/eval/metrics/__init__.py b/src/winml/modelkit/eval/metrics/__init__.py index 358d33b27..872661b70 100644 --- a/src/winml/modelkit/eval/metrics/__init__.py +++ b/src/winml/modelkit/eval/metrics/__init__.py @@ -20,6 +20,7 @@ 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 @@ -39,6 +40,7 @@ "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", @@ -72,6 +74,7 @@ def __dir__() -> list[str]: "KeypointAPMetric", "MAPMetric", "MeanIoUMetric", + "MeanReciprocalRankMetric", "PseudoPerplexityMetric", "RecallAtKMetric", "SpearmanCorrelationMetric", 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/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) From 253a2a79fa25f14b540e55dbadbef991c8f408ec Mon Sep 17 00:00:00 2001 From: Abhilash Chenreddy Date: Wed, 5 Aug 2026 06:03:51 -0400 Subject: [PATCH 5/7] feat(eval): report Recall@K and MRR from image-feature-extraction evaluator Enriches WinMLImageFeatureExtractionEvaluator to report standard retrieval metrics alongside the existing kNN classification accuracy: { 'knn_top1_accuracy': ..., # existing, unchanged 'knn_top5_accuracy': ..., # existing, unchanged 'recall_at_1': ..., # NEW 'recall_at_5': ..., # NEW 'recall_at_10': ..., # NEW 'mrr': ..., # NEW } Reasoning: DINO / DINOv2 / MoCo / MAE and the broader SSL embedding literature report Recall@K and MRR as the standard measures of embedding quality. kNN classification accuracy is a proxy; ranking- based metrics are the direct measure of 'how well do these embeddings separate classes'. Same L2 normalisation, self-exclusion and cosine ranking as KNNAccuracyMetric are reused via a new private helper -- _compute_retrieval_metrics -- so the retrieval numbers describe the same neighbour ordering the kNN classifier voted on. A full descending argsort is added (KNN uses argpartition for top-K only) because MRR needs the position of the first same-class neighbour at arbitrary rank. Model-agnostic per the metric primitives (RecallAtKMetric, MeanReciprocalRankMetric); each treats the query's own class label as the single relevant match, matching the SSL literature convention. Backward-compatible: existing knn_top{1,5}_accuracy keys are still reported. Existing evaluator tests continue to pass unchanged. Tests: extended tests/unit/eval/test_image_feature_extraction_evaluator.py with 4 new tests -- * retrieval_metrics_reported_alongside_knn (output shape) * perfect_clusters_score_max_retrieval (correctness at recall/mrr = 1) * compute_retrieval_metrics_static_helper (helper method works without pipeline mocks) * no_same_class_neighbours_scores_zero (singleton classes -> 0) Full tests/unit/eval/ regression: 642 pass, 1 pre-existing network test skipped. --- .../image_feature_extraction_evaluator.py | 75 ++++++++- ...test_image_feature_extraction_evaluator.py | 159 ++++++++++++++---- 2 files changed, 201 insertions(+), 33 deletions(-) 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/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() - From f191b86b03b02faa0c208382eaa822a22c1bd2e0 Mon Sep 17 00:00:00 2001 From: Abhilash Chenreddy Date: Wed, 5 Aug 2026 06:15:06 -0400 Subject: [PATCH 6/7] feat(datasets): add filter/sample to PromptDataset with _derive extension hook Adds the two standard derived-dataset operations that every mature dataset class exposes: ds.filter(lambda r: len(r['prompt']) < 100) ds.sample(50, seed=42) Both return a NEW PromptDataset (never mutate self), are chainable, and preserve model_name/dataset_name/data_split from the source unless the caller passes an explicit override. Extensibility -- '_derive' is the single hook -------------------------------------------- The two methods share a common shape: transform the record list, build a new dataset instance from the derived records. A private '_derive' method is the extension point: def _derive(self, records, **kwargs) -> PromptDataset: 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) Two properties matter for downstream contributors: * Uses type(self), not PromptDataset, as the constructor -- subclasses automatically receive an instance of themselves. * Subclasses can override _derive to change the derived type, inject extra metadata (dataset naming, source tracking, tags, ...), or add cascade defaults, and the change automatically applies to filter, sample, and any future derived-dataset method. Adding a new derived operation (e.g. 'deduplicate', 'group_by') is now a matter of one method that transforms self._dataset and calls self._derive(new_records) -- no need to touch _derive itself. filter's predicate receives a shallow copy of each record (matches __getitem__ semantics) to prevent accidental mutation of the source dataset from within the predicate. Tests: tests/unit/datasets/test_prompt_dataset.py extended with 18 new tests -- * TestFilter (8): keeps matching / new instance / correct type / metadata inherited / metadata override / predicate copy semantics / rejects all raises / chainable with sample. * TestSample (8): size matches / reproducible seed / different seeds differ / full-size permutation / new instance / oversample rejected / non-positive n rejected / metadata inherited. * TestDeriveExtensionHook (2): subclass receives own type / subclass can override _derive for extra metadata. All 53 PromptDataset tests pass. Full tests/unit/datasets/ regression: 131 pass. --- src/winml/modelkit/datasets/prompt_dataset.py | 106 +++++++++++- tests/unit/datasets/test_prompt_dataset.py | 152 ++++++++++++++++++ 2 files changed, 257 insertions(+), 1 deletion(-) diff --git a/src/winml/modelkit/datasets/prompt_dataset.py b/src/winml/modelkit/datasets/prompt_dataset.py index 45512b6ea..94f31966d 100644 --- a/src/winml/modelkit/datasets/prompt_dataset.py +++ b/src/winml/modelkit/datasets/prompt_dataset.py @@ -51,13 +51,14 @@ 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 Iterable + from collections.abc import Callable, Iterable logger = logging.getLogger(__name__) @@ -360,3 +361,106 @@ def label_col(self) -> str: def __iter__(self): """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/tests/unit/datasets/test_prompt_dataset.py b/tests/unit/datasets/test_prompt_dataset.py index b4ede6bd0..0a49ef527 100644 --- a/tests/unit/datasets/test_prompt_dataset.py +++ b/tests/unit/datasets/test_prompt_dataset.py @@ -296,3 +296,155 @@ 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" From daf8a477a26a692c9e43100893ae607e152c5e4d Mon Sep 17 00:00:00 2001 From: Abhilash Chenreddy Date: Wed, 5 Aug 2026 14:35:38 -0400 Subject: [PATCH 7/7] fix(datasets): add return type to PromptDataset.__iter__ for mypy --- src/winml/modelkit/datasets/prompt_dataset.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/winml/modelkit/datasets/prompt_dataset.py b/src/winml/modelkit/datasets/prompt_dataset.py index 94f31966d..91a68dfaf 100644 --- a/src/winml/modelkit/datasets/prompt_dataset.py +++ b/src/winml/modelkit/datasets/prompt_dataset.py @@ -58,7 +58,7 @@ if TYPE_CHECKING: - from collections.abc import Callable, Iterable + from collections.abc import Callable, Iterable, Iterator logger = logging.getLogger(__name__) @@ -358,7 +358,7 @@ def label_col(self) -> str: # Convenience # --------------------------------------------------------------------- - def __iter__(self): + def __iter__(self) -> Iterator[dict[str, Any]]: """Iterate over the coerced record dicts.""" return iter(self._dataset)