|
| 1 | +# Copyright (c) Microsoft Corporation. |
| 2 | +# Licensed under the MIT license. |
| 3 | + |
| 4 | +import logging |
| 5 | +from enum import Enum |
| 6 | + |
| 7 | +from pyrit.datasets.seed_datasets.remote.remote_dataset_loader import ( |
| 8 | + _RemoteDatasetLoader, |
| 9 | +) |
| 10 | +from pyrit.models import SeedDataset, SeedObjective |
| 11 | + |
| 12 | +logger = logging.getLogger(__name__) |
| 13 | + |
| 14 | + |
| 15 | +_AUTHORS: list[str] = [ |
| 16 | + "Faeze Brahman", |
| 17 | + "Sachin Kumar", |
| 18 | + "Vidhisha Balachandran", |
| 19 | + "Pradeep Dasigi", |
| 20 | + "Valentina Pyatkin", |
| 21 | + "Abhilasha Ravichander", |
| 22 | + "Sarah Wiegreffe", |
| 23 | + "Nouha Dziri", |
| 24 | + "Khyathi Chandu", |
| 25 | + "Jack Hessel", |
| 26 | + "Yulia Tsvetkov", |
| 27 | + "Noah A. Smith", |
| 28 | + "Yejin Choi", |
| 29 | + "Hannaneh Hajishirzi", |
| 30 | +] |
| 31 | + |
| 32 | +_GROUPS: list[str] = ["Allen Institute for AI"] |
| 33 | + |
| 34 | + |
| 35 | +class CoCoNotCategory(Enum): |
| 36 | + """ |
| 37 | + The 5 top-level noncompliance categories defined in the CoCoNot taxonomy. |
| 38 | +
|
| 39 | + Values match the casing used by the upstream HuggingFace dataset so they can be |
| 40 | + used as direct row-filter keys. |
| 41 | + """ |
| 42 | + |
| 43 | + INCOMPLETE = "Incomplete requests" |
| 44 | + UNSUPPORTED = "Unsupported requests" |
| 45 | + INDETERMINATE = "Indeterminate requests" |
| 46 | + HUMANIZING = "Humanizing requests" |
| 47 | + SAFETY = "Requests with safety concerns" |
| 48 | + |
| 49 | + |
| 50 | +class CoCoNotSplit(Enum): |
| 51 | + """Splits available for the upstream ``original`` config.""" |
| 52 | + |
| 53 | + TRAIN = "train" |
| 54 | + TEST = "test" |
| 55 | + |
| 56 | + |
| 57 | +class _CoCoNotBaseDataset(_RemoteDatasetLoader): |
| 58 | + """ |
| 59 | + Shared base for the two CoCoNot sibling loaders. |
| 60 | +
|
| 61 | + CoCoNot (Contextual Noncompliance) is an evaluation suite for refusal calibration in |
| 62 | + LLMs. The dataset is split across two configs that this base wraps with sibling |
| 63 | + subclasses: |
| 64 | +
|
| 65 | + - ``original`` (train+test): prompts the model SHOULD refuse, drawn from 5 |
| 66 | + noncompliance categories (incomplete, unsupported, indeterminate, humanizing, |
| 67 | + safety). |
| 68 | + - ``contrast`` (test only): benign look-alike prompts the model SHOULD comply with, |
| 69 | + used to measure over-refusal behavior. |
| 70 | +
|
| 71 | + Subclasses set ``CONFIG``, ``SPLITS``, ``DEFAULT_DESCRIPTION``, ``size``, and a |
| 72 | + ``dataset_name`` property. |
| 73 | +
|
| 74 | + References: |
| 75 | + - https://huggingface.co/datasets/allenai/coconot |
| 76 | + - https://github.com/allenai/noncompliance |
| 77 | + - [@brahman2024coconot] |
| 78 | +
|
| 79 | + License: ODC-BY 1.0. |
| 80 | + """ |
| 81 | + |
| 82 | + HF_DATASET_NAME: str = "allenai/coconot" |
| 83 | + |
| 84 | + CONFIG: str |
| 85 | + SPLITS: tuple[str, ...] |
| 86 | + DEFAULT_DESCRIPTION: str |
| 87 | + |
| 88 | + harm_categories: list[str] = [m.value.lower() for m in CoCoNotCategory] |
| 89 | + modalities: list[str] = ["text"] |
| 90 | + tags: set[str] = {"safety"} |
| 91 | + |
| 92 | + def __init__(self, *, categories: list[CoCoNotCategory] | None = None) -> None: |
| 93 | + """ |
| 94 | + Initialize the CoCoNot base loader. |
| 95 | +
|
| 96 | + Args: |
| 97 | + categories (list[CoCoNotCategory] | None): Subset of noncompliance categories |
| 98 | + to include. ``None`` (default) loads all 5 categories. |
| 99 | +
|
| 100 | + Raises: |
| 101 | + ValueError: If any value in ``categories`` is not a CoCoNotCategory. |
| 102 | + """ |
| 103 | + if categories is not None: |
| 104 | + self._validate_enums(values=categories, enum_cls=CoCoNotCategory, label="categories") |
| 105 | + self._categories = categories |
| 106 | + |
| 107 | + def _resolved_splits(self) -> tuple[str, ...]: |
| 108 | + """ |
| 109 | + Return the splits to iterate when fetching. |
| 110 | +
|
| 111 | + Subclasses with a user-facing splits filter override this to read the filter. |
| 112 | + Base implementation returns the class-level ``SPLITS`` tuple unchanged. |
| 113 | +
|
| 114 | + Returns: |
| 115 | + tuple[str, ...]: The split names to load. |
| 116 | + """ |
| 117 | + return self.SPLITS |
| 118 | + |
| 119 | + async def fetch_dataset_async(self, *, cache: bool = True) -> SeedDataset: |
| 120 | + """ |
| 121 | + Fetch the CoCoNot subset and return it as a SeedDataset. |
| 122 | +
|
| 123 | + Iterates ``self._resolved_splits()`` and calls the inherited |
| 124 | + ``_fetch_from_huggingface`` once per split, then filters by |
| 125 | + ``self._categories`` if set. |
| 126 | +
|
| 127 | + Args: |
| 128 | + cache (bool): Whether to cache the fetched dataset. Defaults to True. |
| 129 | +
|
| 130 | + Returns: |
| 131 | + SeedDataset: SeedDataset of SeedObjectives. Each seed carries per-row |
| 132 | + metadata: ``id``, ``category``, ``subcategory``, ``subset`` (HF config name), |
| 133 | + ``split``, and ``response`` (populated only on ``original.train`` rows). |
| 134 | +
|
| 135 | + Raises: |
| 136 | + ValueError: If no rows match the category filter. |
| 137 | + """ |
| 138 | + wanted_categories = {c.value for c in self._categories} if self._categories else None |
| 139 | + source_url = f"https://huggingface.co/datasets/{self.HF_DATASET_NAME}" |
| 140 | + seeds: list[SeedObjective] = [] |
| 141 | + |
| 142 | + for split in self._resolved_splits(): |
| 143 | + logger.info(f"Loading CoCoNot rows (config={self.CONFIG}, split={split})") |
| 144 | + rows = await self._fetch_from_huggingface( |
| 145 | + dataset_name=self.HF_DATASET_NAME, |
| 146 | + config=self.CONFIG, |
| 147 | + split=split, |
| 148 | + cache=cache, |
| 149 | + ) |
| 150 | + for row in rows: |
| 151 | + category = row.get("category") |
| 152 | + if wanted_categories is not None and category not in wanted_categories: |
| 153 | + continue |
| 154 | + seeds.append(self._row_to_seed(row=row, split=split, source_url=source_url)) |
| 155 | + |
| 156 | + if not seeds: |
| 157 | + raise ValueError("SeedDataset cannot be empty. Check your filter criteria.") |
| 158 | + |
| 159 | + logger.info(f"Successfully loaded {len(seeds)} objectives from CoCoNot ({self.dataset_name})") |
| 160 | + return SeedDataset(seeds=seeds, dataset_name=self.dataset_name) |
| 161 | + |
| 162 | + def _row_to_seed(self, *, row: dict, split: str, source_url: str) -> SeedObjective: |
| 163 | + """ |
| 164 | + Convert one HF row into a SeedObjective with full per-row metadata. |
| 165 | +
|
| 166 | + Args: |
| 167 | + row (dict): One row from the HuggingFace dataset. |
| 168 | + split (str): The split this row came from (used as ``metadata["split"]``). |
| 169 | + source_url (str): Canonical source URL for the dataset. |
| 170 | +
|
| 171 | + Returns: |
| 172 | + SeedObjective: The constructed seed. |
| 173 | + """ |
| 174 | + category = row.get("category") or "" |
| 175 | + metadata: dict[str, str | int] = { |
| 176 | + "id": row.get("id", ""), |
| 177 | + "category": category, |
| 178 | + "subcategory": row.get("subcategory", ""), |
| 179 | + "subset": self.CONFIG, |
| 180 | + "split": split, |
| 181 | + } |
| 182 | + # response is populated in original.train and empty in test splits. |
| 183 | + response = row.get("response") |
| 184 | + if response: |
| 185 | + metadata["response"] = response |
| 186 | + |
| 187 | + return SeedObjective( |
| 188 | + value=row["prompt"], |
| 189 | + dataset_name=self.dataset_name, |
| 190 | + harm_categories=[category] if category else [], |
| 191 | + description=self.DEFAULT_DESCRIPTION, |
| 192 | + source=source_url, |
| 193 | + authors=_AUTHORS, |
| 194 | + groups=_GROUPS, |
| 195 | + metadata=metadata, |
| 196 | + ) |
| 197 | + |
| 198 | + |
| 199 | +class _CoCoNotRefusalDataset(_CoCoNotBaseDataset): |
| 200 | + """ |
| 201 | + 12,478 prompts (train+test) the model SHOULD refuse. |
| 202 | +
|
| 203 | + Maps to the ``original`` config of ``allenai/coconot``. Combines the ``train`` split |
| 204 | + (11,477 rows, each carrying AI2's reference noncompliant response in |
| 205 | + ``metadata["response"]``) and the ``test`` split (1,001 rows, no reference response). |
| 206 | +
|
| 207 | + Use the ``splits`` constructor argument to restrict to one split. |
| 208 | +
|
| 209 | + Reference: [@brahman2024coconot] |
| 210 | + """ |
| 211 | + |
| 212 | + CONFIG: str = "original" |
| 213 | + SPLITS: tuple[str, ...] = ("train", "test") |
| 214 | + size: str = "huge" |
| 215 | + DEFAULT_DESCRIPTION: str = ( |
| 216 | + "CoCoNot refusal-target set — 12,478 prompts the model should NOT comply with, " |
| 217 | + "drawn from 5 noncompliance categories (incomplete, unsupported, indeterminate, " |
| 218 | + "humanizing, safety). Combines the `original.train` (11,477) and `original.test` " |
| 219 | + "(1,001) splits of `allenai/coconot`." |
| 220 | + ) |
| 221 | + |
| 222 | + def __init__( |
| 223 | + self, |
| 224 | + *, |
| 225 | + categories: list[CoCoNotCategory] | None = None, |
| 226 | + splits: list[CoCoNotSplit] | None = None, |
| 227 | + ) -> None: |
| 228 | + """ |
| 229 | + Initialize the CoCoNot refusal-target loader. |
| 230 | +
|
| 231 | + Args: |
| 232 | + categories (list[CoCoNotCategory] | None): Subset of noncompliance categories |
| 233 | + to include. ``None`` (default) loads all 5 categories. |
| 234 | + splits (list[CoCoNotSplit] | None): Subset of upstream splits to load. ``None`` |
| 235 | + (default) loads both ``train`` (11,477 rows) and ``test`` (1,001 rows). |
| 236 | +
|
| 237 | + Raises: |
| 238 | + ValueError: If any value in ``categories`` or ``splits`` is the wrong enum |
| 239 | + type. |
| 240 | + """ |
| 241 | + super().__init__(categories=categories) |
| 242 | + if splits is not None: |
| 243 | + self._validate_enums(values=splits, enum_cls=CoCoNotSplit, label="splits") |
| 244 | + self._splits = splits |
| 245 | + |
| 246 | + def _resolved_splits(self) -> tuple[str, ...]: |
| 247 | + """ |
| 248 | + Return the splits to load, honoring the user-supplied ``splits`` filter. |
| 249 | +
|
| 250 | + Returns: |
| 251 | + tuple[str, ...]: ``self.SPLITS`` when no filter is set, otherwise the |
| 252 | + user-selected subset in the order they passed. |
| 253 | + """ |
| 254 | + if self._splits is None: |
| 255 | + return self.SPLITS |
| 256 | + return tuple(s.value for s in self._splits) |
| 257 | + |
| 258 | + @property |
| 259 | + def dataset_name(self) -> str: |
| 260 | + """Return the dataset name.""" |
| 261 | + return "coconot_refusal" |
| 262 | + |
| 263 | + |
| 264 | +class _CoCoNotContrastDataset(_CoCoNotBaseDataset): |
| 265 | + """ |
| 266 | + 379 look-alike benign prompts the model SHOULD comply with. |
| 267 | +
|
| 268 | + Maps to the ``contrast.test`` config/split of ``allenai/coconot``. Used to measure |
| 269 | + over-refusal behavior: each prompt is superficially similar to a refusal-target |
| 270 | + prompt but is in fact benign. |
| 271 | +
|
| 272 | + Reference: [@brahman2024coconot] |
| 273 | + """ |
| 274 | + |
| 275 | + CONFIG: str = "contrast" |
| 276 | + SPLITS: tuple[str, ...] = ("test",) |
| 277 | + size: str = "medium" |
| 278 | + tags: set[str] = set() |
| 279 | + DEFAULT_DESCRIPTION: str = ( |
| 280 | + "CoCoNot contrast set — 379 benign prompts that look superficially similar to " |
| 281 | + "refusal-target prompts but should be complied with. Used to measure " |
| 282 | + "over-refusal behavior. Maps to the `contrast.test` config of " |
| 283 | + "`allenai/coconot`." |
| 284 | + ) |
| 285 | + |
| 286 | + @property |
| 287 | + def dataset_name(self) -> str: |
| 288 | + """Return the dataset name.""" |
| 289 | + return "coconot_contrast" |
0 commit comments