diff --git a/doc/api.rst b/doc/api.rst index 3921926fc8..2d9b9b6ebb 100644 --- a/doc/api.rst +++ b/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` ========================= diff --git a/pyrit/common/json_helper.py b/pyrit/common/json_helper.py index 6ee8462665..d2cb239fb0 100644 --- a/pyrit/common/json_helper.py +++ b/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") diff --git a/pyrit/datasets/__init__.py b/pyrit/datasets/__init__.py index 21f6fc4ca5..ed981016d5 100644 --- a/pyrit/datasets/__init__.py +++ b/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", diff --git a/pyrit/datasets/fetch_example_datasets.py b/pyrit/datasets/fetch_example_datasets.py index f6270b110d..8bd2b6aff0 100644 --- a/pyrit/datasets/fetch_example_datasets.py +++ b/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 diff --git a/tests/unit/test_aya_redteaming.py b/tests/unit/test_aya_redteaming.py new file mode 100644 index 0000000000..6e6e20d187 --- /dev/null +++ b/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"] diff --git a/tests/unit/test_fetch_examples.py b/tests/unit/test_fetch_examples.py index f370c658b1..68547fb07a 100644 --- a/tests/unit/test_fetch_examples.py +++ b/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)