diff --git a/examples/recipes/cross-encoder_ms-marco-MiniLM-L6-v2/cpu/cpu/reranking_fp16_config.json b/examples/recipes/cross-encoder_ms-marco-MiniLM-L6-v2/cpu/cpu/reranking_fp16_config.json new file mode 100644 index 000000000..4780f7868 --- /dev/null +++ b/examples/recipes/cross-encoder_ms-marco-MiniLM-L6-v2/cpu/cpu/reranking_fp16_config.json @@ -0,0 +1,65 @@ +{ + "export": { + "opset_version": 17, + "batch_size": 1, + "export_params": true, + "do_constant_folding": true, + "verbose": false, + "dynamo": false, + "enable_hierarchy_tags": true, + "clean_onnx": false, + "hierarchy_tag_format": "full", + "input_tensors": [ + { + "name": "input_ids", + "dtype": "int32", + "shape": [1, 512], + "value_range": [0, 30522] + }, + { + "name": "attention_mask", + "dtype": "int32", + "shape": [1, 512], + "value_range": [0, 2] + }, + { + "name": "token_type_ids", + "dtype": "int32", + "shape": [1, 512], + "value_range": [0, 2] + } + ], + "output_tensors": [ + { + "name": "logits" + } + ] + }, + "optim": {}, + "quant": { + "mode": "fp16", + "samples": 10, + "calibration_method": "minmax", + "weight_type": "uint8", + "activation_type": "uint8", + "per_channel": false, + "symmetric": false, + "weight_symmetric": null, + "activation_symmetric": null, + "save_calibration": false, + "distribution": "uniform", + "seed": null, + "calibration_load_path": null, + "calibration_save_path": null, + "op_types_to_quantize": null, + "nodes_to_exclude": null, + "fp16_keep_io_types": true, + "fp16_op_block_list": null + }, + "compile": null, + "loader": { + "task": "reranking", + "model_class": "AutoModelForSequenceClassification", + "model_type": "bert" + } +} \ No newline at end of file diff --git a/examples/recipes/cross-encoder_ms-marco-MiniLM-L6-v2/cpu/cpu/reranking_fp32_config.json b/examples/recipes/cross-encoder_ms-marco-MiniLM-L6-v2/cpu/cpu/reranking_fp32_config.json new file mode 100644 index 000000000..e6548faa3 --- /dev/null +++ b/examples/recipes/cross-encoder_ms-marco-MiniLM-L6-v2/cpu/cpu/reranking_fp32_config.json @@ -0,0 +1,46 @@ +{ + "export": { + "opset_version": 17, + "batch_size": 1, + "export_params": true, + "do_constant_folding": true, + "verbose": false, + "dynamo": false, + "enable_hierarchy_tags": true, + "clean_onnx": false, + "hierarchy_tag_format": "full", + "input_tensors": [ + { + "name": "input_ids", + "dtype": "int32", + "shape": [1, 512], + "value_range": [0, 30522] + }, + { + "name": "attention_mask", + "dtype": "int32", + "shape": [1, 512], + "value_range": [0, 2] + }, + { + "name": "token_type_ids", + "dtype": "int32", + "shape": [1, 512], + "value_range": [0, 2] + } + ], + "output_tensors": [ + { + "name": "logits" + } + ] + }, + "optim": {}, + "quant": null, + "compile": null, + "loader": { + "task": "reranking", + "model_class": "AutoModelForSequenceClassification", + "model_type": "bert" + } +} \ No newline at end of file diff --git a/scripts/e2e_eval/datasets/build_msmarco_reranking_fixture.py b/scripts/e2e_eval/datasets/build_msmarco_reranking_fixture.py new file mode 100644 index 000000000..54464bfab --- /dev/null +++ b/scripts/e2e_eval/datasets/build_msmarco_reranking_fixture.py @@ -0,0 +1,365 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- + +"""Build a tiny local grouped reranking dataset from authoritative MS MARCO data. + +The fixture is intentionally small: it selects 1-2 real dev queries from the +pinned Hugging Face dataset revision and joins them against the official public +MS MARCO passage-ranking files so ``winml eval --task reranking`` can run fully +offline on a local ``DatasetDict``. + +Saved format: + output/ + dataset_dict.json + Arrow shards via ``DatasetDict.save_to_disk`` + provenance.json + +Each saved row contains: + - ``input``: real query text + - ``expected_output``: list of positive passage IDs present in candidates + - ``metadata``: query/group provenance including the pinned HF row index + - ``candidates``: ordered list of candidate dicts with real passage text + +Usage: + python scripts/e2e_eval/datasets/build_msmarco_reranking_fixture.py --output + python scripts/e2e_eval/datasets/build_msmarco_reranking_fixture.py --output --queries 2 --max-negatives 3 +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import shutil +import tarfile +import urllib.request +from dataclasses import dataclass +from pathlib import Path +from typing import TYPE_CHECKING, Any + + +if TYPE_CHECKING: + from collections.abc import Iterable + + +HF_DATASET_ID = "orgrctera/msmarco_passage_ranking" +HF_REVISION = "a7388b9efd4dd4b87a0db91314e5b3f0e4b0d9e6" +HF_PARQUET_RELATIVE_PATH = "data/dev-00000-of-00001.parquet" +HF_PARQUET_URL = ( + "https://huggingface.co/datasets/" + f"{HF_DATASET_ID}/resolve/{HF_REVISION}/{HF_PARQUET_RELATIVE_PATH}" +) + +OFFICIAL_QRELS_URL = "https://msmarco.z22.web.core.windows.net/msmarcoranking/qrels.dev.tsv" +OFFICIAL_TOP1000_URL = ( + "https://msmarco.z22.web.core.windows.net/msmarcoranking/top1000.dev.tar.gz" +) +OFFICIAL_QUERIES_URL = "https://msmarco.z22.web.core.windows.net/msmarcoranking/queries.tar.gz" + +DEFAULT_CACHE = Path.home() / ".cache" / "winml" / "msmarco_reranking_fixture" + + +@dataclass(frozen=True) +class CandidateRow: + """One candidate passage row from the official top1000 reranking file.""" + + pid: str + query: str + passage: str + rank: int + + +def _sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as handle: + for chunk in iter(lambda: handle.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def _download(url: str, dest: Path) -> Path: + if dest.exists(): + return dest + dest.parent.mkdir(parents=True, exist_ok=True) + with urllib.request.urlopen(url) as response, dest.open("wb") as handle: # noqa: S310 + shutil.copyfileobj(response, handle) + return dest + + +def _parse_json_list(value: Any) -> list[Any]: + if isinstance(value, list): + return value + if isinstance(value, str): + parsed = json.loads(value) + if isinstance(parsed, list): + return parsed + raise ValueError(f"expected JSON list, got {type(value).__name__}") + + +def _parse_json_object(value: Any) -> dict[str, Any]: + if isinstance(value, dict): + return value + if isinstance(value, str): + parsed = json.loads(value) + if isinstance(parsed, dict): + return parsed + raise ValueError(f"expected JSON object, got {type(value).__name__}") + + +def _iter_tar_lines(archive_path: Path, preferred_members: Iterable[str]) -> Iterable[str]: + with tarfile.open(archive_path, "r:gz") as archive: + members = [member for member in archive.getmembers() if member.isfile()] + member = None + preferred = tuple(preferred_members) + for candidate in members: + lower_name = candidate.name.lower() + if any(token in lower_name for token in preferred): + member = candidate + break + if member is None: + if not members: + raise RuntimeError(f"archive {archive_path} has no file members") + member = members[0] + extracted = archive.extractfile(member) + if extracted is None: + raise RuntimeError(f"failed to extract {member.name} from {archive_path}") + for raw_line in extracted: + yield raw_line.decode("utf-8").rstrip("\n") + + +def _load_hf_rows(parquet_path: Path) -> list[dict[str, Any]]: + from datasets import load_dataset + + dataset = load_dataset("parquet", data_files=str(parquet_path), split="train") + return [dict(row) for row in dataset] + + +def _load_queries(archive_path: Path) -> dict[str, str]: + queries: dict[str, str] = {} + for line in _iter_tar_lines(archive_path, preferred_members=("queries.dev", "queries")): + if not line.strip(): + continue + qid, query = line.split("\t", 1) + queries[qid] = query + return queries + + +def _load_qrels(qrels_path: Path) -> dict[str, set[str]]: + qrels: dict[str, set[str]] = {} + for raw_line in qrels_path.read_text(encoding="utf-8").splitlines(): + if not raw_line.strip(): + continue + qid, _unused, pid, rel = raw_line.split() + if int(rel) > 0: + qrels.setdefault(qid, set()).add(pid) + return qrels + + +def _load_top1000(archive_path: Path) -> dict[str, list[CandidateRow]]: + grouped: dict[str, list[CandidateRow]] = {} + for line in _iter_tar_lines(archive_path, preferred_members=("top1000.dev", "top1000")): + if not line.strip(): + continue + qid, pid, query, passage = line.split("\t", 3) + rows = grouped.setdefault(qid, []) + rows.append(CandidateRow(pid=pid, query=query, passage=passage, rank=len(rows) + 1)) + return grouped + + +def _select_rows( + hf_rows: list[dict[str, Any]], + queries: dict[str, str], + qrels: dict[str, set[str]], + top1000: dict[str, list[CandidateRow]], + *, + max_queries: int, + max_negatives: int, +) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]: + selected_rows: list[dict[str, Any]] = [] + selection_provenance: list[dict[str, Any]] = [] + + for row_index, row in enumerate(hf_rows): + metadata = _parse_json_object(row["metadata"]) + qid = str(metadata.get("query_id", "")).strip() + if not qid or qid not in queries or qid not in qrels or qid not in top1000: + continue + + query_text = str(row["input"]) + official_query = queries[qid] + if official_query != query_text: + continue + + hf_positive_ids = {str(value) for value in _parse_json_list(row["expected_output"])} + official_positive_ids = {str(value) for value in qrels[qid]} + + selected_candidates: list[CandidateRow] = [] + selected_positive_ids: list[str] = [] + selected_negative_ids: list[str] = [] + for candidate in top1000[qid]: + if candidate.pid in hf_positive_ids and candidate.pid in official_positive_ids: + selected_candidates.append(candidate) + selected_positive_ids.append(candidate.pid) + elif candidate.pid not in official_positive_ids and len(selected_negative_ids) < max_negatives: + selected_candidates.append(candidate) + selected_negative_ids.append(candidate.pid) + if selected_positive_ids and len(selected_negative_ids) >= max_negatives: + break + + if not selected_positive_ids or not selected_negative_ids: + continue + + candidates = [ + { + "id": candidate.pid, + "text": candidate.passage, + "rank": candidate.rank, + "relevant": candidate.pid in selected_positive_ids, + } + for candidate in selected_candidates + ] + metadata_out = { + **metadata, + "query_id": qid, + "source_row_index": row_index, + "positive_candidate_ids": selected_positive_ids, + "negative_candidate_ids": selected_negative_ids, + "selected_candidate_ids": [candidate.pid for candidate in selected_candidates], + "selection_strategy": "authoritative_top1000_order_with_bounded_negatives", + "candidate_source": "official_top1000.dev", + } + selected_rows.append( + { + "input": query_text, + "expected_output": selected_positive_ids, + "metadata": metadata_out, + "candidates": candidates, + } + ) + selection_provenance.append( + { + "query_id": qid, + "source_row_index": row_index, + "hf_expected_output_ids": sorted(hf_positive_ids), + "official_qrels_positive_ids": sorted(official_positive_ids), + "selected_positive_candidate_ids": selected_positive_ids, + "selected_negative_candidate_ids": selected_negative_ids, + "selected_candidate_ids": [candidate.pid for candidate in selected_candidates], + "selection_strategy": "authoritative_top1000_order_with_bounded_negatives", + "candidate_ranks": {candidate.pid: candidate.rank for candidate in selected_candidates}, + } + ) + if len(selected_rows) >= max_queries: + break + + if len(selected_rows) < max_queries: + raise RuntimeError( + f"Could only materialize {len(selected_rows)} grouped queries; required {max_queries}." + ) + return selected_rows, selection_provenance + + +def build_dataset(output_dir: Path, cache_dir: Path, max_queries: int, max_negatives: int) -> Path: + from datasets import Dataset, DatasetDict + + parquet_path = _download(HF_PARQUET_URL, cache_dir / "hf" / Path(HF_PARQUET_RELATIVE_PATH).name) + qrels_path = _download(OFFICIAL_QRELS_URL, cache_dir / "official" / "qrels.dev.tsv") + top1000_path = _download(OFFICIAL_TOP1000_URL, cache_dir / "official" / "top1000.dev.tar.gz") + queries_path = _download(OFFICIAL_QUERIES_URL, cache_dir / "official" / "queries.tar.gz") + + hf_rows = _load_hf_rows(parquet_path) + queries = _load_queries(queries_path) + qrels = _load_qrels(qrels_path) + top1000 = _load_top1000(top1000_path) + + selected_rows, selection_provenance = _select_rows( + hf_rows, + queries, + qrels, + top1000, + max_queries=max_queries, + max_negatives=max_negatives, + ) + + dataset_dict = DatasetDict({"dev": Dataset.from_list(selected_rows)}) + output_dir.mkdir(parents=True, exist_ok=True) + dataset_dict.save_to_disk(str(output_dir)) + + provenance = { + "schema": "winml/msmarco-reranking-fixture/2", + "builder": { + "script": "scripts/e2e_eval/datasets/build_msmarco_reranking_fixture.py", + "hf_dataset_id": HF_DATASET_ID, + "hf_revision": HF_REVISION, + "max_queries": max_queries, + "max_negatives": max_negatives, + }, + "sources": { + "hf_row_source": { + "url": HF_PARQUET_URL, + "relative_path": HF_PARQUET_RELATIVE_PATH, + "sha256": _sha256(parquet_path), + "local_path": str(parquet_path), + }, + "official_queries": { + "url": OFFICIAL_QUERIES_URL, + "sha256": _sha256(queries_path), + "local_path": str(queries_path), + }, + "official_qrels": { + "url": OFFICIAL_QRELS_URL, + "sha256": _sha256(qrels_path), + "local_path": str(qrels_path), + }, + "official_top1000": { + "url": OFFICIAL_TOP1000_URL, + "sha256": _sha256(top1000_path), + "local_path": str(top1000_path), + }, + }, + "selected_rows": selection_provenance, + "output": { + "dataset_path": str(output_dir), + "split": "dev", + "row_count": len(selected_rows), + }, + } + provenance_path = output_dir / "provenance.json" + provenance_path.write_text(json.dumps(provenance, indent=2), encoding="utf-8") + return provenance_path + + +def main() -> int: + parser = argparse.ArgumentParser(description="Build a tiny grouped MS MARCO reranking fixture.") + parser.add_argument("--output", type=Path, required=True, help="Output dataset directory.") + parser.add_argument( + "--cache-dir", + type=Path, + default=DEFAULT_CACHE, + help="Directory for downloaded source artifacts.", + ) + parser.add_argument( + "--queries", + type=int, + default=2, + help="Number of grouped dev queries to materialize.", + ) + parser.add_argument( + "--max-negatives", + type=int, + default=3, + help="Maximum real negative candidates to retain per query.", + ) + args = parser.parse_args() + provenance_path = build_dataset( + output_dir=args.output, + cache_dir=args.cache_dir, + max_queries=args.queries, + max_negatives=args.max_negatives, + ) + print(provenance_path) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/src/winml/modelkit/datasets/__init__.py b/src/winml/modelkit/datasets/__init__.py index 66bee45de..275ca1c2e 100644 --- a/src/winml/modelkit/datasets/__init__.py +++ b/src/winml/modelkit/datasets/__init__.py @@ -41,6 +41,7 @@ "image-classification": ImageDataset, "image-feature-extraction": ImageDataset, "object-detection": ObjectDetectionDataset, + "reranking": TextDataset, "text-classification": TextDataset, "text-feature-extraction": TextDataset, "feature-extraction": TextDataset, diff --git a/src/winml/modelkit/eval/base_evaluator.py b/src/winml/modelkit/eval/base_evaluator.py index 49665d735..7b5a1fc50 100644 --- a/src/winml/modelkit/eval/base_evaluator.py +++ b/src/winml/modelkit/eval/base_evaluator.py @@ -103,7 +103,17 @@ def prepare_data(self) -> Dataset: try: ds_path = Path(ds.path).expanduser() if ds.path else None if ds_path and ds_path.is_dir(): - dataset = load_from_disk(str(ds_path)) + loaded = load_from_disk(str(ds_path)) + if isinstance(loaded, Dataset): + dataset = loaded + else: + available_splits = sorted(str(name) for name in loaded) + if ds.split not in loaded: + raise DatasetValidationError( + f"Local dataset '{ds.path}' has splits {available_splits}, " + f"but split '{ds.split}' was requested" + ) + dataset = loaded[ds.split] else: dataset = load_dataset( ds.path, diff --git a/src/winml/modelkit/eval/evaluate.py b/src/winml/modelkit/eval/evaluate.py index 07eefd6ee..541f36a3b 100644 --- a/src/winml/modelkit/eval/evaluate.py +++ b/src/winml/modelkit/eval/evaluate.py @@ -65,6 +65,8 @@ def _select_model_loader(config: WinMLEvaluationConfig) -> _ModelLoaderKind: _EVALUATOR_REGISTRY: dict[str, str] = { "image-classification": "winml.modelkit.eval.base_evaluator:WinMLEvaluator", + "reranking": + "winml.modelkit.eval.reranking_evaluator:WinMLRerankingEvaluator", "text-classification": "winml.modelkit.eval.text_classification_evaluator:WinMLTextClassificationEvaluator", "sequence-classification": @@ -191,6 +193,19 @@ def _validate_pytorch_runtime_config(config: WinMLEvaluationConfig) -> None: "second_input_column": "sentence2", }, }, + "reranking": { + "path": "mteb/scidocs-reranking", + "split": "test", + "revision": "56a6d0140cf6356659e2a7c1413286a774468d44", + "streaming": True, + "shuffle": False, + "columns_mapping": { + "query_column": "query", + "positive_column": "positive", + "negative_column": "negative", + "max_candidates": "10", + }, + }, "token-classification": { "path": "BramVanroy/conll2003", "split": "validation", diff --git a/src/winml/modelkit/eval/metrics/__init__.py b/src/winml/modelkit/eval/metrics/__init__.py index 0488db527..db6e9bdb2 100644 --- a/src/winml/modelkit/eval/metrics/__init__.py +++ b/src/winml/modelkit/eval/metrics/__init__.py @@ -20,6 +20,7 @@ from .mean_average_precision import MAPMetric from .mean_iou import IGNORE_INDEX, MeanIoUMetric from .pseudo_perplexity import PseudoPerplexityMetric + from .ranking import RerankingMetric from .spearman_correlation import SpearmanCorrelationMetric from .top_k_accuracy import TopKAccuracyMetric @@ -37,6 +38,7 @@ "MAPMetric": ".mean_average_precision:MAPMetric", "MeanIoUMetric": ".mean_iou:MeanIoUMetric", "PseudoPerplexityMetric": ".pseudo_perplexity:PseudoPerplexityMetric", + "RerankingMetric": ".ranking:RerankingMetric", "SpearmanCorrelationMetric": ".spearman_correlation:SpearmanCorrelationMetric", "TopKAccuracyMetric": ".top_k_accuracy:TopKAccuracyMetric", } @@ -68,6 +70,7 @@ def __dir__() -> list[str]: "MAPMetric", "MeanIoUMetric", "PseudoPerplexityMetric", + "RerankingMetric", "SpearmanCorrelationMetric", "TopKAccuracyMetric", ] diff --git a/src/winml/modelkit/eval/metrics/ranking.py b/src/winml/modelkit/eval/metrics/ranking.py new file mode 100644 index 000000000..e3ea06030 --- /dev/null +++ b/src/winml/modelkit/eval/metrics/ranking.py @@ -0,0 +1,66 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- + +"""Ranking metrics for reranking evaluators.""" + +from __future__ import annotations + +from dataclasses import dataclass + + +@dataclass +class RerankingMetric: + """Aggregate MRR@K and Recall@K over grouped candidates. + + Ties are broken stably by the candidates' original order. Groups with no + positive labels are counted separately and excluded from the metric + denominator so malformed supervision never scores as a silent miss. + """ + + recall_ks: tuple[int, ...] = (1, 10) + mrr_k: int = 10 + + def __post_init__(self) -> None: + ordered = sorted({int(k) for k in self.recall_ks if int(k) > 0}) + if not ordered: + raise ValueError("RerankingMetric requires at least one positive Recall@K value.") + self.recall_ks = tuple(ordered) + if self.mrr_k <= 0: + raise ValueError("mrr_k must be positive.") + self._scored_groups = 0 + self._groups_without_positive = 0 + self._mrr_sum = 0.0 + self._recall_hits = dict.fromkeys(self.recall_ks, 0) + + def update(self, scores: list[float], labels: list[bool]) -> None: + """Update the aggregate with one ranked group.""" + if len(scores) != len(labels): + raise ValueError("scores and labels must have the same length.") + ranked = sorted(range(len(scores)), key=lambda index: (-scores[index], index)) + positive_ranks = [rank for rank, index in enumerate(ranked, start=1) if labels[index]] + if not positive_ranks: + self._groups_without_positive += 1 + return + + self._scored_groups += 1 + best_rank = positive_ranks[0] + if best_rank <= self.mrr_k: + self._mrr_sum += 1.0 / best_rank + + for k in self.recall_ks: + if any(rank <= k for rank in positive_ranks): + self._recall_hits[k] += 1 + + def compute(self) -> dict[str, float | int]: + """Return aggregated ranking metrics and accounting.""" + denom = self._scored_groups + metrics: dict[str, float | int] = { + f"mrr@{self.mrr_k}": round(self._mrr_sum / denom, 6) if denom else 0.0, + "scored_groups": denom, + "groups_without_positive": self._groups_without_positive, + } + for k in self.recall_ks: + metrics[f"recall@{k}"] = round(self._recall_hits[k] / denom, 6) if denom else 0.0 + return metrics diff --git a/src/winml/modelkit/eval/reranking_evaluator.py b/src/winml/modelkit/eval/reranking_evaluator.py new file mode 100644 index 000000000..2ae9b2ba4 --- /dev/null +++ b/src/winml/modelkit/eval/reranking_evaluator.py @@ -0,0 +1,414 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- + +"""Reranking evaluator for grouped query-document candidates.""" + +from __future__ import annotations + +import json +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, cast + +from transformers.pipelines.text_classification import TextClassificationPipeline + +from ..utils.eval_utils import detect_reranking_dataset_mode, get_dataset_column_names +from .base_evaluator import WinMLEvaluator + + +if TYPE_CHECKING: + import torch + from datasets import Dataset + from transformers.pipelines.base import Pipeline + from transformers.utils.generic import ModelOutput + + from ..models.winml.base import WinMLPreTrainedModel + from .config import DatasetConfig, WinMLEvaluationConfig + + +@dataclass(frozen=True) +class _Candidate: + candidate_id: str + text: str + relevant: bool + + +@dataclass(frozen=True) +class _Group: + group_id: str + query: str + candidates: tuple[_Candidate, ...] + + +class _RawRerankingPipeline(TextClassificationPipeline): + """Compatibility pipeline that preserves raw model outputs. + + Reranking evaluation reads raw relevance logits directly instead of the + text-classification pipeline's label/score postprocessing. + """ + + def postprocess( + self, + model_outputs: ModelOutput, + function_to_apply: Any = None, + top_k: Any = 1, + _legacy: Any = True, + **postprocess_parameters: dict[Any, Any], + ) -> ModelOutput: + return model_outputs + + +class WinMLRerankingEvaluator(WinMLEvaluator): + """Evaluator for cross-encoder reranking checkpoints.""" + + def __init__( + self, + config: WinMLEvaluationConfig, + model: WinMLPreTrainedModel, + ) -> None: + from transformers import AutoTokenizer + + from ..utils.eval_utils import get_default + + mapping = config.dataset.columns_mapping + task = "reranking" + self._query_col = mapping.get("query_column", get_default(task, "query_column") or "input") + self._expected_output_col = mapping.get( + "expected_output_column", + get_default(task, "expected_output_column") or "expected_output", + ) + self._metadata_col = mapping.get( + "metadata_column", + get_default(task, "metadata_column") or "metadata", + ) + self._candidates_col = mapping.get("candidates_column") + self._positive_col = mapping.get("positive_column") + self._negative_col = mapping.get("negative_column") + self._document_col = mapping.get("document_column") + self._group_col = mapping.get("group_column") + self._label_col = mapping.get("label_column") + self._candidate_id_col = mapping.get("candidate_id_column") + self._candidate_text_key = mapping.get( + "candidate_text_key", + get_default(task, "candidate_text_key") or "text", + ) + self._candidate_id_key = mapping.get( + "candidate_id_key", + get_default(task, "candidate_id_key") or "id", + ) + self._metadata_group_key = mapping.get( + "metadata_group_key", + get_default(task, "metadata_group_key") or "query_id", + ) + self._recall_ks = self._parse_recall_ks( + mapping.get("recall_ks", get_default(task, "recall_ks") or "1,10") + ) + self._max_candidates = self._parse_max_candidates( + mapping.get("max_candidates", get_default(task, "max_candidates") or "10") + ) + self._tokenizer = AutoTokenizer.from_pretrained( + config.model_id, + trust_remote_code=config.trust_remote_code, + ) + super().__init__(config, model) + + def prepare_pipeline(self) -> Pipeline: + """Return a Pipeline-compatible object without classification postprocessing.""" + from transformers import pipeline + + pipeline_kwargs: dict[str, Any] = { + "device": self.config.pipeline_device, + "pipeline_class": _RawRerankingPipeline, + "function_to_apply": "none", + } + if self.config.trust_remote_code: + pipeline_kwargs["trust_remote_code"] = True + + return cast( + "Pipeline", + pipeline( + "text-classification", + model=self.model, + tokenizer=self.config.model_id, + **pipeline_kwargs, + ), + ) + + def align_labels(self, dataset: Dataset, ds_config: DatasetConfig) -> Dataset: + """No class-label alignment for grouped relevance judgments.""" + return dataset + + def compute(self) -> dict[str, Any]: + """Score grouped candidates with raw relevance logits.""" + from .metrics import RerankingMetric + + groups = self._materialize_groups() + metric = RerankingMetric(recall_ks=self._recall_ks) + processed_groups = 0 + skipped_groups = 0 + processed_pairs = 0 + + for group in groups: + if not group.candidates: + skipped_groups += 1 + continue + scores: list[float] = [] + labels: list[bool] = [] + for candidate in group.candidates: + scores.append(self._score_pair(group.query, candidate.text)) + labels.append(candidate.relevant) + processed_groups += 1 + processed_pairs += len(group.candidates) + metric.update(scores, labels) + + result = metric.compute() + result.update( + { + "requested_rows": len(self.data), + "processed_groups": processed_groups, + "skipped_groups": skipped_groups, + "processed_pairs": processed_pairs, + "expanded_pairs": processed_pairs, + } + ) + return result + + def _materialize_groups(self) -> list[_Group]: + column_names = set(get_dataset_column_names(self.data)) + dataset_mode = detect_reranking_dataset_mode( + column_names, + self.config.dataset.columns_mapping, + ) + if dataset_mode == "pairwise": + return self._groups_from_pairwise_rows(column_names) + if dataset_mode == "grouped-text": + return self._groups_from_text_lists() + return self._groups_from_grouped_rows(column_names) + + def _groups_from_text_lists(self) -> list[_Group]: + from ..utils.eval_utils import DatasetValidationError + + assert self._positive_col is not None + assert self._negative_col is not None + groups: list[_Group] = [] + for row_index, sample in enumerate(self.data): + query = str(sample[self._query_col]) + if not query.strip(): + continue + positives = self._parse_json_sequence(sample[self._positive_col]) + negatives = self._parse_json_sequence(sample[self._negative_col]) + if not positives: + raise DatasetValidationError( + f"reranking group {row_index!r} has no positive passages" + ) + + candidates = [ + _Candidate( + candidate_id=f"{row_index}:positive:{candidate_index}", + text=str(text), + relevant=True, + ) + for candidate_index, text in enumerate(positives[: self._max_candidates]) + if str(text).strip() + ] + if not candidates: + raise DatasetValidationError( + f"reranking group {row_index!r} has no non-empty positive passages" + ) + remaining = max(self._max_candidates - len(candidates), 0) + candidates.extend( + _Candidate( + candidate_id=f"{row_index}:negative:{candidate_index}", + text=str(text), + relevant=False, + ) + for candidate_index, text in enumerate(negatives[:remaining]) + if str(text).strip() + ) + groups.append( + _Group(group_id=str(row_index), query=query, candidates=tuple(candidates)) + ) + return groups + + def _groups_from_pairwise_rows(self, column_names: set[str]) -> list[_Group]: + from ..utils.eval_utils import DatasetValidationError + + required = [self._query_col, self._document_col, self._group_col, self._label_col] + required_names = [name for name in required if name is not None] + missing = [name for name in required_names if name not in column_names] + if missing: + raise DatasetValidationError( + f"pairwise reranking dataset is missing required column(s): {sorted(missing)}" + ) + + grouped: dict[str, list[_Candidate]] = {} + queries: dict[str, str] = {} + for row_index, sample in enumerate(self.data): + group_id = str(sample[self._group_col]) + query = str(sample[self._query_col]) + document = str(sample[self._document_col]) + if not query.strip() or not document.strip(): + continue + previous_query = queries.get(group_id) + if previous_query is not None and previous_query != query: + raise DatasetValidationError( + f"group {group_id!r} contains inconsistent query text across rows" + ) + queries[group_id] = query + grouped.setdefault(group_id, []).append( + _Candidate( + candidate_id=str( + sample.get(self._candidate_id_col, f"{group_id}:{row_index}") + if self._candidate_id_col + else f"{group_id}:{row_index}" + ), + text=document, + relevant=self._parse_label(sample[self._label_col]), + ) + ) + + return [ + _Group(group_id=group_id, query=queries[group_id], candidates=tuple(candidates)) + for group_id, candidates in grouped.items() + ] + + def _groups_from_grouped_rows(self, column_names: set[str]) -> list[_Group]: + from ..utils.eval_utils import DatasetValidationError + + required = [self._query_col, self._expected_output_col, self._metadata_col] + missing = [name for name in required if name not in column_names] + if missing: + raise DatasetValidationError( + "reranking datasets require either pairwise columns " + "(query/document/group/label) or grouped authoritative columns " + f"({sorted(required)}); missing {sorted(missing)}" + ) + + if self._candidates_col is None: + raise DatasetValidationError( + "grouped reranking rows require --column candidates_column= or a " + "dataset script that materializes inline candidate passages; the authoritative " + "MS MARCO snapshot only stores relevant passage IDs, not candidate text." + ) + if self._candidates_col not in column_names: + raise DatasetValidationError( + f"grouped reranking dataset is missing candidates column {self._candidates_col!r}" + ) + + groups: list[_Group] = [] + for row_index, sample in enumerate(self.data): + query = str(sample[self._query_col]) + if not query.strip(): + continue + relevant_ids = set(self._parse_json_sequence(sample[self._expected_output_col])) + metadata = self._parse_json_object(sample[self._metadata_col]) + group_id = str(metadata.get(self._metadata_group_key, row_index)) + raw_candidates = self._parse_json_sequence(sample[self._candidates_col]) + candidates: list[_Candidate] = [] + for candidate in raw_candidates: + if not isinstance(candidate, dict): + raise DatasetValidationError( + f"group {group_id!r} has malformed candidate entry {candidate!r}" + ) + candidate_id = candidate.get(self._candidate_id_key) + text = candidate.get(self._candidate_text_key) + if candidate_id is None or text is None or not str(text).strip(): + raise DatasetValidationError( + f"group {group_id!r} candidates must expose non-empty " + f"{self._candidate_id_key!r} and {self._candidate_text_key!r} fields" + ) + candidate_id_text = str(candidate_id) + candidates.append( + _Candidate( + candidate_id=candidate_id_text, + text=str(text), + relevant=candidate_id_text in relevant_ids, + ) + ) + groups.append(_Group(group_id=group_id, query=query, candidates=tuple(candidates))) + return groups + + def _score_pair(self, query: str, document: str) -> float: + import torch + + tokenizer_kwargs: dict[str, Any] = { + "truncation": True, + "return_tensors": "pt", + } + max_length = self._fixed_seq_length() + if max_length is not None: + tokenizer_kwargs["padding"] = "max_length" + tokenizer_kwargs["max_length"] = max_length + + encoding = self._tokenizer(query, document, **tokenizer_kwargs) + encoding = self._pad_or_truncate(encoding, self._tokenizer) + tensor_encoding = { + name: value.to(self.config.pipeline_device) + for name, value in encoding.items() + if isinstance(value, torch.Tensor) + } + with torch.no_grad(): + outputs = self.model(**tensor_encoding) + return self._extract_relevance_score(outputs) + + def _extract_relevance_score(self, outputs: Any) -> float: + logits = outputs["logits"] if isinstance(outputs, dict) else outputs.logits + tensor = cast("torch.Tensor", logits) + if tensor.numel() != 1: + raise ValueError( + "reranking expects exactly one logit per query-document pair; " + f"got shape {tuple(tensor.shape)}" + ) + return float(tensor.reshape(-1)[0].item()) + + @staticmethod + def _parse_label(value: Any) -> bool: + if isinstance(value, bool): + return value + if isinstance(value, (int, float)): + return bool(value) + text = str(value).strip().lower() + if text in {"1", "true", "yes", "positive", "relevant"}: + return True + if text in {"0", "false", "no", "negative", "irrelevant"}: + return False + raise ValueError(f"unsupported reranking relevance label: {value!r}") + + @staticmethod + def _parse_json_sequence(value: Any) -> list[Any]: + if isinstance(value, list): + return value + if isinstance(value, tuple): + return list(value) + if isinstance(value, str): + parsed = json.loads(value) + if not isinstance(parsed, list): + raise TypeError(f"expected a JSON list, got {type(parsed).__name__}") + return parsed + raise TypeError(f"expected a list or JSON list string, got {type(value).__name__}") + + @staticmethod + def _parse_json_object(value: Any) -> dict[str, Any]: + if isinstance(value, dict): + return value + if isinstance(value, str): + parsed = json.loads(value) + if not isinstance(parsed, dict): + raise TypeError(f"expected a JSON object, got {type(parsed).__name__}") + return cast("dict[str, Any]", parsed) + raise TypeError(f"expected a dict or JSON object string, got {type(value).__name__}") + + @staticmethod + def _parse_recall_ks(raw: str) -> tuple[int, ...]: + ks = tuple(sorted({int(part.strip()) for part in raw.split(",") if part.strip()})) + if not ks or any(k <= 0 for k in ks): + raise ValueError(f"invalid recall_ks setting: {raw!r}") + return ks + + @staticmethod + def _parse_max_candidates(raw: str) -> int: + value = int(raw) + if value <= 0: + raise ValueError(f"invalid max_candidates setting: {raw!r}") + return value diff --git a/src/winml/modelkit/inference/pipeline.py b/src/winml/modelkit/inference/pipeline.py index 36b56ee58..83b85b7b7 100644 --- a/src/winml/modelkit/inference/pipeline.py +++ b/src/winml/modelkit/inference/pipeline.py @@ -31,6 +31,7 @@ # Mapped to their HF pipeline equivalent before calling ``pipeline()``. _HF_PIPELINE_TASK_MAP: dict[str, str] = { "image-to-text": "image-text-to-text", + "reranking": "text-classification", "next-sentence-prediction": "text-classification", "sequence-classification": "text-classification", "sentence-similarity": "feature-extraction", diff --git a/src/winml/modelkit/inference/tasks.py b/src/winml/modelkit/inference/tasks.py index b2d1aa208..fa3168367 100644 --- a/src/winml/modelkit/inference/tasks.py +++ b/src/winml/modelkit/inference/tasks.py @@ -362,6 +362,18 @@ def _postprocess_sentence_similarity( mapping=PipelineMapping(pipe_input="text"), ), # -- Text + text -------------------------------------------------------- + "reranking": TaskInputSpec( + user_inputs=[ + InputField(name="query", type="text", required=True, description="Search query"), + InputField( + name="document", + type="text", + required=True, + description="Candidate document to score", + ), + ], + mapping=PipelineMapping(pipe_input=["query", "document"]), + ), "question-answering": TaskInputSpec( user_inputs=[ InputField( diff --git a/src/winml/modelkit/loader/resolution.py b/src/winml/modelkit/loader/resolution.py index a1a372353..7fdba271d 100644 --- a/src/winml/modelkit/loader/resolution.py +++ b/src/winml/modelkit/loader/resolution.py @@ -414,6 +414,24 @@ def resolve_composite_load_task( _SEQ2SEQ_GENERATION_TASK = "text2text-generation" +def _surface_detected_task(config: PretrainedConfig, opt_task: str, model_id: str | None) -> str: + """Return the surfaced WinML task for a detected Optimum task. + + Keeps export/model-class resolution on Optimum's canonical task while allowing + user-facing task semantics to upgrade when authoritative metadata carries a + narrower meaning. + """ + surfaced = _resolve_task_modality(config, opt_task) + if surfaced != "text-classification" or not model_id: + return surfaced + + from ..utils.hub_utils import get_pipeline_tag + + if normalize_task(get_pipeline_tag(model_id) or "") == "reranking": + return "reranking" + return surfaced + + def _infer_task_from_architecture(config: PretrainedConfig) -> str: """Optimum task inferred from ``config.architectures[0]``. @@ -518,8 +536,9 @@ def resolve_task( # pixel_values arch). (b) is a no-op for non-feature-extraction tasks, so (a) # is preserved. Consistent with the inferred branch below and USER_TASK — # adding --model-class must not collapse the modality. - opt_task = normalize_task(task) - surfaced = _resolve_task_modality(config, opt_task) + surfaced_task = normalize_task(task) + opt_task = to_optimum_task(surfaced_task) + surfaced = _resolve_task_modality(config, surfaced_task) else: # Task inferred from the architecture: surface it modality-aware, consistent # with the detection path (Stage 3), so e.g. a ViT backbone is @@ -556,6 +575,7 @@ def resolve_task( if task is not None: original = task normalized = normalize_task(task) + optimum_task = to_optimum_task(normalized) # Exact-key composite lookup on the ORIGINAL user string: registration keys are # `summarization` / `table-question-answering`, never the normalized # `text2text-generation`. So `--task summarization` tags the composite while @@ -569,7 +589,7 @@ def resolve_task( if resolved is None: try: resolved = TasksManager.get_model_class_for_task( - normalized, framework="pt", model_type=model_type or None + optimum_task, framework="pt", model_type=model_type or None ) except KeyError as e: if composite is not None: @@ -583,8 +603,13 @@ def resolve_task( f"Task '{normalized}' not supported by TasksManager. " f"Check optimum documentation for supported tasks." ) from e + surfaced_task = normalized if original == "text-ranking" else original return TaskResolution( - original, to_optimum_task(original), resolved, TaskSource.USER_TASK, composite + surfaced_task, + to_optimum_task(surfaced_task), + resolved, + TaskSource.USER_TASK, + composite, ) # --- Stage 1: detection ----------------------------------------------- @@ -664,7 +689,7 @@ def resolve_task( resolved = _resolve_model_class_from_config(config) # arch fallback # --- Stage 3: modality upgrade (surfaced task only) ------------------- - surfaced = _resolve_task_modality(config, opt_task) + surfaced = _surface_detected_task(config, opt_task, model_id) # --- Stage 4: composite tag (detection path) -------------------------- composite = _composite_components_for_task(model_type, opt_task) if model_type else None diff --git a/src/winml/modelkit/loader/task.py b/src/winml/modelkit/loader/task.py index 0ae0ca6be..11dcb4ecf 100644 --- a/src/winml/modelkit/loader/task.py +++ b/src/winml/modelkit/loader/task.py @@ -65,6 +65,7 @@ "inpainting": None, "text-to-image": None, # NLP + "reranking": "rerank", "text-classification": "txtcls", "token-classification": "tokcls", "question-answering": "qa", @@ -236,6 +237,9 @@ def normalize_task(task: str) -> str: Returns: Canonical task name """ + if task in {"text-ranking", "reranking"}: + return "reranking" + from optimum.exporters.tasks import TasksManager return cast("str", TasksManager.map_from_synonym(task)) @@ -244,6 +248,7 @@ def normalize_task(task: str) -> str: # WinML task-synonym extensions — extend Optimum's ``TasksManager.map_from_synonym`` # for tasks it does not recognize or mis-maps. Entries here take priority over Optimum. TASK_SYNONYM_EXTENSIONS: dict[str, str] = { + "reranking": "text-classification", # NOTE: do NOT add "image-feature-extraction" here. This set is also consulted by # commands.build._validate_task_supported_for_model (its "WinML extension" branch), # so adding it would silence the cross-modality visibility warning. Its Optimum-synonym diff --git a/src/winml/modelkit/models/winml/__init__.py b/src/winml/modelkit/models/winml/__init__.py index 251c40dad..5bd8ca8ac 100644 --- a/src/winml/modelkit/models/winml/__init__.py +++ b/src/winml/modelkit/models/winml/__init__.py @@ -34,6 +34,7 @@ TASK_TO_WINML_CLASS: dict[str, str] = { # Implemented "image-classification": "WinMLModelForImageClassification", + "reranking": "WinMLModelForSequenceClassification", "text-classification": "WinMLModelForSequenceClassification", "sequence-classification": "WinMLModelForSequenceClassification", "next-sentence-prediction": "WinMLModelForSequenceClassification", diff --git a/src/winml/modelkit/utils/eval_utils.py b/src/winml/modelkit/utils/eval_utils.py index f85c672b4..efda25a30 100644 --- a/src/winml/modelkit/utils/eval_utils.py +++ b/src/winml/modelkit/utils/eval_utils.py @@ -12,6 +12,7 @@ from __future__ import annotations +from collections.abc import Mapping, Sequence from dataclasses import dataclass from typing import Literal, TypeAlias, get_args @@ -79,6 +80,96 @@ class TaskSchema: ), ) +_RERANKING_SCHEMA = TaskSchema( + columns=( + SchemaItem( + "query_column", + "query text (grouped rows typically use the authoritative 'input' column)", + default="input", + remap_hint="", + ), + SchemaItem( + "expected_output_column", + "JSON/list of relevant candidate IDs for grouped rows", + default="expected_output", + remap_hint="", + ), + SchemaItem( + "metadata_column", + "metadata dict/JSON for grouped rows (used for query_id and provenance)", + default="metadata", + remap_hint="", + ), + SchemaItem( + "candidates_column", + "inline candidate list for grouped rows; each item must expose text and ID fields", + remap_hint="", + ), + SchemaItem( + "positive_column", + "relevant passage text list for grouped rows", + remap_hint="", + ), + SchemaItem( + "negative_column", + "non-relevant passage text list for grouped rows", + remap_hint="", + ), + SchemaItem( + "document_column", + "candidate document text for pre-expanded pairwise rows", + remap_hint="", + ), + SchemaItem( + "group_column", + "group/query identifier for pre-expanded pairwise rows", + remap_hint="", + ), + SchemaItem( + "label_column", + "binary relevance flag for pre-expanded pairwise rows", + remap_hint="", + ), + SchemaItem( + "candidate_id_column", + "candidate identifier for pre-expanded pairwise rows", + remap_hint="", + ), + ), + params=( + SchemaItem( + "candidate_text_key", + "candidate text field inside grouped-row candidates", + default="text", + remap_hint="", + ), + SchemaItem( + "candidate_id_key", + "candidate ID field inside grouped-row candidates", + default="id", + remap_hint="", + ), + SchemaItem( + "metadata_group_key", + "group/query identifier field inside grouped-row metadata", + default="query_id", + remap_hint="", + ), + SchemaItem( + "recall_ks", + "comma-separated K values for Recall@K", + default="1,10", + remap_hint="", + ), + SchemaItem( + "max_candidates", + "maximum candidates materialized from positive/negative passage lists", + default="10", + remap_hint="", + ), + ), +) + _TOKEN_CLASSIFICATION_SCHEMA = TaskSchema( columns=( SchemaItem( @@ -450,6 +541,7 @@ class TaskSchema: TASK_SCHEMAS: dict[str, TaskSchema] = { "image-classification": _IMAGE_CLASSIFICATION_SCHEMA, + "reranking": _RERANKING_SCHEMA, "text-classification": _TEXT_CLASSIFICATION_SCHEMA, "sequence-classification": _TEXT_CLASSIFICATION_SCHEMA, "next-sentence-prediction": _TEXT_CLASSIFICATION_SCHEMA, @@ -490,6 +582,87 @@ class DatasetValidationError(Exception): """Dataset failed schema validation against a task's expected columns.""" +RerankingDatasetMode: TypeAlias = Literal[ + "pairwise", + "grouped-inline", + "grouped-text", + "grouped-authoritative", +] + + +def get_dataset_column_names(dataset: object) -> tuple[str, ...]: + """Best-effort column-name extraction for datasets and list-backed test fixtures.""" + column_names = getattr(dataset, "column_names", None) + if isinstance(column_names, (list, tuple)): + return tuple(str(name) for name in column_names) + if isinstance(dataset, Sequence) and not isinstance(dataset, (str, bytes, bytearray)): + names: set[str] = set() + for row in dataset: + if isinstance(row, Mapping): + names.update(str(name) for name in row) + return tuple(sorted(names)) + return () + + +def _resolved_reranking_column(mapping: dict[str, str], key: str) -> str | None: + return mapping.get(key, get_default("reranking", key)) + + +def detect_reranking_dataset_mode( + column_names: set[str] | list[str] | tuple[str, ...], + columns_mapping: dict[str, str] | None = None, +) -> RerankingDatasetMode: + """Resolve reranking datasets to pairwise, grouped-inline, or grouped-authoritative.""" + mapping = columns_mapping or {} + actual = set(column_names) + + query_col = _resolved_reranking_column(mapping, "query_column") + expected_output_col = _resolved_reranking_column(mapping, "expected_output_column") + metadata_col = _resolved_reranking_column(mapping, "metadata_column") + document_col = mapping.get("document_column") + group_col = mapping.get("group_column") + label_col = mapping.get("label_column") + candidates_col = mapping.get("candidates_column") + positive_col = mapping.get("positive_column") + negative_col = mapping.get("negative_column") + + grouped_required = tuple( + name for name in (query_col, expected_output_col, metadata_col) if name is not None + ) + pairwise_required = tuple( + name for name in (query_col, document_col, group_col, label_col) if name is not None + ) + + has_grouped_core = len(grouped_required) == 3 and all( + name in actual for name in grouped_required + ) + has_pairwise = len(pairwise_required) == 4 and all(name in actual for name in pairwise_required) + has_grouped_text = ( + query_col is not None + and positive_col is not None + and negative_col is not None + and all(name in actual for name in (query_col, positive_col, negative_col)) + ) + + if has_grouped_core and candidates_col is not None and candidates_col in actual: + return "grouped-inline" + if has_pairwise: + return "pairwise" + if has_grouped_text: + return "grouped-text" + if has_grouped_core: + return "grouped-authoritative" + + grouped_missing = sorted(name for name in grouped_required if name not in actual) + pairwise_missing = sorted(name for name in pairwise_required if name not in actual) + raise DatasetValidationError( + "reranking datasets require pairwise columns " + f"{sorted(pairwise_required)} or grouped authoritative columns {sorted(grouped_required)}; " + f"missing pairwise={pairwise_missing} grouped={grouped_missing}; " + f"dataset has {sorted(actual)}" + ) + + def validate_dataset_columns( dataset: object, task: str, @@ -507,6 +680,9 @@ def validate_dataset_columns( return mapping = columns_mapping or {} actual = set(column_names) + if task == "reranking": + detect_reranking_dataset_mode(actual, mapping) + return missing = [ (item.name, mapping.get(item.name, item.default)) for item in schema.columns diff --git a/tests/unit/commands/test_eval.py b/tests/unit/commands/test_eval.py index 509594389..6399527f9 100644 --- a/tests/unit/commands/test_eval.py +++ b/tests/unit/commands/test_eval.py @@ -8,7 +8,7 @@ from __future__ import annotations import json -from unittest.mock import patch +from unittest.mock import MagicMock, patch import click import pytest @@ -1017,6 +1017,7 @@ def compute(self): ("zero-shot-classification", "fancyzhx/ag_news", "test"), ("zero-shot-image-classification", "uoft-cs/cifar100", "test"), ("image-classification", "timm/mini-imagenet", "test"), + ("reranking", "mteb/scidocs-reranking", "test"), ], ) def test_per_task_default_split_reaches_evaluator( @@ -1078,6 +1079,96 @@ def test_user_split_ignored_when_default_dataset_used( ) assert cfg.dataset.split == "test" # the default's split wins + def test_reranking_default_runs_real_evaluator_with_bounded_candidates( + self, + runner: CliRunner, + onnx_file, + ) -> None: + from types import SimpleNamespace + + import torch + + from winml.modelkit.commands.eval import eval as eval_cmd + from winml.modelkit.eval.reranking_evaluator import WinMLRerankingEvaluator + + public_row = { + "query": "A Direct Search Method to solve Economic Dispatch Problem", + "positive": [f"relevant passage {index}" for index in range(5)], + "negative": [f"negative passage {index}" for index in range(25)], + } + + class _Dataset: + def __init__(self, rows): + self.rows = rows + self.column_names = ["query", "positive", "negative"] + self.features = None + + def __len__(self): + return len(self.rows) + + def __iter__(self): + return iter(self.rows) + + def shuffle(self, **_kwargs): + return self + + def take(self, count): + return iter(self.rows[:count]) + + def select(self, indices): + return _Dataset([self.rows[index] for index in indices]) + + class _Tokenizer: + def __call__(self, _query, _document, **_kwargs): + values = torch.ones((1, 4), dtype=torch.int64) + return {"input_ids": values, "attention_mask": values} + + def pad(self, encoding, **_kwargs): + return encoding + + model = MagicMock() + model.io_config = {"input_shapes": [[1, 4]]} + model.return_value = SimpleNamespace(logits=torch.tensor([[0.5]])) + + with ( + patch("winml.modelkit.models.WinMLAutoModel.from_onnx", return_value=model), + patch("winml.modelkit.loader.load_hf_config", return_value=MagicMock()), + patch("datasets.load_dataset", return_value=_Dataset([public_row])) as load, + patch("datasets.Dataset.from_list", return_value=_Dataset([public_row])), + patch("transformers.AutoTokenizer.from_pretrained", return_value=_Tokenizer()), + patch.object(WinMLRerankingEvaluator, "prepare_pipeline", return_value=object()), + patch("winml.modelkit.commands.eval._resolve_device", return_value=None), + patch("winml.modelkit.commands.eval._write_and_display") as display, + ): + result = runner.invoke( + eval_cmd, + [ + "-m", + str(onnx_file), + "--model-id", + "cross-encoder/ms-marco-MiniLM-L6-v2", + "--task", + "reranking", + "--ep", + "cpu", + "--device", + "cpu", + "--samples", + "1", + ], + obj={"debug": False}, + ) + + assert result.exit_code == 0, result.output + assert load.call_args.args == ("mteb/scidocs-reranking",) + assert load.call_args.kwargs["revision"] == "56a6d0140cf6356659e2a7c1413286a774468d44" + assert load.call_args.kwargs["streaming"] is True + assert model.call_count == 10 + eval_result = display.call_args.args[0] + assert eval_result.metrics["processed_groups"] == 1 + assert eval_result.metrics["processed_pairs"] == 10 + assert eval_result.metrics["groups_without_positive"] == 0 + def test_user_column_merged_when_default_dataset_used( self, runner: CliRunner, diff --git a/tests/unit/eval/test_eval.py b/tests/unit/eval/test_eval.py index e41ad6e8e..31d411e54 100644 --- a/tests/unit/eval/test_eval.py +++ b/tests/unit/eval/test_eval.py @@ -304,6 +304,12 @@ def test_registered_task_returns_class(self): assert cls.__module__ == module_path assert cls.__name__ == class_name + def test_text_classification_still_uses_classification_evaluator(self) -> None: + from winml.modelkit.eval import WinMLEvaluationConfig, get_evaluator_class + + cls = get_evaluator_class(WinMLEvaluationConfig(task="text-classification")) + assert cls.__name__ == "WinMLTextClassificationEvaluator" + def test_unsupported_task_raises_value_error(self): from winml.modelkit.eval import WinMLEvaluationConfig, get_evaluator_class @@ -702,6 +708,76 @@ def test_revision_defaults_to_none( mock_load_ds.assert_called_once() assert mock_load_ds.call_args.kwargs["revision"] is None + @patch("evaluate.evaluator") + @patch("transformers.pipeline") + @patch("datasets.load_from_disk") + def test_local_dataset_dict_uses_requested_split( + self, + mock_load_from_disk, + mock_pipeline, + mock_hf_eval, + tmp_path, + ): + from winml.modelkit.eval import WinMLEvaluator + + local_dir = tmp_path / "fixture" + local_dir.mkdir() + + dev_ds = MagicMock() + dev_ds.__len__ = lambda self: 3 + dev_ds.shuffle.return_value = dev_ds + dev_ds.select.return_value = dev_ds + dev_ds.column_names = ["image", "label"] + + train_ds = MagicMock() + train_ds.__len__ = lambda self: 5 + train_ds.shuffle.return_value = train_ds + train_ds.select.return_value = train_ds + train_ds.column_names = ["image", "label"] + + mock_load_from_disk.return_value = {"train": train_ds, "dev": dev_ds} + mock_pipeline.return_value = MagicMock() + mock_hf_eval.return_value = MagicMock(compute=MagicMock(return_value={})) + + model = MagicMock() + model.config.label2id = None + + config = WinMLEvaluationConfig( + model_id="test/model", + task="image-classification", + dataset=DatasetConfig(path=str(local_dir), split="dev", samples=2), + ) + + WinMLEvaluator(config, model) + + dev_ds.select.assert_called_once_with(range(2)) + train_ds.select.assert_not_called() + + @patch("datasets.load_from_disk") + def test_local_dataset_dict_missing_split_raises( + self, + mock_load_from_disk, + tmp_path, + ): + from winml.modelkit.eval import WinMLEvaluator + from winml.modelkit.utils.eval_utils import DatasetValidationError + + local_dir = tmp_path / "fixture" + local_dir.mkdir() + + mock_load_from_disk.return_value = {"train": MagicMock()} + + model = MagicMock() + model.config.label2id = None + config = WinMLEvaluationConfig( + model_id="test/model", + task="image-classification", + dataset=DatasetConfig(path=str(local_dir), split="dev", samples=1), + ) + + with pytest.raises(DatasetValidationError, match="has splits"): + WinMLEvaluator(config, model) + @patch("evaluate.evaluator") @patch("transformers.pipeline") @patch("datasets.load_dataset") diff --git a/tests/unit/eval/test_reranking_evaluator.py b/tests/unit/eval/test_reranking_evaluator.py new file mode 100644 index 000000000..70e9fef7d --- /dev/null +++ b/tests/unit/eval/test_reranking_evaluator.py @@ -0,0 +1,390 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- + +from __future__ import annotations + +import importlib.util +import sys +from pathlib import Path +from types import SimpleNamespace + +import pytest +import torch + +from winml.modelkit.eval import DatasetConfig, WinMLEvaluationConfig +from winml.modelkit.eval.metrics.ranking import RerankingMetric +from winml.modelkit.eval.reranking_evaluator import WinMLRerankingEvaluator +from winml.modelkit.utils.eval_utils import ( + DatasetValidationError, + detect_reranking_dataset_mode, +) + + +_FIXTURE_BUILDER_PATH = ( + Path(__file__).resolve().parents[3] + / "scripts" + / "e2e_eval" + / "datasets" + / "build_msmarco_reranking_fixture.py" +) +_FIXTURE_BUILDER_SPEC = importlib.util.spec_from_file_location( + "build_msmarco_reranking_fixture", + _FIXTURE_BUILDER_PATH, +) +assert _FIXTURE_BUILDER_SPEC is not None +assert _FIXTURE_BUILDER_SPEC.loader is not None +_FIXTURE_BUILDER = importlib.util.module_from_spec(_FIXTURE_BUILDER_SPEC) +sys.modules[_FIXTURE_BUILDER_SPEC.name] = _FIXTURE_BUILDER +_FIXTURE_BUILDER_SPEC.loader.exec_module(_FIXTURE_BUILDER) + +CandidateRow = _FIXTURE_BUILDER.CandidateRow +select_rows = _FIXTURE_BUILDER._select_rows + + +class _FakeTokenizer: + def __call__(self, query: str, document: str, **_kwargs): + width = max(len(query), len(document), 1) + values = torch.ones((1, min(width, 4)), dtype=torch.int64) + return { + "input_ids": values, + "attention_mask": torch.ones_like(values), + } + + def pad(self, encoding, **_kwargs): + return encoding + + +class _FakeModel: + def __init__(self, scores: list[float]): + self._scores = list(scores) + self.io_config = {"input_shapes": [[1, 4]]} + + def __call__(self, **_kwargs): + score = self._scores.pop(0) + return SimpleNamespace(logits=torch.tensor([[score]], dtype=torch.float32)) + + +def _make_evaluator(data, scores: list[float]) -> WinMLRerankingEvaluator: + evaluator = WinMLRerankingEvaluator.__new__(WinMLRerankingEvaluator) + evaluator.config = WinMLEvaluationConfig( + model_id="cross-encoder/ms-marco-MiniLM-L6-v2", + task="reranking", + dataset=DatasetConfig( + path="dummy", + columns_mapping={ + "query_column": "query", + "document_column": "document", + "group_column": "group_id", + "label_column": "label", + "candidate_id_column": "candidate_id", + "recall_ks": "1,2,10", + }, + ), + ) + evaluator.model = _FakeModel(scores) + evaluator.data = data + evaluator._query_col = "query" + evaluator._expected_output_col = "expected_output" + evaluator._metadata_col = "metadata" + evaluator._candidates_col = None + evaluator._positive_col = None + evaluator._negative_col = None + evaluator._document_col = "document" + evaluator._group_col = "group_id" + evaluator._label_col = "label" + evaluator._candidate_id_col = "candidate_id" + evaluator._candidate_text_key = "text" + evaluator._candidate_id_key = "id" + evaluator._metadata_group_key = "query_id" + evaluator._recall_ks = (1, 2, 10) + evaluator._max_candidates = 10 + evaluator._tokenizer = _FakeTokenizer() + return evaluator + + +def test_reranking_metric_handles_ties_and_no_positive_groups() -> None: + metric = RerankingMetric(recall_ks=(1, 2, 10)) + metric.update([0.9, 0.9, 0.1], [False, True, False]) + metric.update([0.2, 0.1], [False, False]) + + result = metric.compute() + + assert result["mrr@10"] == 0.5 + assert result["recall@1"] == 0.0 + assert result["recall@2"] == 1.0 + assert result["groups_without_positive"] == 1 + assert result["scored_groups"] == 1 + + +def test_reranking_metric_ties_preserve_authoritative_candidate_order() -> None: + metric = RerankingMetric(recall_ks=(1, 2, 10)) + + metric.update([0.9, 0.9, 0.9], [False, False, True]) + + result = metric.compute() + + assert result["mrr@10"] == pytest.approx(1 / 3) + assert result["recall@1"] == 0.0 + assert result["recall@2"] == 0.0 + assert result["recall@10"] == 1.0 + + +def test_reranking_evaluator_scores_single_logits_and_accounts_for_groups() -> None: + evaluator = _make_evaluator( + [ + { + "query": "what is pcnt", + "document": "negative passage", + "group_id": "q1", + "label": 0, + "candidate_id": "n1", + }, + { + "query": "what is pcnt", + "document": "positive passage", + "group_id": "q1", + "label": 1, + "candidate_id": "p1", + }, + { + "query": "cost of endless pools/swim spa", + "document": "positive first hit", + "group_id": "q2", + "label": 1, + "candidate_id": "p2", + }, + ], + scores=[0.2, 0.8, 0.7], + ) + + result = evaluator.compute() + + assert result["mrr@10"] == 1.0 + assert result["recall@1"] == 1.0 + assert result["processed_groups"] == 2 + assert result["processed_pairs"] == 3 + assert result["expanded_pairs"] == 3 + assert result["skipped_groups"] == 0 + + +def test_reranking_evaluator_rejects_grouped_rows_without_candidates() -> None: + evaluator = _make_evaluator( + [ + { + "query": "what is pcnt", + "expected_output": '["7187227"]', + "metadata": '{"query_id": "q1"}', + } + ], + scores=[], + ) + evaluator._query_col = "query" + evaluator._document_col = None + evaluator._group_col = None + evaluator._label_col = None + + with pytest.raises(DatasetValidationError, match="candidates_column"): + evaluator.compute() + + +def test_reranking_evaluator_scores_grouped_rows_with_inline_candidates() -> None: + evaluator = _make_evaluator( + [ + { + "query": "what is pcnt", + "expected_output": ["7187227"], + "metadata": {"query_id": "1048579", "source_row_index": 1}, + "candidates": [ + {"id": "n1", "text": "negative passage"}, + {"id": "7187227", "text": "positive passage"}, + ], + } + ], + scores=[0.1, 0.9], + ) + evaluator._document_col = None + evaluator._group_col = None + evaluator._label_col = None + evaluator._candidate_id_col = None + evaluator._candidates_col = "candidates" + result = evaluator.compute() + + assert result["mrr@10"] == 1.0 + assert result["recall@1"] == 1.0 + assert result["processed_groups"] == 1 + assert result["processed_pairs"] == 2 + + +def test_reranking_evaluator_materializes_bounded_positive_and_negative_text() -> None: + evaluator = _make_evaluator( + [ + { + "query": "economic dispatch", + "positive": ["relevant one", "relevant two"], + "negative": ["negative one", "negative two", "negative three"], + } + ], + scores=[0.9, 0.8, 0.1], + ) + evaluator._positive_col = "positive" + evaluator._negative_col = "negative" + evaluator._document_col = None + evaluator._group_col = None + evaluator._label_col = None + evaluator._max_candidates = 3 + evaluator.config.dataset.columns_mapping = { + "query_column": "query", + "positive_column": "positive", + "negative_column": "negative", + } + + groups = evaluator._materialize_groups() + result = evaluator.compute() + + assert [candidate.candidate_id for candidate in groups[0].candidates] == [ + "0:positive:0", + "0:positive:1", + "0:negative:0", + ] + assert [candidate.relevant for candidate in groups[0].candidates] == [True, True, False] + assert result["processed_pairs"] == 3 + assert result["recall@1"] == 1.0 + + +def test_reranking_evaluator_does_not_cap_materialized_candidate_column() -> None: + evaluator = _make_evaluator( + [ + { + "query": "what is pcnt", + "expected_output": ["p1"], + "metadata": {"query_id": "q1"}, + "candidates": [ + {"id": "n1", "text": "negative one"}, + {"id": "n2", "text": "negative two"}, + {"id": "p1", "text": "positive"}, + ], + } + ], + scores=[0.1, 0.2, 0.9], + ) + evaluator._document_col = None + evaluator._group_col = None + evaluator._label_col = None + evaluator._candidates_col = "candidates" + evaluator._max_candidates = 1 + + result = evaluator.compute() + + assert result["processed_pairs"] == 3 + assert result["recall@1"] == 1.0 + + +def test_reranking_evaluator_grouped_inline_ties_keep_original_candidate_order() -> None: + evaluator = _make_evaluator( + [ + { + "query": "what is pcnt", + "expected_output": ["7187227"], + "metadata": {"query_id": "1048579", "source_row_index": 1}, + "candidates": [ + {"id": "n1", "text": "negative one"}, + {"id": "n2", "text": "negative two"}, + {"id": "7187227", "text": "positive passage"}, + ], + } + ], + scores=[0.5, 0.5, 0.5], + ) + evaluator._document_col = None + evaluator._group_col = None + evaluator._label_col = None + evaluator._candidate_id_col = None + evaluator._candidates_col = "candidates" + + result = evaluator.compute() + + assert result["mrr@10"] == pytest.approx(1 / 3) + assert result["recall@1"] == 0.0 + assert result["recall@2"] == 0.0 + assert result["recall@10"] == 1.0 + assert result["processed_groups"] == 1 + assert result["processed_pairs"] == 3 + + +def test_fixture_builder_preserves_authoritative_order_when_negative_precedes_positive() -> None: + hf_rows = [ + { + "input": "what is pcnt", + "expected_output": ["p1"], + "metadata": {"query_id": "q1"}, + } + ] + queries = {"q1": "what is pcnt"} + qrels = {"q1": {"p1"}} + top1000 = { + "q1": [ + CandidateRow(pid="n1", query="what is pcnt", passage="negative one", rank=1), + CandidateRow(pid="n2", query="what is pcnt", passage="negative two", rank=2), + CandidateRow(pid="p1", query="what is pcnt", passage="positive", rank=3), + CandidateRow(pid="n3", query="what is pcnt", passage="negative three", rank=4), + ] + } + + selected_rows, provenance = select_rows( + hf_rows, + queries, + qrels, + top1000, + max_queries=1, + max_negatives=2, + ) + + assert [candidate["id"] for candidate in selected_rows[0]["candidates"]] == ["n1", "n2", "p1"] + assert [candidate["relevant"] for candidate in selected_rows[0]["candidates"]] == [ + False, + False, + True, + ] + assert selected_rows[0]["metadata"]["selected_candidate_ids"] == ["n1", "n2", "p1"] + assert selected_rows[0]["metadata"]["positive_candidate_ids"] == ["p1"] + assert selected_rows[0]["metadata"]["negative_candidate_ids"] == ["n1", "n2"] + assert provenance[0]["selected_candidate_ids"] == ["n1", "n2", "p1"] + assert provenance[0]["candidate_ranks"] == {"n1": 1, "n2": 2, "p1": 3} + + +def test_reranking_dataset_mode_prefers_grouped_inline_candidates() -> None: + mode = detect_reranking_dataset_mode( + ["input", "expected_output", "metadata", "candidates"], + { + "query_column": "input", + "expected_output_column": "expected_output", + "metadata_column": "metadata", + "candidates_column": "candidates", + }, + ) + + assert mode == "grouped-inline" + + +def test_reranking_dataset_mode_accepts_pairwise_rows_without_grouped_columns() -> None: + mode = detect_reranking_dataset_mode( + ["query", "document", "group_id", "label"], + { + "query_column": "query", + "document_column": "document", + "group_column": "group_id", + "label_column": "label", + }, + ) + + assert mode == "pairwise" + + +def test_reranking_evaluator_rejects_multi_logit_classification_outputs() -> None: + evaluator = _make_evaluator([], scores=[]) + outputs = SimpleNamespace(logits=torch.tensor([[0.1, 0.9]], dtype=torch.float32)) + + with pytest.raises(ValueError, match="exactly one logit"): + evaluator._extract_relevance_score(outputs) diff --git a/tests/unit/inference/test_pipeline.py b/tests/unit/inference/test_pipeline.py index 2482e3218..3d66dce8e 100644 --- a/tests/unit/inference/test_pipeline.py +++ b/tests/unit/inference/test_pipeline.py @@ -162,7 +162,7 @@ def test_image_to_text_maps_to_transformers_5_name(self) -> None: @pytest.mark.parametrize( "task", - ["sequence-classification", "next-sentence-prediction"], + ["reranking", "sequence-classification", "next-sentence-prediction"], ) def test_classification_aliases_map_to_transformers_task(self, task: str) -> None: assert _HF_PIPELINE_TASK_MAP[task] == "text-classification" diff --git a/tests/unit/loader/test_detect_task.py b/tests/unit/loader/test_detect_task.py index bdb0046a9..e638136d5 100644 --- a/tests/unit/loader/test_detect_task.py +++ b/tests/unit/loader/test_detect_task.py @@ -252,6 +252,19 @@ def test_resolve_task_uses_pipeline_tag_when_architecture_fails() -> None: assert r.source == TaskSource.PIPELINE_TAG +def test_resolve_task_surfaces_reranking_from_pipeline_tag() -> None: + """A text-ranking Hub pipeline tag upgrades a sequence classifier's surfaced task.""" + cfg = _FakeConfig("bert", name_or_path="cross-encoder/ms-marco-MiniLM-L6-v2") + with ( + patch(_INFER, return_value="text-classification"), + patch(_GET_PIPELINE_TAG, return_value="text-ranking"), + ): + r = resolve_task(cfg) + assert r.task == "reranking" + assert r.optimum_task == "text-classification" + assert r.source == TaskSource.TASKS_MANAGER + + def test_resolve_task_pipeline_tag_skips_non_exportable_task() -> None: """A pipeline_tag that is not in the model-type's ONNX-exportable set (e.g. a HuggingFace pipeline label like text-to-image with no export path) is rejected and diff --git a/tests/unit/loader/test_resolve_task_and_model_class.py b/tests/unit/loader/test_resolve_task_and_model_class.py index 48454682d..84a4e60c7 100644 --- a/tests/unit/loader/test_resolve_task_and_model_class.py +++ b/tests/unit/loader/test_resolve_task_and_model_class.py @@ -118,6 +118,23 @@ def test_task_normalized_with_model_class(self): assert r.task == "fill-mask" assert r.source == TaskSource.USER_CLASS + def test_reranking_preserved_with_model_class(self): + """Explicit reranking should stay surfaced even though export uses classification.""" + config = MagicMock() + config.model_type = "bert" + config.architectures = ["BertForSequenceClassification"] + config._name_or_path = "" + + r = resolve_task( + config, + task="reranking", + model_class="AutoModelForSequenceClassification", + ) + + assert r.task == "reranking" + assert r.optimum_task == "text-classification" + assert r.source == TaskSource.USER_CLASS + class TestUnderscoreModelTypePassedToTasksManager: """Regression: model_type with underscores must reach TasksManager un-normalized. diff --git a/tests/unit/loader/test_task_boundary.py b/tests/unit/loader/test_task_boundary.py index df183cfe0..6dee243c7 100644 --- a/tests/unit/loader/test_task_boundary.py +++ b/tests/unit/loader/test_task_boundary.py @@ -22,6 +22,8 @@ [ # Optimum collapses modality (image-feature-extraction -> feature-extraction). ("image-feature-extraction", "feature-extraction"), + # WinML canonical reranking still exports through sequence classification. + ("reranking", "text-classification"), # WinML extension: routed to its Optimum-canonical target. ("next-sentence-prediction", "text-classification"), # WinML extension preserved as-is (Optimum would mis-map it otherwise).