From d03bdf3dfd429c0eed74c4864a8050abe0572d12 Mon Sep 17 00:00:00 2001
From: Paulina Kalicka <71526180+paulinek13@users.noreply.github.com>
Date: Mon, 24 Mar 2025 15:25:51 +0100
Subject: [PATCH 1/6] feat: add MLCommons AILuminate dataset

---
 doc/api.rst                                   |  1 +
 pyrit/datasets/__init__.py                    |  2 +
 .../datasets/mlcommons_ailuminate_dataset.py  | 48 +++++++++++++++++++
 .../datasets/test_fetch_datasets.py           |  2 +
 4 files changed, 53 insertions(+)
 create mode 100644 pyrit/datasets/mlcommons_ailuminate_dataset.py

diff --git a/doc/api.rst b/doc/api.rst
index 0d99494815..e956f9cfb9 100644
--- a/doc/api.rst
+++ b/doc/api.rst
@@ -122,6 +122,7 @@ API Reference
     fetch_librAI_do_not_answer_dataset
     fetch_llm_latent_adversarial_training_harmful_dataset
     fetch_many_shot_jailbreaking_dataset
+    fetch_mlcommons_ailuminate_dataset
     fetch_pku_safe_rlhf_dataset
     fetch_seclists_bias_testing_dataset
     fetch_tdc23_redteaming_dataset
diff --git a/pyrit/datasets/__init__.py b/pyrit/datasets/__init__.py
index f230f61726..d2bf1f9aa3 100644
--- a/pyrit/datasets/__init__.py
+++ b/pyrit/datasets/__init__.py
@@ -14,6 +14,7 @@
     fetch_llm_latent_adversarial_training_harmful_dataset,
 )
 from pyrit.datasets.many_shot_jailbreaking_dataset import fetch_many_shot_jailbreaking_dataset
+from pyrit.datasets.mlcommons_ailuminate_dataset import fetch_mlcommons_ailuminate_dataset
 from pyrit.datasets.pku_safe_rlhf_dataset import fetch_pku_safe_rlhf_dataset
 from pyrit.datasets.red_team_social_bias_dataset import fetch_red_team_social_bias_dataset
 from pyrit.datasets.seclists_bias_testing_dataset import fetch_seclists_bias_testing_dataset
@@ -33,6 +34,7 @@
     "fetch_librAI_do_not_answer_dataset",
     "fetch_llm_latent_adversarial_training_harmful_dataset",
     "fetch_many_shot_jailbreaking_dataset",
+    "fetch_mlcommons_ailuminate_dataset",
     "fetch_pku_safe_rlhf_dataset",
     "fetch_red_team_social_bias_dataset",
     "fetch_seclists_bias_testing_dataset",
diff --git a/pyrit/datasets/mlcommons_ailuminate_dataset.py b/pyrit/datasets/mlcommons_ailuminate_dataset.py
new file mode 100644
index 0000000000..3234a93ec4
--- /dev/null
+++ b/pyrit/datasets/mlcommons_ailuminate_dataset.py
@@ -0,0 +1,48 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT license.
+
+from pathlib import Path
+from typing import Literal, Optional
+
+from pyrit.datasets.dataset_helper import fetch_examples
+from pyrit.models import SeedPromptDataset
+from pyrit.models.seed_prompt import SeedPrompt
+
+
+def fetch_mlcommons_ailuminate_dataset(
+    source: str = (
+        "https://raw.githubusercontent.com/mlcommons/ailuminate/refs/heads/main/"
+        "airr_official_1.0_demo_en_us_prompt_set_release.csv"
+    ),
+    source_type: Literal["public_url"] = "public_url",
+    cache: bool = True,
+    data_home: Optional[Path] = None,
+) -> SeedPromptDataset:
+    """
+    Fetch MLCommons AI-Luminate examples and create a SeedPromptDataset.
+
+    Returns:
+        SeedPromptDataset: A SeedPromptDataset containing the examples.
+    """
+    examples = fetch_examples(source, source_type, cache, data_home)
+
+    seed_prompts = [
+        SeedPrompt(
+            value=example["prompt_text"],
+            data_type="text",
+            name="",
+            dataset_name="AILuminate v1.0 DEMO Prompt Set",
+            harm_categories=[example["hazard"]],
+            description=(
+                "This dataset contains the DEMO prompt library of the AILuminate 1.0 prompt dataset, created by"
+                " MLCommons AI Risk & Reliability working group. It contains 1,200 human-generated prompts that"
+                " cover 12 different hazard categories. The dataset is a mix of adversarial and non-adversarial"
+                " prompts. This dataset is a 10% subset of the full AILuminate training dataset."
+            ),
+            source="https://github.com/mlcommons/ailuminate",
+        )
+        for example in examples
+    ]
+
+    seed_prompt_dataset = SeedPromptDataset(prompts=seed_prompts)
+    return seed_prompt_dataset
diff --git a/tests/integration/datasets/test_fetch_datasets.py b/tests/integration/datasets/test_fetch_datasets.py
index 97d4cb3295..183a01d5d8 100644
--- a/tests/integration/datasets/test_fetch_datasets.py
+++ b/tests/integration/datasets/test_fetch_datasets.py
@@ -14,6 +14,7 @@
     fetch_librAI_do_not_answer_dataset,
     fetch_llm_latent_adversarial_training_harmful_dataset,
     fetch_many_shot_jailbreaking_dataset,
+    fetch_mlcommons_ailuminate_dataset,
     fetch_pku_safe_rlhf_dataset,
     fetch_red_team_social_bias_dataset,
     fetch_seclists_bias_testing_dataset,
@@ -37,6 +38,7 @@
         (fetch_librAI_do_not_answer_dataset, True),
         (fetch_llm_latent_adversarial_training_harmful_dataset, True),
         (fetch_many_shot_jailbreaking_dataset, False),
+        (fetch_mlcommons_ailuminate_dataset, True),
         (fetch_pku_safe_rlhf_dataset, True),
         (fetch_red_team_social_bias_dataset, True),
         (fetch_seclists_bias_testing_dataset, True),

From 337608ff22731874161cfc1949c407a4e2563f22 Mon Sep 17 00:00:00 2001
From: Paulina Kalicka <71526180+paulinek13@users.noreply.github.com>
Date: Tue, 25 Mar 2025 10:09:32 +0100
Subject: [PATCH 2/6] docstring update

---
 pyrit/datasets/mlcommons_ailuminate_dataset.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pyrit/datasets/mlcommons_ailuminate_dataset.py b/pyrit/datasets/mlcommons_ailuminate_dataset.py
index 3234a93ec4..403aca0899 100644
--- a/pyrit/datasets/mlcommons_ailuminate_dataset.py
+++ b/pyrit/datasets/mlcommons_ailuminate_dataset.py
@@ -19,7 +19,7 @@ def fetch_mlcommons_ailuminate_dataset(
     data_home: Optional[Path] = None,
 ) -> SeedPromptDataset:
     """
-    Fetch MLCommons AI-Luminate examples and create a SeedPromptDataset.
+    Fetch examples from AILuminate v1.0 DEMO Prompt Set and create a SeedPromptDataset.
 
     Returns:
         SeedPromptDataset: A SeedPromptDataset containing the examples.

From 7b94dacbbc5114f9200119d64923bcecce6ba345 Mon Sep 17 00:00:00 2001
From: Paulina Kalicka <71526180+paulinek13@users.noreply.github.com>
Date: Tue, 25 Mar 2025 15:22:24 +0100
Subject: [PATCH 3/6] renaming

---
 doc/api.rst                                                   | 2 +-
 pyrit/datasets/__init__.py                                    | 4 ++--
 ...minate_dataset.py => mlcommons_ailuminate_demo_dataset.py} | 2 +-
 tests/integration/datasets/test_fetch_datasets.py             | 4 ++--
 4 files changed, 6 insertions(+), 6 deletions(-)
 rename pyrit/datasets/{mlcommons_ailuminate_dataset.py => mlcommons_ailuminate_demo_dataset.py} (97%)

diff --git a/doc/api.rst b/doc/api.rst
index e956f9cfb9..f4ee85dae4 100644
--- a/doc/api.rst
+++ b/doc/api.rst
@@ -122,7 +122,7 @@ API Reference
     fetch_librAI_do_not_answer_dataset
     fetch_llm_latent_adversarial_training_harmful_dataset
     fetch_many_shot_jailbreaking_dataset
-    fetch_mlcommons_ailuminate_dataset
+    fetch_mlcommons_ailuminate_demo_dataset
     fetch_pku_safe_rlhf_dataset
     fetch_seclists_bias_testing_dataset
     fetch_tdc23_redteaming_dataset
diff --git a/pyrit/datasets/__init__.py b/pyrit/datasets/__init__.py
index d2bf1f9aa3..2b2a6f4408 100644
--- a/pyrit/datasets/__init__.py
+++ b/pyrit/datasets/__init__.py
@@ -14,7 +14,7 @@
     fetch_llm_latent_adversarial_training_harmful_dataset,
 )
 from pyrit.datasets.many_shot_jailbreaking_dataset import fetch_many_shot_jailbreaking_dataset
-from pyrit.datasets.mlcommons_ailuminate_dataset import fetch_mlcommons_ailuminate_dataset
+from pyrit.datasets.mlcommons_ailuminate_demo_dataset import fetch_mlcommons_ailuminate_demo_dataset
 from pyrit.datasets.pku_safe_rlhf_dataset import fetch_pku_safe_rlhf_dataset
 from pyrit.datasets.red_team_social_bias_dataset import fetch_red_team_social_bias_dataset
 from pyrit.datasets.seclists_bias_testing_dataset import fetch_seclists_bias_testing_dataset
@@ -34,7 +34,7 @@
     "fetch_librAI_do_not_answer_dataset",
     "fetch_llm_latent_adversarial_training_harmful_dataset",
     "fetch_many_shot_jailbreaking_dataset",
-    "fetch_mlcommons_ailuminate_dataset",
+    "fetch_mlcommons_ailuminate_demo_dataset",
     "fetch_pku_safe_rlhf_dataset",
     "fetch_red_team_social_bias_dataset",
     "fetch_seclists_bias_testing_dataset",
diff --git a/pyrit/datasets/mlcommons_ailuminate_dataset.py b/pyrit/datasets/mlcommons_ailuminate_demo_dataset.py
similarity index 97%
rename from pyrit/datasets/mlcommons_ailuminate_dataset.py
rename to pyrit/datasets/mlcommons_ailuminate_demo_dataset.py
index 403aca0899..f0ea78c932 100644
--- a/pyrit/datasets/mlcommons_ailuminate_dataset.py
+++ b/pyrit/datasets/mlcommons_ailuminate_demo_dataset.py
@@ -9,7 +9,7 @@
 from pyrit.models.seed_prompt import SeedPrompt
 
 
-def fetch_mlcommons_ailuminate_dataset(
+def fetch_mlcommons_ailuminate_demo_dataset(
     source: str = (
         "https://raw.githubusercontent.com/mlcommons/ailuminate/refs/heads/main/"
         "airr_official_1.0_demo_en_us_prompt_set_release.csv"
diff --git a/tests/integration/datasets/test_fetch_datasets.py b/tests/integration/datasets/test_fetch_datasets.py
index 183a01d5d8..d930be3cf1 100644
--- a/tests/integration/datasets/test_fetch_datasets.py
+++ b/tests/integration/datasets/test_fetch_datasets.py
@@ -14,7 +14,7 @@
     fetch_librAI_do_not_answer_dataset,
     fetch_llm_latent_adversarial_training_harmful_dataset,
     fetch_many_shot_jailbreaking_dataset,
-    fetch_mlcommons_ailuminate_dataset,
+    fetch_mlcommons_ailuminate_demo_dataset,
     fetch_pku_safe_rlhf_dataset,
     fetch_red_team_social_bias_dataset,
     fetch_seclists_bias_testing_dataset,
@@ -38,7 +38,7 @@
         (fetch_librAI_do_not_answer_dataset, True),
         (fetch_llm_latent_adversarial_training_harmful_dataset, True),
         (fetch_many_shot_jailbreaking_dataset, False),
-        (fetch_mlcommons_ailuminate_dataset, True),
+        (fetch_mlcommons_ailuminate_demo_dataset, True),
         (fetch_pku_safe_rlhf_dataset, True),
         (fetch_red_team_social_bias_dataset, True),
         (fetch_seclists_bias_testing_dataset, True),

From 770a97f9c514a9ca57004424614a55046a2aa35f Mon Sep 17 00:00:00 2001
From: Paulina Kalicka <71526180+paulinek13@users.noreply.github.com>
Date: Tue, 25 Mar 2025 16:59:35 +0100
Subject: [PATCH 4/6] add hazard categories mapping

---
 .../mlcommons_ailuminate_demo_dataset.py      | 22 ++++++++++++++++++-
 1 file changed, 21 insertions(+), 1 deletion(-)

diff --git a/pyrit/datasets/mlcommons_ailuminate_demo_dataset.py b/pyrit/datasets/mlcommons_ailuminate_demo_dataset.py
index f0ea78c932..b93e79342e 100644
--- a/pyrit/datasets/mlcommons_ailuminate_demo_dataset.py
+++ b/pyrit/datasets/mlcommons_ailuminate_demo_dataset.py
@@ -8,6 +8,25 @@
 from pyrit.models import SeedPromptDataset
 from pyrit.models.seed_prompt import SeedPrompt
 
+# https://github.com/mlcommons/ailuminate#hazard-categories-and-abbreviations
+_hazard_categories_map = {
+    "vcr": "Violent Crimes",
+    "src": "Sex-Related Crimes",
+    "cse": "Child Sexual Exploitation",
+    "ssh": "Suicide & Self-Harm",
+    "iwp": "Indiscriminate Weapons (CBRNE)",
+    "ipv": "Intellectual Property Violations",
+    "dfm": "Defamation",
+    "ncr": "Non-Violent Crimes",
+    "hte": "Hate",
+    "prv": "Privacy Violations",
+    "spc_ele": "Specialized Advice; elections",
+    "spc_fin": "Specialized Advice; financial",
+    "spc_hlt": "Specialized Advice; health",
+    "sxc_prn": "Sexual Content; pornographic",
+    "spc_lgl": "Specialized Advice; legal",
+}
+
 
 def fetch_mlcommons_ailuminate_demo_dataset(
     source: str = (
@@ -32,7 +51,8 @@ def fetch_mlcommons_ailuminate_demo_dataset(
             data_type="text",
             name="",
             dataset_name="AILuminate v1.0 DEMO Prompt Set",
-            harm_categories=[example["hazard"]],
+            # Saving both the full hazard category name and the abbreviation
+            harm_categories=[_hazard_categories_map[example["hazard"]], example["hazard"]],
             description=(
                 "This dataset contains the DEMO prompt library of the AILuminate 1.0 prompt dataset, created by"
                 " MLCommons AI Risk & Reliability working group. It contains 1,200 human-generated prompts that"

From 16dfc75bd3880cd6a4669c673e1c10cad0b1edad Mon Sep 17 00:00:00 2001
From: Paulina Kalicka <71526180+paulinek13@users.noreply.github.com>
Date: Tue, 25 Mar 2025 17:16:54 +0100
Subject: [PATCH 5/6] add a note

---
 pyrit/datasets/mlcommons_ailuminate_demo_dataset.py | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/pyrit/datasets/mlcommons_ailuminate_demo_dataset.py b/pyrit/datasets/mlcommons_ailuminate_demo_dataset.py
index b93e79342e..aba4da4923 100644
--- a/pyrit/datasets/mlcommons_ailuminate_demo_dataset.py
+++ b/pyrit/datasets/mlcommons_ailuminate_demo_dataset.py
@@ -42,6 +42,12 @@ def fetch_mlcommons_ailuminate_demo_dataset(
 
     Returns:
         SeedPromptDataset: A SeedPromptDataset containing the examples.
+
+    Note:
+        For more information about the dataset and related materials, visit:
+        - https://github.com/mlcommons/ailuminate \n
+        - https://mlcommons.org/ailuminate/ \n
+        - https://doi.org/10.48550/arXiv.2404.12241
     """
     examples = fetch_examples(source, source_type, cache, data_home)
 

From c621deb91906294935f620ec147b7d0a69d214ff Mon Sep 17 00:00:00 2001
From: Paulina Kalicka <71526180+paulinek13@users.noreply.github.com>
Date: Tue, 25 Mar 2025 17:26:55 +0100
Subject: [PATCH 6/6] tiny fix in the docstring

---
 pyrit/datasets/mlcommons_ailuminate_demo_dataset.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pyrit/datasets/mlcommons_ailuminate_demo_dataset.py b/pyrit/datasets/mlcommons_ailuminate_demo_dataset.py
index aba4da4923..a98ce4bd4c 100644
--- a/pyrit/datasets/mlcommons_ailuminate_demo_dataset.py
+++ b/pyrit/datasets/mlcommons_ailuminate_demo_dataset.py
@@ -44,7 +44,7 @@ def fetch_mlcommons_ailuminate_demo_dataset(
         SeedPromptDataset: A SeedPromptDataset containing the examples.
 
     Note:
-        For more information about the dataset and related materials, visit:
+        For more information about the dataset and related materials, visit: \n
         - https://github.com/mlcommons/ailuminate \n
         - https://mlcommons.org/ailuminate/ \n
         - https://doi.org/10.48550/arXiv.2404.12241