diff --git a/doc/code/datasets/1_loading_datasets.ipynb b/doc/code/datasets/1_loading_datasets.ipynb index e692089df..e95be7fa7 100644 --- a/doc/code/datasets/1_loading_datasets.ipynb +++ b/doc/code/datasets/1_loading_datasets.ipynb @@ -20,6 +20,14 @@ "id": "1", "metadata": {}, "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "/mnt/c/Users/warisgill/Documents/PyRIT/.venv/lib/python3.12/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n", + " from .autonotebook import tqdm as notebook_tqdm\n" + ] + }, { "data": { "text/plain": [ @@ -40,6 +48,7 @@ " 'airt_violence',\n", " 'aya_redteaming',\n", " 'babelscape_alert',\n", + " 'cbt_bench',\n", " 'ccp_sensitive_prompts',\n", " 'dark_bench',\n", " 'equitymedqa',\n", @@ -100,40 +109,72 @@ "name": "stderr", "output_type": "stream", "text": [ - "\r\n", - "Loading datasets - this can take a few minutes: 0%| | 0/49 [00:00 str: + """Return the dataset name.""" + return "cbt_bench" + + async def fetch_dataset(self, *, cache: bool = True) -> SeedDataset: + """ + Fetch CBT-Bench dataset from HuggingFace and return as SeedDataset. + + Args: + cache: Whether to cache the fetched dataset. Defaults to True. + + Returns: + SeedDataset: A SeedDataset containing CBT-Bench examples. + + Raises: + ValueError: If the dataset is empty after processing. + Exception: If the dataset cannot be loaded or processed. + """ + logger.info(f"Loading CBT-Bench dataset from {self.source} (config={self.config})") + + data = await self._fetch_from_huggingface( + dataset_name=self.source, + config=self.config, + split=self.split, + cache=cache, + ) + + authors = [ + "Mian Zhang", + "Xianjun Yang", + "Xinlu Zhang", + "Travis Labrum", + "Jamie C Chiu", + "Shaun M Eack", + "Fei Fang", + "William Yang Wang", + "Zhiyu Zoey Chen", + ] + description = ( + "CBT-Bench is a benchmark designed to evaluate the proficiency of Large Language Models " + "in assisting Cognitive Behavioral Therapy (CBT). The dataset covers basic CBT knowledge, " + "cognitive model understanding, and therapeutic response generation." + ) + + seed_prompts = [] + + for item in data: + situation = item.get("situation", "").strip() + thoughts = item.get("thoughts", "").strip() + + # Combine situation and thoughts as the prompt value + if situation and thoughts: + value = f"Situation: {situation}\n\nThoughts: {thoughts}" + elif situation: + value = situation + elif thoughts: + value = thoughts + else: + logger.warning("[CBT-Bench] Skipping item with no situation or thoughts") + continue + + # Extract core beliefs for metadata + core_beliefs = item.get("core_belief_fine_grained", []) + + metadata: dict[str, Any] = { + "config": self.config, + } + + if core_beliefs: + metadata["core_belief_fine_grained"] = core_beliefs + + seed_prompt = SeedPrompt( + value=value, + data_type="text", + dataset_name=self.dataset_name, + harm_categories=["psycho-social harms"], + description=description, + source=f"https://huggingface.co/datasets/{self.source}", + authors=authors, + metadata=metadata, + ) + + seed_prompts.append(seed_prompt) + + if not seed_prompts: + raise ValueError("SeedDataset cannot be empty.") + + logger.info(f"Successfully loaded {len(seed_prompts)} examples from CBT-Bench dataset") + + return SeedDataset(seeds=seed_prompts, dataset_name=self.dataset_name) diff --git a/tests/unit/datasets/test_cbt_bench_dataset.py b/tests/unit/datasets/test_cbt_bench_dataset.py new file mode 100644 index 000000000..9702e3719 --- /dev/null +++ b/tests/unit/datasets/test_cbt_bench_dataset.py @@ -0,0 +1,150 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +from unittest.mock import patch + +import pytest + +from pyrit.datasets.seed_datasets.remote.cbt_bench_dataset import _CBTBenchDataset +from pyrit.models import SeedDataset, SeedPrompt + + +@pytest.fixture +def mock_cbt_bench_data(): + """Mock data for CBT-Bench dataset.""" + return [ + { + "id": "1", + "ori_text": "Original text 1", + "situation": "I feel guilty for lying to my boyfriend.", + "thoughts": "I feel ashamed and afraid of losing his trust.", + "core_belief_fine_grained": ["I am unlovable", "I am immoral"], + }, + { + "id": "2", + "ori_text": "Original text 2", + "situation": "I am concerned I may have schizophrenia.", + "thoughts": "I hear voices and see things. I am scared.", + "core_belief_fine_grained": ["I am powerless, weak, vulnerable", "I am out of control"], + }, + ] + + +@pytest.fixture +def mock_cbt_bench_data_missing_thoughts(): + """Mock data with missing thoughts field.""" + return [ + { + "id": "1", + "situation": "A situation without thoughts.", + "thoughts": "", + "core_belief_fine_grained": ["I am helpless"], + }, + ] + + +@pytest.fixture +def mock_cbt_bench_data_empty(): + """Mock data with all empty fields.""" + return [ + { + "id": "1", + "situation": "", + "thoughts": "", + "core_belief_fine_grained": [], + }, + ] + + +class TestCBTBenchDataset: + """Test the CBT-Bench dataset loader.""" + + def test_dataset_name(self): + """Test dataset_name property.""" + loader = _CBTBenchDataset() + assert loader.dataset_name == "cbt_bench" + + @pytest.mark.asyncio + async def test_fetch_dataset(self, mock_cbt_bench_data): + """Test fetching CBT-Bench dataset with mocked data.""" + loader = _CBTBenchDataset() + + with patch.object(loader, "_fetch_from_huggingface", return_value=mock_cbt_bench_data): + dataset = await loader.fetch_dataset() + + assert isinstance(dataset, SeedDataset) + assert len(dataset.seeds) == 2 + assert all(isinstance(p, SeedPrompt) for p in dataset.seeds) + + # Check first prompt combines situation and thoughts + first_prompt = dataset.seeds[0] + assert "I feel guilty for lying to my boyfriend." in first_prompt.value + assert "I feel ashamed and afraid of losing his trust." in first_prompt.value + assert first_prompt.value.startswith("Situation:") + assert "Thoughts:" in first_prompt.value + assert first_prompt.data_type == "text" + assert first_prompt.dataset_name == "cbt_bench" + assert first_prompt.harm_categories == ["psycho-social harms"] + assert first_prompt.metadata["core_belief_fine_grained"] == ["I am unlovable", "I am immoral"] + + @pytest.mark.asyncio + async def test_fetch_dataset_with_custom_config(self, mock_cbt_bench_data): + """Test fetching with custom HuggingFace config and split.""" + loader = _CBTBenchDataset( + source="custom/cbt-bench", + config="core_major_seed", + split="test", + ) + + with patch.object(loader, "_fetch_from_huggingface", return_value=mock_cbt_bench_data) as mock_fetch: + dataset = await loader.fetch_dataset(cache=False) + + assert len(dataset.seeds) == 2 + mock_fetch.assert_called_once() + call_kwargs = mock_fetch.call_args.kwargs + assert call_kwargs["dataset_name"] == "custom/cbt-bench" + assert call_kwargs["config"] == "core_major_seed" + assert call_kwargs["split"] == "test" + assert call_kwargs["cache"] is False + + @pytest.mark.asyncio + async def test_fetch_dataset_situation_only(self, mock_cbt_bench_data_missing_thoughts): + """Test that items with only situation (no thoughts) still work.""" + loader = _CBTBenchDataset() + + with patch.object(loader, "_fetch_from_huggingface", return_value=mock_cbt_bench_data_missing_thoughts): + dataset = await loader.fetch_dataset() + + assert len(dataset.seeds) == 1 + assert dataset.seeds[0].value == "A situation without thoughts." + + @pytest.mark.asyncio + async def test_fetch_dataset_empty_raises(self, mock_cbt_bench_data_empty): + """Test that an empty dataset raises ValueError.""" + loader = _CBTBenchDataset() + + with patch.object(loader, "_fetch_from_huggingface", return_value=mock_cbt_bench_data_empty): + with pytest.raises(ValueError, match="SeedDataset cannot be empty"): + await loader.fetch_dataset() + + @pytest.mark.asyncio + async def test_fetch_dataset_metadata_includes_config(self, mock_cbt_bench_data): + """Test that metadata includes the config name.""" + loader = _CBTBenchDataset(config="distortions_seed") + + with patch.object(loader, "_fetch_from_huggingface", return_value=mock_cbt_bench_data): + dataset = await loader.fetch_dataset() + + for seed in dataset.seeds: + assert seed.metadata["config"] == "distortions_seed" + + @pytest.mark.asyncio + async def test_fetch_dataset_source_url(self, mock_cbt_bench_data): + """Test that source URL is correctly set.""" + loader = _CBTBenchDataset() + + with patch.object(loader, "_fetch_from_huggingface", return_value=mock_cbt_bench_data): + dataset = await loader.fetch_dataset() + + for seed in dataset.seeds: + assert seed.source == "https://huggingface.co/datasets/Psychotherapy-LLM/CBT-Bench"