FEAT: add a fetch function for Aya Red-teaming Dataset

doc/api.rst

@@ -111,6 +111,7 @@ API Reference
     fetch_forbidden_questions_dataset
     fetch_llm_latent_adversarial_training_harmful_dataset
     fetch_tdc23_redteaming_dataset
+    fetch_aya_redteaming_dataset
 
 :py:mod:`pyrit.embedding`
 =========================

pyrit/common/json_helper.py

@@ -11,3 +11,12 @@ def read_json(file) -> List[Dict[str, str]]:
 
 def write_json(file, examples: List[Dict[str, str]]):
     json.dump(examples, file)
+
+
+def read_jsonl(file) -> List[Dict[str, str]]:
+    return [json.loads(line) for line in file]
+
+
+def write_jsonl(file, examples: List[Dict[str, str]]):
+    for example in examples:
+        file.write(json.dumps(example) + "\n")

pyrit/datasets/__init__.py

@@ -3,6 +3,7 @@
 
 from .fetch_example_datasets import (
     fetch_adv_bench_dataset,
+    fetch_aya_redteaming_dataset,
     fetch_decoding_trust_stereotypes_dataset,
     fetch_examples,
     fetch_forbidden_questions_dataset,
@@ -18,6 +19,7 @@
 )
 
 __all__ = [
+    "fetch_aya_redteaming_dataset",
     "fetch_decoding_trust_stereotypes_dataset",
     "fetch_examples",
     "fetch_harmbench_dataset",

pyrit/datasets/fetch_example_datasets.py

@@ -13,7 +13,7 @@
 from datasets import load_dataset
 
 from pyrit.common.csv_helper import read_csv, write_csv
-from pyrit.common.json_helper import read_json, write_json
+from pyrit.common.json_helper import read_json, read_jsonl, write_json, write_jsonl
 from pyrit.common.path import DATASETS_PATH, DB_DATA_PATH
 from pyrit.common.text_helper import read_txt, write_txt
 from pyrit.models import (
@@ -30,6 +30,7 @@
 
 FILE_TYPE_HANDLERS: Dict[str, Dict[str, Callable]] = {
     "json": {"read": read_json, "write": write_json},
+    "jsonl": {"read": read_jsonl, "write": write_jsonl},
     "csv": {"read": read_csv, "write": write_csv},
     "txt": {"read": read_txt, "write": write_txt},
 }
@@ -844,3 +845,89 @@ def fetch_wmdp_dataset(category: Optional[str] = None) -> QuestionAnsweringDatas
     )
 
     return dataset
+
+
+def fetch_aya_redteaming_dataset(
+    cache: bool = True,
+    data_home: Optional[Path] = None,
+    language: Literal["English", "Hindi", "French", "Spanish", "Arabic", "Russian", "Serbian", "Tagalog"] = "English",
+    harm_categories: Optional[
+        List[
+            Literal[
+                "Bullying & Harassment",
+                "Discrimination & Injustice",
+                "Graphic material",
+                "Harms of Representation Allocation and Quality of Service",
+                "Hate Speech",
+                "Non-consensual sexual content",
+                "Profanity",
+                "Self-Harm",
+                "Violence, Threats & Incitement",
+            ]
+        ]
+    ] = None,
+    harm_scope: Optional[Literal["global", "local"]] = None,
+) -> SeedPromptDataset:
+    """
+    Fetch examples from the Aya Red-teaming dataset with optional filtering and create a SeedPromptDataset.
+
+    Args:
+        cache (bool): Whether to cache the fetched examples. Defaults to True.
+        data_home: Directory to store cached data. Defaults to None.
+        language: The language to fetch prompts for. Defaults to "English".
+        harm_categories: List of harm categories to filter the examples.
+            Defaults to None, which means all categories are included.
+            Otherwise, only prompts with at least one matching category are included.
+        harm_scope: Whether to fetch globally or locally harmful prompts.
+            Defaults to None, which means all examples are included.
+
+    Returns:
+        SeedPromptDataset: A SeedPromptDataset containing the filtered examples.
+
+    Note:
+        For more information and access to the original dataset and related materials, visit:
+        https://huggingface.co/datasets/CohereForAI/aya_redteaming/blob/main/README.md \n
+        Related paper: https://arxiv.org/abs/2406.18682 \n
+        The dataset license: Apache 2.0
+
+    Warning:
+        Due to the nature of these prompts, it may be advisable to consult your relevant legal
+        department before testing them with LLMs to ensure compliance and reduce potential risks.
+    """
+    _lang = {
+        "English": "eng",
+        "Hindi": "hin",
+        "French": "fra",
+        "Spanish": "spa",
+        "Arabic": "arb",
+        "Russian": "rus",
+        "Serbian": "srp",
+        "Tagalog": "tgl",
+    }
+
+    examples = fetch_examples(
+        source=f"https://huggingface.co/datasets/CohereForAI/aya_redteaming/raw/main/aya_{_lang[language]}.jsonl",
+        source_type="public_url",
+        cache=cache,
+        data_home=data_home,
+    )
+
+    seed_prompts = []
+
+    for example in examples:
+        categories = eval(example["harm_category"])
+        if harm_categories is None or any(cat in categories for cat in harm_categories):
+            if harm_scope is None or example["global_or_local"] == harm_scope:
+                seed_prompts.append(
+                    SeedPrompt(
+                        value=example["prompt"],
+                        data_type="text",
+                        name="Aya Red-teaming Examples",
+                        dataset_name="Aya Red-teaming Examples",
+                        harm_categories=categories,
+                        source="https://huggingface.co/datasets/CohereForAI/aya_redteaming",
+                    )
+                )
+
+    seed_prompt_dataset = SeedPromptDataset(prompts=seed_prompts)
+    return seed_prompt_dataset

tests/unit/test_aya_redteaming.py

@@ -0,0 +1,115 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT license.
+
+from unittest.mock import patch
+
+import pytest
+
+from pyrit.datasets import fetch_aya_redteaming_dataset
+from pyrit.models import SeedPromptDataset
+
+
+@pytest.fixture
+def mock_aya_redteaming_data():
+    return [
+        {
+            "prompt": "[Prompt #001] ...",
+            "language": "English",
+            "harm_category": '["Discrimination & Injustice"]',
+            "global_or_local": "global",
+        },
+        {
+            "prompt": "[Prompt #002] ...",
+            "language": "English",
+            "harm_category": '["Harms of Representation Allocation and Quality of Service", '
+            '"Discrimination & Injustice"]',
+            "global_or_local": "global",
+        },
+        {
+            "prompt": "[Prompt #003] ...",
+            "language": "English",
+            "harm_category": '["Non-consensual sexual content"]',
+            "global_or_local": "local",
+        },
+        {
+            "prompt": "[Prompt #004] ...",
+            "language": "English",
+            "harm_category": '["Self-Harm", "Violence, Threats & Incitement"]',
+            "global_or_local": "local",
+        },
+        {
+            "prompt": "[Prompt #005] ...",
+            "language": "English",
+            "harm_category": '["Violence, Threats & Incitement", "Bullying & Harassment", "Profanity"]',
+            "global_or_local": "global",
+        },
+        {
+            "prompt": "[Prompt #006] ...",
+            "language": "English",
+            "harm_category": '["Violence, Threats & Incitement", "Graphic material"]',
+            "global_or_local": "global",
+        },
+        {
+            "prompt": "[Prompt #007] ...",
+            "language": "English",
+            "harm_category": '["Bullying & Harassment"]',
+            "global_or_local": "global",
+        },
+    ]
+
+
+@patch("pyrit.datasets.fetch_example_datasets.fetch_examples")
+def test_fetch_aya_redteaming_dataset(mock_fetch_examples, mock_aya_redteaming_data):
+    mock_fetch_examples.return_value = mock_aya_redteaming_data
+
+    # Test fetching the dataset without any filters
+
+    dataset = fetch_aya_redteaming_dataset()
+
+    assert isinstance(dataset, SeedPromptDataset)
+    assert len(dataset.prompts) == 7
+
+    first_prompt = dataset.prompts[0]
+    assert first_prompt.value == "[Prompt #001] ..."
+    assert first_prompt.data_type == "text"
+    assert first_prompt.name == "Aya Red-teaming Examples"
+    assert first_prompt.dataset_name == "Aya Red-teaming Examples"
+    assert first_prompt.harm_categories == ["Discrimination & Injustice"]
+    assert first_prompt.source == "https://huggingface.co/datasets/CohereForAI/aya_redteaming"
+
+    assert dataset.prompts[1].value == "[Prompt #002] ..."
+    assert dataset.prompts[3].harm_categories == ["Self-Harm", "Violence, Threats & Incitement"]
+
+    # Test fetching the dataset with a `harm_categories` filter
+
+    dataset = fetch_aya_redteaming_dataset(harm_categories=["Bullying & Harassment"])
+
+    assert len(dataset.prompts) == 2
+    assert dataset.prompts[0].value == "[Prompt #005] ..."
+
+    dataset = fetch_aya_redteaming_dataset(harm_categories=["Discrimination & Injustice", "Graphic material"])
+
+    assert len(dataset.prompts) == 3
+    assert dataset.prompts[0].value == "[Prompt #001] ..."
+    assert dataset.prompts[2].value == "[Prompt #006] ..."
+
+    # Test fetching the dataset with a `harm_scope` filter
+
+    dataset = fetch_aya_redteaming_dataset(harm_scope="local")
+
+    assert len(dataset.prompts) == 2
+    assert dataset.prompts[0].value == "[Prompt #003] ..."
+
+    # Test fetching the dataset with multiple filters
+
+    dataset = fetch_aya_redteaming_dataset(harm_categories=["Violence, Threats & Incitement"], harm_scope="global")
+
+    assert len(dataset.prompts) == 2
+    assert dataset.prompts[0].value == "[Prompt #005] ..."
+
+    # Test fetching the dataset with a `language` filter
+
+    dataset = fetch_aya_redteaming_dataset(language="French")
+
+    call_args = mock_fetch_examples.call_args
+    assert "aya_fra.jsonl" in call_args.kwargs["source"]

tests/unit/test_fetch_examples.py

@@ -22,12 +22,13 @@
 # These URLs are placeholders and will be mocked in tests
 SOURCE_URLS = {
     "json": "https://example.com/examples.json",
+    "jsonl": "https://example.com/examples.jsonl",
     "csv": "https://example.com/examples.csv",
     "txt": "https://example.com/examples.txt",
 }
 
 
-FILE_TYPES = ["json", "csv", "txt"]
+FILE_TYPES = ["json", "jsonl", "csv", "txt"]
 UNSUPPORTED_FILE_TYPES = ["xml", "pdf", "docx"]  # Unsupported file types for testing
 
 
@@ -117,7 +118,7 @@ def test_fetch_from_public_url_unsupported(file_type):
     mock_response.text = "example content"
 
     with patch("requests.get", return_value=mock_response):
-        with pytest.raises(ValueError, match="Invalid file_type. Expected one of: json, csv, txt."):
+        with pytest.raises(ValueError, match="Invalid file_type. Expected one of: json, jsonl, csv, txt."):
             _fetch_from_public_url(url, file_type)
 
 
@@ -133,7 +134,7 @@ def test_read_cache_unsupported(file_type):
     cache_file_dir.mkdir(parents=True, exist_ok=True)
     cache_file.touch()
 
-    with pytest.raises(ValueError, match="Invalid file_type. Expected one of: json, csv, txt."):
+    with pytest.raises(ValueError, match="Invalid file_type. Expected one of: json, jsonl, csv, txt."):
         _read_cache(cache_file, file_type)
 
     # Cleanup the created file after the test
@@ -148,7 +149,7 @@ def test_write_cache_unsupported(file_type):
     cache_file = DB_DATA_PATH / ".pyrit_test" / "datasets" / f"cache_file.{file_type}"
     examples = [{"prompt": "example"}]
 
-    with pytest.raises(ValueError, match="Invalid file_type. Expected one of: json, csv, txt."):
+    with pytest.raises(ValueError, match="Invalid file_type. Expected one of: json, jsonl, csv, txt."):
         _write_cache(cache_file, examples, file_type)