Feat/text to image eval primitives: add prompt-driven evaluation primitives (PromptDataset + CLIPScore + Recall@K + MRR) - #1284
Conversation
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).
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).
…luation
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.
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.
…luator
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.
…sion 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.
| 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) |
There was a problem hiding this comment.
This counts every occurrence of a relevant label in top_k, so Recall@K can exceed 1.0 when the ranked list contains duplicates. That happens with the new image-feature evaluator because it passes ranked_labels; for example, five same-class neighbors in the top 5 would report recall_at_5 = 5.0 for a scalar ground truth. Please count unique retrieved relevant IDs (or treat scalar ground truth as a boolean hit) and add a duplicate-label regression test.
| 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 |
There was a problem hiding this comment.
Since these are being exposed from the package root, please also add PromptDataset and PromptRecord to __all__ below. Otherwise from winml.modelkit.datasets import * and code/docs that rely on the declared public API won't include the new prompt dataset primitives, unlike the other dataset classes exported here.
Qiong Wu (qiowu) (DingmaomaoBJTU)
left a comment
Qiong Wu (qiowu) (DingmaomaoBJTU)
left a comment
There was a problem hiding this comment.
Looking at this as the foundation for a future Stable Diffusion/text-to-image evaluator: can you show how these primitives are intended to be exercised through winml eval today? For example, is there a command/config that loads a prompt dataset and reports CLIPScore / Recall@K / MRR, or are these library-only primitives until the follow-up SD evaluator PR wires them into the eval registry?
| f"PromptDataset record must be a dict or PromptRecord, got {type(raw).__name__}" | ||
| ) | ||
|
|
||
| unknown = set(raw) - _ALLOWED_KEYS |
There was a problem hiding this comment.
This strict unknown-key check may make real SD prompt corpora harder to use than necessary. Prompt sets often carry extra fields such as category, style, seed, source, or difficulty; requiring callers to pre-pack all of those into metadata means common JSONL/HF prompt datasets will fail before the evaluator can use them. Could we either collect unknown top-level keys into metadata by default, or expose a non-strict mode for evaluator-facing prompt datasets?
| return cls(records, **kwargs) | ||
|
|
||
| @classmethod | ||
| def from_hf( |
There was a problem hiding this comment.
For the future SD evaluator, I think we should be careful not to create a parallel dataset-loading path that bypasses WinMLEvaluator.prepare_data() / DatasetConfig. The existing eval flow already owns path,
ame, split, samples, shuffle, streaming,
evision, default datasets, and --schema validation. Could the eventual text-to-image evaluator load via the existing DatasetConfig path and normalize rows into prompt records there, so this stays testable through winml eval instead of only through direct PromptDataset.from_hf/from_jsonl calls?
Qiong Wu (qiowu) (DingmaomaoBJTU)
left a comment
Qiong Wu (qiowu) (DingmaomaoBJTU)
left a comment
There was a problem hiding this comment.
One broader design question: why do we need PromptDataset as a first-class dataset abstraction here, instead of keeping prompt handling as a text-to-image evaluator normalization step over the existing DatasetConfig / HF dataset flow? I can see value in a canonical PromptRecord shape, but I want to understand the intended boundary before we build more evaluator code on top of it.
| # --------------------------------------------------------------------------- | ||
|
|
||
|
|
||
| class PromptDataset(BaseTaskDataset): |
There was a problem hiding this comment.
Could we narrow this abstraction so PromptDataset is primarily a row normalizer / typed view rather than the owner of dataset loading? The eval stack already has DatasetConfig + task schema + WinMLEvaluator.prepare_data() to own path,
ame, split, samples, shuffle, streaming,
evision, defaults, and CLI schema validation. For a future text-to-image evaluator, I think the main path should load through that existing lifecycle, then normalize rows into prompt records via something like PromptDataset.from_rows() / rom_dataset(). rom_jsonl() / rom_hf() can still exist as library conveniences, but they shouldn't become the winml eval integration path unless they deliberately mirror all of DatasetConfig semantics.
| # --------------------------------------------------------------------------- | ||
|
|
||
|
|
||
| class PromptDataset(BaseTaskDataset): |
There was a problem hiding this comment.
Stepping back from the future SD evaluator use case: do we need a first-class PromptDataset abstraction yet? A text-to-image evaluator could already load a regular HF datasets.Dataset through the existing DatasetConfig / WinMLEvaluator.prepare_data() path, then read prompt,
egative_prompt, reference columns, and metadata via columns_mapping. That keeps split/sampling/shuffle/streaming/schema behavior aligned with the rest of winml eval. Unless we expect multiple prompt-driven evaluators to share this immediately, a smaller PromptRecord/row-normalization helper may be enough for now and would avoid introducing a parallel dataset abstraction.
Zhenchao Ni (zhenchaoni)
left a comment
Zhenchao Ni (zhenchaoni)
left a comment
There was a problem hiding this comment.
Reviewing the new prompt-driven evaluation primitives for correctness and alignment with the existing evaluator architecture.
| 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) |
There was a problem hiding this comment.
ranked_predictions may contain duplicate labels, as it does in the image feature evaluator. Counting every occurrence means [0, 0, 0] with ground truth 0 produces Recall@K = 3.0. Please count the unique overlap, or use any(...) for scalar ground truth, and add a duplicate-label test.
| """ | ||
| # Coerce and validate before the base class initialises so that | ||
| # _initialize can just assign the list. | ||
| self._raw_records = list(records) |
There was a problem hiding this comment.
max_samples is applied only after list(records) and full validation, so a large or streaming dataset is still consumed entirely. Please limit the iterable before materializing it, and use the appropriate select/take path for Hugging Face datasets.
| norms = np.maximum(norms, 1e-9) | ||
| normalized = embeddings / norms | ||
|
|
||
| similarity = normalized @ normalized.T |
There was a problem hiding this comment.
KNNAccuracyMetric.compute() already normalizes the embeddings and builds an N×N similarity matrix. Repeating that work here, followed by a full sort, significantly increases cost for large datasets. Could the shared ranking/similarity computation live in the metrics layer and be reused by kNN, Recall@K, and MRR?
| 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 |
There was a problem hiding this comment.
Could we wire these primitives into the evaluator flow in this PR? PromptDataset and CLIPScoreMetric are currently only exported; no evaluator or task registration selects them or runs inference through the model abstraction, so winml eval cannot use this new functionality. If this PR is intentionally primitives-only, that scope should be made explicit and the evaluator integration tracked separately.
| @@ -35,7 +38,7 @@ class BaseTaskDataset(ABC): | |||
|
|
|||
| def __init__( | |||
| self, | |||
There was a problem hiding this comment.
Should PromptDataset inherit from BaseTaskDataset if doing so requires making model_name optional for every task dataset? This weakens the base contract and shifts validation to all existing subclasses. A model-agnostic dataset abstraction, or a separate prompt dataset base, may preserve the current task-dataset invariant more safely.
3 participants
This PR has the prerequisites needed to support stable diffusion models. In order to do it cleanly, SD work will be done in 2 PRs where we keep this PR general enough while adding the necessary components for SD work.
What's in here:
While I was working on the evaluation components, I noticed WinMLImageFeatureExtractionEvaluator already computes the full cosine similarity matrix and ranks neighbours for the k-NN classifier, but only reports classification accuracy. The SSL / embedding-quality literature (DINO, DINOv2, MoCo, MAE) reports Recall@K and MRR -those are the direct measures; k-NN accuracy is a proxy. So this PR also has it report recall_at_{1,5,10} and mrr alongside knn_top{1,5}_accuracy. Same similarity matrix, no meaningful extra compute, backward-compatible.
One non-trivial base-class change: widened BaseTaskDataset.model_name from str to str | None. PromptDataset is task-agnostic so it doesn't have a model to bind to. All 7 existing subclasses still receive a str from their callers, I ran their tests to confirm nothing behaves differently.
Things I deliberately left out:
No shared cosine kernel across metrics. Each metric implements its own math (matches what every other metric here does).
Testing:
115 new tests. 773 pass in tests/unit/{datasets,eval}/ (131 datasets + 642 eval). Broader unit sweep (~2,440 tests across build/quant/commands/inference/models/loader) has no new failures relative to upstream/main. CI-parity (tests/regression/ + tests/cli/ with the CI -m filter) is bit-identical to upstream/main. Ruff check + format clean on all 13 touched files.