diff --git a/src/winml/modelkit/datasets/__init__.py b/src/winml/modelkit/datasets/__init__.py index 66bee45de..b8f9f231a 100644 --- a/src/winml/modelkit/datasets/__init__.py +++ b/src/winml/modelkit/datasets/__init__.py @@ -26,6 +26,7 @@ from .processor_utils import get_image_processor_config from .random_dataset import RandomDataset from .text import TextDataset +from .zero_shot_classification import ZeroShotClassificationDataset if TYPE_CHECKING: @@ -47,7 +48,7 @@ "sentence-similarity": TextDataset, "next-sentence-prediction": TextDataset, "fill-mask": TextDataset, - "zero-shot-classification": TextDataset, + "zero-shot-classification": ZeroShotClassificationDataset, "image-segmentation": ImageSegmentationDataset, "mask-generation": MaskGenerationDataset, "depth-estimation": DepthEstimationDataset, @@ -337,6 +338,7 @@ def __len__(self) -> int: "ObjectDetectionDataset", "RandomDataset", "TextDataset", + "ZeroShotClassificationDataset", # Utilities "format_data", "get_image_processor_config", diff --git a/src/winml/modelkit/datasets/zero_shot_classification.py b/src/winml/modelkit/datasets/zero_shot_classification.py new file mode 100644 index 000000000..3991b7f60 --- /dev/null +++ b/src/winml/modelkit/datasets/zero_shot_classification.py @@ -0,0 +1,203 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- +"""Zero-shot classification calibration dataset. + +Zero-shot text classifiers consume natural-language premise/hypothesis pairs. +This specialization keeps calibration aligned with that input distribution +instead of reusing generic text-classification defaults. +""" + +from __future__ import annotations + +import logging +from math import ceil +from numbers import Integral +from random import Random +from typing import TYPE_CHECKING, Any, cast + +from datasets import load_dataset +from transformers import AutoTokenizer + +from .base import BaseTaskDataset + + +if TYPE_CHECKING: + from collections.abc import Sequence + + +logger = logging.getLogger(__name__) + +DEFAULT_ZERO_SHOT_CLASSIFICATION_DATASET = "fancyzhx/ag_news" +DEFAULT_ZERO_SHOT_CLASSIFICATION_SPLIT = "test" +DEFAULT_ZERO_SHOT_TEXT_COLUMN = "text" +DEFAULT_ZERO_SHOT_CANDIDATE_LABELS = ("World", "Sports", "Business", "Sci/Tech") +DEFAULT_ZERO_SHOT_HYPOTHESIS_TEMPLATE = "This text is about {}." + + +class ZeroShotClassificationDataset(BaseTaskDataset): + """Calibration dataset for zero-shot text classification models.""" + + DEFAULT_SEQ_LEN = 128 + + def __init__( + self, + model_name: str, + dataset_name: str | None = None, + max_samples: int | None = None, + data_split: str | None = None, + *, + candidate_labels: str | Sequence[str] | None = None, + hypothesis_template: str | None = None, + input_column: str | None = None, + max_length: int | None = None, + io_config: dict | None = None, + io_mapping: dict | None = None, + **kwargs: Any, + ) -> None: + self._candidate_labels = self._normalize_candidate_labels(candidate_labels) + self._hypothesis_template = hypothesis_template or DEFAULT_ZERO_SHOT_HYPOTHESIS_TEMPLATE + self._input_column = input_column or DEFAULT_ZERO_SHOT_TEXT_COLUMN + self._max_length = max_length + self._io_config = io_config + self._io_mapping = io_mapping or {} + + super().__init__( + model_name=model_name, + dataset_name=dataset_name, + max_samples=max_samples, + data_split=data_split, + io_config=io_config, + io_mapping=io_mapping, + **kwargs, + ) + + @staticmethod + def _normalize_candidate_labels(candidate_labels: str | Sequence[str] | None) -> list[str]: + """Return non-empty labels from comma-separated or sequence input.""" + if candidate_labels is None: + raw_labels: Sequence[str] = DEFAULT_ZERO_SHOT_CANDIDATE_LABELS + elif isinstance(candidate_labels, str): + raw_labels = candidate_labels.split(",") + else: + raw_labels = candidate_labels + + labels = [str(label).strip() for label in raw_labels if str(label).strip()] + if not labels: + raise ValueError("candidate_labels must contain at least one label") + return labels + + def _get_default_dataset(self) -> None: + """Set the built-in zero-shot dataset defaults.""" + if self._dataset_name is None: + self._dataset_name = DEFAULT_ZERO_SHOT_CLASSIFICATION_DATASET + if self._data_split is None: + self._data_split = DEFAULT_ZERO_SHOT_CLASSIFICATION_SPLIT + + def _resolve_max_length(self) -> None: + """Resolve tokenizer max length from ONNX io_config when available.""" + if self._max_length is None: + self._max_length = self.DEFAULT_SEQ_LEN + + if not self._io_config: + return + + onnx_name = self._io_mapping.get("input_ids", "input_ids") + input_config = self._io_config.get(onnx_name) + if not input_config: + return + + shape = input_config.get("shape", []) + if len(shape) > 1 and isinstance(shape[1], Integral): + self._max_length = int(shape[1]) + logger.info("max_length=%d from io_config[%s]", self._max_length, onnx_name) + + def _source_sample_count(self) -> int | None: + """Number of source rows needed to produce max_samples calibration pairs.""" + if self._max_samples is None: + return None + return ceil(self._max_samples / len(self._candidate_labels)) + + def _load_and_sample(self) -> Any: + """Load the source text dataset and cap rows before pair expansion.""" + subset = self._config.get("subset") or self._config.get("dataset_config_name") + load_args = [self._dataset_name] + if subset: + load_args.append(subset) + + logger.info( + "Loading zero-shot calibration dataset: %s (split=%s)", + load_args, + self._data_split, + ) + dataset = load_dataset(*load_args, split=self._data_split) + + source_count = self._source_sample_count() + shuffle = self._config.get("shuffle", False) + seed = self._config.get("seed", 42) + + if source_count is not None: + n = min(source_count, len(dataset)) + indices = Random(seed).sample(range(len(dataset)), n) if shuffle else list(range(n)) + dataset = dataset.select(indices) + elif shuffle: + dataset = dataset.shuffle(seed=seed) + + return dataset + + def _apply_io_mapping(self, sample: dict[str, Any]) -> dict[str, Any]: + """Rename tokenizer output fields to ONNX input names when requested.""" + if not self._io_mapping: + return sample + return {self._io_mapping.get(key, key): value for key, value in sample.items()} + + def _initialize(self) -> None: + """Load text rows and tokenize them as zero-shot premise/hypothesis pairs.""" + self._get_default_dataset() + self._resolve_max_length() + dataset = self._load_and_sample() + tokenizer = AutoTokenizer.from_pretrained(self._model_name, use_fast=True) + + samples: list[dict[str, Any]] = [] + for example in dataset: + if self._input_column not in example: + raise ValueError( + f"Column '{self._input_column}' not found in {self._dataset_name}; " + f"available columns: {sorted(example)}" + ) + premise = str(example[self._input_column]) + for label in self._candidate_labels: + hypothesis = self._hypothesis_template.format(label) + tokenized = dict( + tokenizer( + premise, + hypothesis, + padding="max_length", + truncation=True, + max_length=self._max_length, + return_tensors="pt", + ) + ) + samples.append(self._apply_io_mapping(tokenized)) + if self._max_samples is not None and len(samples) >= self._max_samples: + self._dataset = samples + return + + self._dataset = samples + logger.info("Initialized zero-shot calibration dataset with %d pairs", len(samples)) + + def __getitem__(self, idx: int) -> dict[str, Any]: + """Get a tokenized premise/hypothesis calibration sample.""" + return cast("dict[str, Any]", self._dataset[idx]) + + @property + def label_col(self) -> str: + """Zero-shot calibration does not consume dataset labels.""" + return "" + + @property + def max_length(self) -> int: + """Sequence length.""" + assert self._max_length is not None, "max_length not resolved" + return self._max_length diff --git a/tests/integration/datasets/test_zero_shot_classification.py b/tests/integration/datasets/test_zero_shot_classification.py new file mode 100644 index 000000000..b82ad25ea --- /dev/null +++ b/tests/integration/datasets/test_zero_shot_classification.py @@ -0,0 +1,89 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- +"""Tests for zero-shot-classification calibration datasets.""" + +from __future__ import annotations + +from typing import Any +from unittest.mock import MagicMock, patch + +import torch +from datasets.features import ClassLabel, Value + + +class _FakeDataset: + def __init__(self, rows: list[dict[str, Any]]) -> None: + self._rows = rows + self.features = { + "text": Value("string"), + "label": ClassLabel(names=["World", "Sports", "Business", "Sci/Tech"]), + } + + def __len__(self) -> int: + return len(self._rows) + + def __getitem__(self, idx: int) -> dict[str, Any]: + return self._rows[idx] + + def select(self, indices: list[int]) -> _FakeDataset: + return _FakeDataset([self._rows[i] for i in indices]) + + +class _FakeTokenizer: + def __init__(self) -> None: + self.calls: list[tuple[str, str]] = [] + + def __call__(self, premise: str, hypothesis: str, **_: Any) -> dict[str, torch.Tensor]: + self.calls.append((premise, hypothesis)) + return { + "input_ids": torch.ones(1, 8, dtype=torch.int64), + "attention_mask": torch.ones(1, 8, dtype=torch.int64), + } + + +class TestZeroShotClassificationDataset: + @patch("winml.modelkit.datasets.zero_shot_classification.load_dataset") + @patch("winml.modelkit.datasets.zero_shot_classification.AutoTokenizer") + def test_defaults_use_ag_news_sentence_pairs( + self, + mock_tokenizer_cls: MagicMock, + mock_load_dataset: MagicMock, + ) -> None: + from winml.modelkit.datasets import ZeroShotClassificationDataset + from winml.modelkit.datasets.zero_shot_classification import ( + DEFAULT_ZERO_SHOT_CLASSIFICATION_DATASET, + DEFAULT_ZERO_SHOT_CLASSIFICATION_SPLIT, + ) + + tokenizer = _FakeTokenizer() + mock_tokenizer_cls.from_pretrained.return_value = tokenizer + mock_load_dataset.return_value = _FakeDataset( + [{"text": "Markets rallied after the earnings report.", "label": 2}] + ) + + dataset = ZeroShotClassificationDataset( + model_name="cross-encoder/nli-deberta-v3-small", + max_samples=4, + ) + + assert dataset.dataset_name == DEFAULT_ZERO_SHOT_CLASSIFICATION_DATASET + assert dataset.data_split == DEFAULT_ZERO_SHOT_CLASSIFICATION_SPLIT + mock_load_dataset.assert_called_once_with( + DEFAULT_ZERO_SHOT_CLASSIFICATION_DATASET, + split=DEFAULT_ZERO_SHOT_CLASSIFICATION_SPLIT, + ) + assert tokenizer.calls == [ + ("Markets rallied after the earnings report.", "This text is about World."), + ("Markets rallied after the earnings report.", "This text is about Sports."), + ("Markets rallied after the earnings report.", "This text is about Business."), + ("Markets rallied after the earnings report.", "This text is about Sci/Tech."), + ] + assert len(dataset) == 4 + assert set(dataset[0]) == {"input_ids", "attention_mask"} + + def test_zero_shot_task_uses_specialized_dataset(self) -> None: + from winml.modelkit.datasets import TASK_DATASET_MAPPING, ZeroShotClassificationDataset + + assert TASK_DATASET_MAPPING["zero-shot-classification"] is ZeroShotClassificationDataset