|
| 1 | +# ------------------------------------------------------------------------- |
| 2 | +# Copyright (c) Microsoft Corporation. All rights reserved. |
| 3 | +# Licensed under the MIT License. |
| 4 | +# -------------------------------------------------------------------------- |
| 5 | + |
| 6 | +"""Zero-shot classification evaluator for NLI checkpoints. |
| 7 | +
|
| 8 | +Computes accuracy and macro-F1 via ClassificationMetric. |
| 9 | +HF evaluate library has no zero-shot-classification evaluator, so this |
| 10 | +class runs the metric loop manually: each text is scored against every |
| 11 | +candidate label as a (premise, hypothesis) NLI pair, and the label with |
| 12 | +the top entailment score wins. |
| 13 | +
|
| 14 | +Candidate labels come from ``columns_mapping["candidate_labels"]`` if set |
| 15 | +(comma-separated), otherwise from ``dataset.features[label_column].names`` |
| 16 | +when the column is a ``ClassLabel``. |
| 17 | +""" |
| 18 | + |
| 19 | +from __future__ import annotations |
| 20 | + |
| 21 | +from typing import TYPE_CHECKING, Any |
| 22 | + |
| 23 | +from transformers.pipelines.zero_shot_classification import ZeroShotClassificationPipeline |
| 24 | + |
| 25 | +from .base_evaluator import WinMLEvaluator |
| 26 | + |
| 27 | + |
| 28 | +if TYPE_CHECKING: |
| 29 | + from datasets import Dataset |
| 30 | + from transformers.pipelines.base import Pipeline |
| 31 | + |
| 32 | + from ..datasets.config import DatasetConfig |
| 33 | + from ..models.winml.base import WinMLPreTrainedModel |
| 34 | + from .config import WinMLEvaluationConfig |
| 35 | + |
| 36 | + |
| 37 | +class _FixedShapeZeroShotPipeline(ZeroShotClassificationPipeline): |
| 38 | + """Pad to ``tokenizer.model_max_length`` for fixed-shape ONNX exports.""" |
| 39 | + |
| 40 | + def _parse_and_tokenize(self, sequence_pairs: Any, **kwargs: Any) -> Any: |
| 41 | + kwargs["padding"] = "max_length" |
| 42 | + kwargs.setdefault("truncation", True) |
| 43 | + return super()._parse_and_tokenize(sequence_pairs, **kwargs) |
| 44 | + |
| 45 | + |
| 46 | +class WinMLZeroShotClassificationEvaluator(WinMLEvaluator): |
| 47 | + """Evaluator for zero-shot text classification using NLI models.""" |
| 48 | + |
| 49 | + @classmethod |
| 50 | + def schema_info(cls) -> list: |
| 51 | + """Return expected dataset schema for zero-shot classification.""" |
| 52 | + from .config import SchemaColumn |
| 53 | + |
| 54 | + return [ |
| 55 | + SchemaColumn("text", "Value(string)", "input_column", description="input text"), |
| 56 | + SchemaColumn( |
| 57 | + "label", |
| 58 | + "ClassLabel", |
| 59 | + "label_column", |
| 60 | + description="gold label (ClassLabel or string)", |
| 61 | + ), |
| 62 | + SchemaColumn( |
| 63 | + "<candidate_labels>", |
| 64 | + "comma-separated str", |
| 65 | + "candidate_labels", |
| 66 | + required=False, |
| 67 | + description="override candidate labels (required for non-ClassLabel columns)", |
| 68 | + ), |
| 69 | + SchemaColumn( |
| 70 | + "<hypothesis_template>", |
| 71 | + "Value(string)", |
| 72 | + "hypothesis_template", |
| 73 | + required=False, |
| 74 | + description='NLI hypothesis template (default: "This example is {}.")', |
| 75 | + ), |
| 76 | + ] |
| 77 | + |
| 78 | + def __init__( |
| 79 | + self, |
| 80 | + config: WinMLEvaluationConfig, |
| 81 | + model: WinMLPreTrainedModel, |
| 82 | + ) -> None: |
| 83 | + mapping = config.dataset.columns_mapping |
| 84 | + self._input_col = mapping.get("input_column", "text") |
| 85 | + self._label_col = mapping.get("label_column", "label") |
| 86 | + self._candidate_labels_override = mapping.get("candidate_labels") |
| 87 | + self._hypothesis_template = mapping.get("hypothesis_template") |
| 88 | + super().__init__(config, model) |
| 89 | + |
| 90 | + def prepare_pipeline(self) -> Pipeline: |
| 91 | + """Create pipeline with fixed-length tokenization for ONNX.""" |
| 92 | + from transformers import pipeline |
| 93 | + |
| 94 | + io_config = getattr(self.model, "io_config", None) or {} |
| 95 | + shapes = io_config.get("input_shapes", [[]]) |
| 96 | + max_length: int | None = None |
| 97 | + if shapes and len(shapes[0]) > 1 and isinstance(shapes[0][1], int): |
| 98 | + max_length = shapes[0][1] |
| 99 | + |
| 100 | + pipe = pipeline( |
| 101 | + "zero-shot-classification", |
| 102 | + model=self.model, |
| 103 | + framework="pt", |
| 104 | + tokenizer=self.config.model_id, |
| 105 | + device="cpu", |
| 106 | + pipeline_class=_FixedShapeZeroShotPipeline, |
| 107 | + ) |
| 108 | + |
| 109 | + if pipe.tokenizer is not None and max_length is not None: |
| 110 | + pipe.tokenizer.model_max_length = max_length |
| 111 | + |
| 112 | + # Drop tokenizer keys the ONNX graph does not accept |
| 113 | + # (e.g. DeBERTa-v3 MNLI exports omit token_type_ids). |
| 114 | + input_names = io_config.get("input_names", []) |
| 115 | + if input_names: |
| 116 | + filtered = [n for n in pipe.tokenizer.model_input_names if n in input_names] |
| 117 | + if filtered: |
| 118 | + pipe.tokenizer.model_input_names = filtered |
| 119 | + |
| 120 | + return pipe |
| 121 | + |
| 122 | + def align_labels( |
| 123 | + self, |
| 124 | + dataset: Dataset, |
| 125 | + ds_config: DatasetConfig, |
| 126 | + ) -> Dataset: |
| 127 | + """Validate input and label columns. |
| 128 | +
|
| 129 | + Base-class label alignment is bypassed: NLI ``label2id`` identifies |
| 130 | + entailment/neutral/contradiction classes, which are unrelated to the |
| 131 | + ground-truth labels used for accuracy. |
| 132 | + """ |
| 133 | + col_names = set(dataset.column_names) |
| 134 | + for col in (self._input_col, self._label_col): |
| 135 | + if col not in col_names: |
| 136 | + raise ValueError(f"Column '{col}' not found in dataset: {sorted(col_names)}") |
| 137 | + return dataset |
| 138 | + |
| 139 | + def _resolve_candidate_labels(self, dataset: Dataset) -> list[str]: |
| 140 | + """Return candidate labels from user override or dataset ``ClassLabel``.""" |
| 141 | + if self._candidate_labels_override: |
| 142 | + labels = [s.strip() for s in self._candidate_labels_override.split(",") if s.strip()] |
| 143 | + if not labels: |
| 144 | + raise ValueError("candidate_labels override must not be empty.") |
| 145 | + return labels |
| 146 | + |
| 147 | + names = getattr(dataset.features.get(self._label_col), "names", None) |
| 148 | + if names: |
| 149 | + return list(names) |
| 150 | + |
| 151 | + raise ValueError( |
| 152 | + f"Column '{self._label_col}' is not a ClassLabel; pass " |
| 153 | + f'--column "candidate_labels=a,b,...".', |
| 154 | + ) |
| 155 | + |
| 156 | + def compute(self) -> dict[str, Any]: |
| 157 | + """Compute accuracy and macro-F1 over all samples.""" |
| 158 | + from .metrics import ClassificationMetric |
| 159 | + |
| 160 | + candidate_labels = self._resolve_candidate_labels(self.data) |
| 161 | + class_names = getattr(self.data.features.get(self._label_col), "names", None) |
| 162 | + |
| 163 | + pipe_kwargs: dict[str, Any] = {"candidate_labels": candidate_labels} |
| 164 | + if self._hypothesis_template is not None: |
| 165 | + pipe_kwargs["hypothesis_template"] = self._hypothesis_template |
| 166 | + |
| 167 | + predictions: list[str] = [] |
| 168 | + references: list[str] = [] |
| 169 | + for sample in self.data: |
| 170 | + result = self.pipe(sample[self._input_col], **pipe_kwargs) |
| 171 | + predictions.append(result["labels"][0]) |
| 172 | + raw = sample[self._label_col] |
| 173 | + references.append(class_names[int(raw)] if class_names else str(raw)) |
| 174 | + |
| 175 | + return ClassificationMetric().compute(predictions, references, candidate_labels) |
0 commit comments