FEAT Overlay harm-category standardization across dataset loaders (microsoft#2158)

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Author: romanlutz

pyrit/datasets/seed_datasets/remote/aegis_ai_content_safety_dataset.py

@@ -10,6 +10,7 @@
     _RemoteDatasetLoader,
 )
 from pyrit.models import Modality, SeedDataset, SeedPrompt, SeedUnion
+from pyrit.models.harm_category import HarmCategory
 
 logger = logging.getLogger(__name__)
 
@@ -142,6 +143,49 @@ async def fetch_dataset_async(self, *, cache: bool = True) -> SeedDataset:
             cache=cache,
         )
 
+        # Map AEGIS-specific categories to PyRIT harm categories
+        alias_overrides: dict[str, list[HarmCategory]] = {
+            "Controlled/Regulated Substances": [HarmCategory.DRUG_USE, HarmCategory.REGULATED_GOODS],
+            "Copyright/Trademark/Plagiarism": [
+                HarmCategory.COPYRIGHT,
+                HarmCategory.TRADEMARK,
+                HarmCategory.PLAGIARISM,
+            ],
+            "Criminal Planning/Confessions": [HarmCategory.COORDINATION_HARM],
+            "Fraud/Deception": [HarmCategory.SCAMS, HarmCategory.DECEPTION],
+            "Guns and Illegal Weapons": [
+                HarmCategory.REGULATED_GOODS,
+                HarmCategory.COORDINATION_HARM,
+                HarmCategory.VIOLENT_CONTENT,
+            ],
+            "Hate/Identity Hate": [HarmCategory.HATE_SPEECH, HarmCategory.REPRESENTATIONAL],
+            "High Risk Gov Decision Making": [HarmCategory.HIGH_RISK_GOVERNMENT],
+            "Illegal Activity": [HarmCategory.COORDINATION_HARM],
+            "Immoral/Unethical": [HarmCategory.OTHER],
+            "Manipulation": [HarmCategory.DECEPTION],
+            "Needs Caution": [HarmCategory.OTHER],
+            "PII/Privacy": [HarmCategory.PPI],
+            "Political/Misinformation/Conspiracy": [
+                HarmCategory.INFO_INTEGRITY,
+                HarmCategory.CURRENT_EVENTS_MISINFO,
+                HarmCategory.CAMPAIGNING,
+            ],
+            "Sexual": [HarmCategory.SEXUAL_CONTENT],
+            "Sexual (minor)": [HarmCategory.SEXUALIZATION, HarmCategory.CHILD_LEAKAGE, HarmCategory.SEXUAL_CONTENT],
+            "Suicide and Self Harm": [HarmCategory.SUICIDE, HarmCategory.SELF_HARM],
+            "Threat": [HarmCategory.VIOLENT_THREATS],
+            "Unauthorized Advice": [
+                HarmCategory.FINANCIAL_ADVICE,
+                HarmCategory.HEALTH_DIAGNOSIS,
+                HarmCategory.LEGAL_ADVICE,
+            ],
+            "Violence": [
+                HarmCategory.VIOLENT_CONTENT,
+                HarmCategory.VIOLENT_THREATS,
+                HarmCategory.COORDINATION_HARM,
+            ],
+        }
+
         seed_prompts: list[SeedUnion] = []
 
         for split_name in hf_dataset:
@@ -162,6 +206,10 @@ async def fetch_dataset_async(self, *, cache: bool = True) -> SeedDataset:
                     if violated_categories
                     else []
                 )
+                standardized_categories = self._standardize_harm_categories(
+                    prompt_harm_categories,
+                    alias_overrides=alias_overrides,
+                )
 
                 # Filter by harm_categories if specified
                 if self._selected_category_values is not None and not any(
@@ -174,7 +222,7 @@ async def fetch_dataset_async(self, *, cache: bool = True) -> SeedDataset:
                         value=prompt_value,
                         data_type="text",
                         dataset_name=self.dataset_name,
-                        harm_categories=prompt_harm_categories if prompt_harm_categories else None,
+                        harm_categories=standardized_categories if standardized_categories else None,
                         source=self.source,
                         authors=self._AUTHORS,
                         groups=self._GROUPS,
@@ -184,6 +232,7 @@ async def fetch_dataset_async(self, *, cache: bool = True) -> SeedDataset:
                             "response_label": example.get("response_label"),
                             "prompt_label_source": example.get("prompt_label_source"),
                             "response_label_source": example.get("response_label_source"),
+                            "aegis_violated_categories": ", ".join(prompt_harm_categories),
                         },
                     )
                 )

pyrit/datasets/seed_datasets/remote/agent_threat_rules_dataset.py

@@ -111,21 +111,23 @@ class _AgentThreatRulesDataset(_RemoteDatasetLoader):
     upstream metadata fields (``original_rule_id``, ``technique``,
     ``detection_field``, ``variation_type``) are preserved on
     ``SeedPrompt.metadata`` so downstream consumers can route, filter, or
-    score by them. ``harm_categories`` is set to the rule's ATR category
-    (single-element list).
+    score by them. ATR categories are agent-security technique labels rather
+    than content-harm labels, so ``harm_categories`` is intentionally empty.
 
     The optional ``categories``, ``techniques``, ``detection_fields``, and
     ``variation_types`` arguments narrow the dataset client-side after fetch.
     Passing an empty list is rejected — pass ``None`` to disable a filter.
     """
 
-    # Class-attribute metadata picked up by SeedDatasetMetadata. Derived from
-    # _RULE_ID_TO_CATEGORY so the categories the loader claims to cover always
-    # match the categories it actually produces — same single-source-of-truth
-    # rationale as the enum-typed dict above.
-    harm_categories: list[str] = sorted({c.value for c in _RULE_ID_TO_CATEGORY.values()})
+    # ATR categories are agent-security technique labels, not content-harm
+    # labels, so this loader intentionally leaves harm categories empty.
+    harm_categories: list[str] = []
     modalities: list[str] = ["text"]
     size: str = "large"  # 1,054 seeds
+    # ATR's upstream corpus grows over time, but this loader pins to a specific commit
+    # by default (see ``source`` below) so a default fetch returns a static, reproducible
+    # snapshot. It is therefore intentionally NOT tagged as a "feed"; unpin the source
+    # (pass ``main`` or a newer ref) to track upstream additions.
     tags: set[str] = {"safety", "agent_security", "prompt_injection"}
 
     def __init__(
@@ -275,6 +277,7 @@ async def fetch_dataset_async(self, *, cache: bool = True) -> SeedDataset:
                 "detection_field": example["detection_field"],
                 "variation_type": example["variation_type"],
                 "atr_id": example["id"],
+                "atr_category": category_value,
             }
 
             # Per-rule description so downstream consumers reading metadata see
@@ -291,7 +294,7 @@ async def fetch_dataset_async(self, *, cache: bool = True) -> SeedDataset:
                     data_type="text",
                     name=rule_id,
                     dataset_name=self.dataset_name,
-                    harm_categories=[category_value],
+                    harm_categories=[],
                     description=description,
                     authors=authors,
                     groups=groups,

pyrit/datasets/seed_datasets/remote/aya_redteaming_dataset.py

@@ -11,6 +11,7 @@
     _RemoteDatasetLoader,
 )
 from pyrit.models import Modality, SeedDataset, SeedPrompt
+from pyrit.models.harm_category import HarmCategory
 
 logger = logging.getLogger(__name__)
 
@@ -123,10 +124,30 @@ async def fetch_dataset_async(self, *, cache: bool = True) -> SeedDataset:
             cache=cache,
         )
 
-        seed_prompts = []
+        harm_category_alias_overrides: dict[str, list[HarmCategory]] = {
+            "bullying & harassment": [HarmCategory.HARASSMENT],
+            "discrimination & injustice": [HarmCategory.HATE_SPEECH, HarmCategory.REPRESENTATIONAL],
+            "graphic material": [HarmCategory.VIOLENT_CONTENT],
+            "harms of representation allocation and quality of service": [
+                HarmCategory.REPRESENTATIONAL,
+                HarmCategory.ALLOCATION,
+                HarmCategory.QUALITY_OF_SERVICE,
+            ],
+            "non-consensual sexual content": [HarmCategory.SEXUAL_CONTENT, HarmCategory.NONCONSENSUAL_UPLOAD],
+            "violence, threats & incitement": [
+                HarmCategory.VIOLENT_CONTENT,
+                HarmCategory.VIOLENT_THREATS,
+                HarmCategory.COORDINATION_HARM,
+            ],
+        }
+        seed_prompts: list[SeedPrompt] = []
 
         for example in examples:
             categories = ast.literal_eval(example["harm_category"])
+            standardized_categories = self._standardize_harm_categories(
+                categories,
+                alias_overrides=harm_category_alias_overrides,
+            )
 
             # Apply filters
             if self.harm_categories_filter is not None and not any(
@@ -142,7 +163,11 @@ async def fetch_dataset_async(self, *, cache: bool = True) -> SeedDataset:
                     value=example["prompt"],
                     data_type="text",
                     dataset_name=self.dataset_name,
-                    harm_categories=categories,
+                    harm_categories=standardized_categories if standardized_categories else None,
+                    metadata={
+                        "aya_redteaming_categories": ", ".join(categories),
+                        "aya_redteaming_scope": example["global_or_local"],
+                    },
                     source="https://huggingface.co/datasets/CohereForAI/aya_redteaming",
                     authors=self._AUTHORS,
                     groups=self._GROUPS,

pyrit/datasets/seed_datasets/remote/babelscape_alert_dataset.py

@@ -10,6 +10,7 @@
     _RemoteDatasetLoader,
 )
 from pyrit.models import Modality, SeedDataset, SeedPrompt
+from pyrit.models.harm_category import HarmCategory
 
 if TYPE_CHECKING:
     from pyrit.models.seeds.seed_group import SeedUnion
@@ -28,6 +29,41 @@ class _BabelscapeAlertDataset(_RemoteDatasetLoader):
     Reference: [@tedeschi2024alert]
     """
 
+    HARM_CATEGORY_ALIAS_OVERRIDES: dict[str, list[HarmCategory]] = {
+        "crime_cyber": [HarmCategory.MALWARE, HarmCategory.COORDINATION_HARM],
+        "crime_injury": [HarmCategory.VIOLENT_CONTENT],
+        "crime_kidnapp": [HarmCategory.VIOLENT_THREATS, HarmCategory.COORDINATION_HARM],
+        "crime_other": [HarmCategory.COORDINATION_HARM],
+        "crime_privacy": [HarmCategory.PPI],
+        "crime_propaganda": [HarmCategory.INFO_INTEGRITY],
+        "crime_tax": [HarmCategory.COORDINATION_HARM],
+        "crime_theft": [HarmCategory.COORDINATION_HARM],
+        "hate_body": [HarmCategory.HATE_SPEECH, HarmCategory.REPRESENTATIONAL],
+        "hate_disabled": [HarmCategory.HATE_SPEECH, HarmCategory.REPRESENTATIONAL],
+        "hate_ethnic": [HarmCategory.HATE_SPEECH, HarmCategory.REPRESENTATIONAL],
+        "hate_lgbtq+": [HarmCategory.HATE_SPEECH, HarmCategory.REPRESENTATIONAL],
+        "hate_other": [HarmCategory.HATE_SPEECH, HarmCategory.REPRESENTATIONAL],
+        "hate_poor": [HarmCategory.HATE_SPEECH, HarmCategory.REPRESENTATIONAL],
+        "hate_religion": [HarmCategory.HATE_SPEECH, HarmCategory.REPRESENTATIONAL],
+        "hate_women": [HarmCategory.HATE_SPEECH, HarmCategory.REPRESENTATIONAL],
+        "self_harm_other": [HarmCategory.SELF_HARM],
+        "self_harm_suicide": [HarmCategory.SUICIDE, HarmCategory.SELF_HARM],
+        "self_harm_thin": [HarmCategory.EATING_DISORDERS, HarmCategory.SELF_HARM],
+        "sex_harrasment": [HarmCategory.HARASSMENT, HarmCategory.SEXUAL_SOLICITATION],
+        "sex_other": [HarmCategory.SEXUAL_CONTENT],
+        "sex_porn": [HarmCategory.SEXUAL_CONTENT],
+        "substance_alcohol": [HarmCategory.DRUG_USE, HarmCategory.REGULATED_GOODS],
+        "substance_cannabis": [HarmCategory.DRUG_USE, HarmCategory.REGULATED_GOODS],
+        "substance_drug": [HarmCategory.DRUG_USE],
+        "substance_other": [HarmCategory.DRUG_USE],
+        "substance_tobacco": [HarmCategory.DRUG_USE, HarmCategory.REGULATED_GOODS],
+        "weapon_biological": [HarmCategory.CBRN],
+        "weapon_chemical": [HarmCategory.CBRN],
+        "weapon_firearm": [HarmCategory.REGULATED_GOODS],
+        "weapon_other": [HarmCategory.REGULATED_GOODS],
+        "weapon_radioactive": [HarmCategory.CBRN],
+    }
+
     _AUTHORS = [
         "Simone Tedeschi",
         "Felix Friedrich",
@@ -112,7 +148,10 @@ async def fetch_dataset_async(self, *, cache: bool = True) -> SeedDataset:
         seed_prompts: list[SeedUnion] = [
             SeedPrompt(
                 value=prompt,
-                harm_categories=[category],
+                harm_categories=self._standardize_harm_categories(
+                    category,
+                    alias_overrides=self.HARM_CATEGORY_ALIAS_OVERRIDES,
+                ),
                 data_type="text",
                 dataset_name=self.dataset_name,
                 description=(
@@ -121,6 +160,7 @@ async def fetch_dataset_async(self, *, cache: bool = True) -> SeedDataset:
                     "red teaming prompts."
                 ),
                 source=f"https://huggingface.co/datasets/{self.source}",
+                metadata={"category": category},
                 authors=self._AUTHORS,
                 groups=self._GROUPS,
             )

pyrit/datasets/seed_datasets/remote/beaver_tails_dataset.py

@@ -1,12 +1,14 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT license.
 
+import json
 import logging
 
 from pyrit.datasets.seed_datasets.remote.remote_dataset_loader import (
     _RemoteDatasetLoader,
 )
 from pyrit.models import Modality, SeedDataset, SeedPrompt
+from pyrit.models.harm_category import HarmCategory
 
 logger = logging.getLogger(__name__)
 
@@ -30,6 +32,34 @@ class _BeaverTailsDataset(_RemoteDatasetLoader):
     """
 
     HF_DATASET_NAME: str = "PKU-Alignment/BeaverTails"
+    HARM_CATEGORY_ALIAS_OVERRIDES: dict[str, list[HarmCategory]] = {
+        "aiding_and_abetting": [HarmCategory.COORDINATION_HARM],
+        "animal_abuse": [HarmCategory.OTHER],
+        "banned_substance": [HarmCategory.DRUG_USE, HarmCategory.REGULATED_GOODS],
+        "child_abuse": [HarmCategory.GROOMING, HarmCategory.SEXUAL_CONTENT, HarmCategory.CHILD_LEAKAGE],
+        "controversial_topics": [HarmCategory.OTHER],
+        "discrimination": [HarmCategory.REPRESENTATIONAL, HarmCategory.HATE_SPEECH],
+        "drug_abuse": [HarmCategory.DRUG_USE],
+        "financial_crime": [HarmCategory.SCAMS, HarmCategory.COORDINATION_HARM],
+        "hate_speech": [HarmCategory.HATE_SPEECH],
+        "incitement": [HarmCategory.VIOLENT_THREATS],
+        "injustice": [HarmCategory.REPRESENTATIONAL, HarmCategory.HATE_SPEECH],
+        "laws_and_safety": [HarmCategory.INFO_INTEGRITY],
+        "misinformation_regarding_ethics": [HarmCategory.INFO_INTEGRITY],
+        "non_violent_unethical_behavior": [HarmCategory.OTHER],
+        "offensive_language": [HarmCategory.HATE_SPEECH],
+        "organized_crime": [HarmCategory.COORDINATION_HARM],
+        "politics": [HarmCategory.OTHER],
+        "privacy_violation": [HarmCategory.PPI],
+        "property_crime": [HarmCategory.COORDINATION_HARM],
+        "self_harm": [HarmCategory.SELF_HARM],
+        "sexually_explicit": [HarmCategory.SEXUAL_CONTENT],
+        "stereotype": [HarmCategory.REPRESENTATIONAL, HarmCategory.HATE_SPEECH],
+        "terrorism": [HarmCategory.VIOLENT_EXTREMISM],
+        "theft": [HarmCategory.COORDINATION_HARM],
+        "violence": [HarmCategory.VIOLENT_CONTENT, HarmCategory.VIOLENT_THREATS, HarmCategory.COORDINATION_HARM],
+        "weapons": [HarmCategory.REGULATED_GOODS],
+    }
 
     # Metadata
     modalities: tuple[Modality, ...] = (Modality.TEXT,)
@@ -102,7 +132,13 @@ async def fetch_dataset_async(self, *, cache: bool = True) -> SeedDataset:
             if self.unsafe_only and item["is_safe"]:
                 continue
 
-            harm_categories = [k for k, v in item["category"].items() if v]
+            raw_harm_categories = [
+                part.strip() for k, v in item["category"].items() if v for part in k.split(",") if part.strip()
+            ]
+            harm_categories = self._standardize_harm_categories(
+                raw_harm_categories,
+                alias_overrides=self.HARM_CATEGORY_ALIAS_OVERRIDES,
+            )
 
             seed_prompts.append(
                 SeedPrompt(
@@ -114,6 +150,10 @@ async def fetch_dataset_async(self, *, cache: bool = True) -> SeedDataset:
                     source=source_url,
                     authors=authors,
                     groups=groups,
+                    metadata={
+                        "beaver_tails_categories": ",".join(raw_harm_categories),
+                        "beaver_tails_category_flags": json.dumps(item["category"], sort_keys=True),
+                    },
                 )
             )
 

pyrit/datasets/seed_datasets/remote/categorical_harmful_qa_dataset.py

@@ -8,6 +8,7 @@
     _RemoteDatasetLoader,
 )
 from pyrit.models import SeedDataset, SeedObjective, SeedUnion
+from pyrit.models.harm_category import HarmCategory
 
 logger = logging.getLogger(__name__)
 
@@ -107,13 +108,55 @@ async def fetch_dataset_async(self, *, cache: bool = True) -> SeedDataset:
         )
         source_url = f"https://huggingface.co/datasets/{self.HF_DATASET_NAME}"
         groups = ["DeCLaRe Lab, Singapore University of Technology and Design"]
+        harm_category_alias_overrides: dict[str, list[HarmCategory]] = {
+            "adult content": [HarmCategory.SEXUAL_CONTENT],
+            "child abuse": [HarmCategory.GROOMING, HarmCategory.SEXUAL_CONTENT, HarmCategory.CHILD_LEAKAGE],
+            "economic harm": [HarmCategory.SCAMS, HarmCategory.DECEPTION],
+            "fraud/deception": [HarmCategory.SCAMS, HarmCategory.DECEPTION],
+            "hate/harass/violence": [
+                HarmCategory.HATE_SPEECH,
+                HarmCategory.HARASSMENT,
+                HarmCategory.VIOLENT_CONTENT,
+            ],
+            "illegal activity": [HarmCategory.COORDINATION_HARM],
+            "malware viruses": [HarmCategory.MALWARE],
+            "physical harm": [HarmCategory.VIOLENT_CONTENT, HarmCategory.COORDINATION_HARM],
+            "political campaigning": [HarmCategory.CAMPAIGNING],
+            "privacy violation activity": [HarmCategory.PPI],
+            "tailored financial advice": [HarmCategory.FINANCIAL_ADVICE],
+        }
+        # "Child Abuse" spans both sexual and non-sexual sub-harms, so it is mapped at
+        # the subcategory level; otherwise non-sexual abuse (physical, emotional,
+        # neglect) would be incorrectly tagged with CSAM-related categories. Unknown
+        # subcategories fall back to the category-level mapping above.
+        child_abuse_subcategory_overrides: dict[str, list[HarmCategory]] = {
+            "content sexual exploitation": [
+                HarmCategory.GROOMING,
+                HarmCategory.SEXUAL_CONTENT,
+                HarmCategory.CHILD_LEAKAGE,
+            ],
+            "online predation": [HarmCategory.GROOMING, HarmCategory.SEXUAL_CONTENT, HarmCategory.CHILD_LEAKAGE],
+            "physical abuse": [HarmCategory.VIOLENT_CONTENT, HarmCategory.COORDINATION_HARM],
+            "emotional and psychological abuse": [HarmCategory.HARASSMENT],
+            "neglect": [HarmCategory.DANGEROUS_SITUATIONS],
+        }
+
+        def resolve_harm_categories(category: str | None, subcategory: str | None) -> list[str]:
+            if (
+                category is not None
+                and category.strip().lower() == "child abuse"
+                and subcategory is not None
+                and subcategory.strip().lower() in child_abuse_subcategory_overrides
+            ):
+                return self._standardize_harm_categories(subcategory, alias_overrides=child_abuse_subcategory_overrides)
+            return self._standardize_harm_categories(category, alias_overrides=harm_category_alias_overrides)
 
         seed_objectives: list[SeedUnion] = [
             SeedObjective(
                 value=item["Question"],
                 name="CategoricalHarmfulQA",
                 dataset_name=self.dataset_name,
-                harm_categories=[item["Category"]] if item.get("Category") else [],
+                harm_categories=resolve_harm_categories(item.get("Category"), item.get("Subcategory")),
                 description=description,
                 source=source_url,
                 authors=authors,