From 84a4d1708e42e8059cf3150e5c35dcd821c0b39e Mon Sep 17 00:00:00 2001 From: "Liao, Jie" <28374150+leo0481@users.noreply.github.com> Date: Thu, 18 Dec 2025 23:32:04 +0800 Subject: [PATCH] Enhance project configuration and structure: Added build-system requirements in pyproject.toml, updated README for installation instructions, and refactored attack classes to utilize a new create_test_case method. Introduced a plugins.yaml for component registration and removed the legacy registry initialization file. --- README.md | 42 +++-- attacks/bap/attack.py | 18 +-- attacks/cs_dj/attack.py | 20 +-- attacks/figstep/attack.py | 19 +-- attacks/hades/attack.py | 22 +-- attacks/himrd/attack.py | 21 +-- attacks/imgjp/attack.py | 18 +-- attacks/jood/attack.py | 19 +-- attacks/mml/attack.py | 22 +-- attacks/pba/attack.py | 18 +-- attacks/qr/attack.py | 24 ++- attacks/si/attack.py | 35 ++--- attacks/umk/attack.py | 17 +- attacks/viscra/attack.py | 18 +-- attacks/visual_adv/attack.py | 18 +-- config/plugins.yaml | 55 +++++++ core/base_classes.py | 161 ++++++++++++++++++- core/registry_init.py | 69 -------- core/unified_registry.py | 153 +++++++++++++----- defenses/__init__.py | 2 +- defenses/adashield.py | 2 +- defenses/base_defense.py | 57 ------- defenses/cider.py | 15 +- defenses/coca.py | 12 +- defenses/dps.py | 2 +- defenses/ecso.py | 10 +- defenses/guardreasoner_vl.py | 15 +- defenses/hiddendetect.py | 21 ++- defenses/jailguard.py | 10 +- defenses/llama_guard_3.py | 18 +-- defenses/llama_guard_4.py | 12 +- defenses/llavaguard.py | 15 +- defenses/mllm_protector.py | 11 +- defenses/qguard.py | 15 +- defenses/shieldlm.py | 17 +- defenses/uniguard.py | 2 +- defenses/vlguard.py | 10 +- evaluators/__init__.py | 2 +- evaluators/base_evaluator.py | 34 ---- evaluators/default_judge.py | 12 +- models/base_model.py | 146 ++++++++++++++++- pipeline/base_pipeline.py | 166 +++++++++++++++----- pipeline/batch_strategies.py | 115 +++++++------- pipeline/evaluate_results.py | 268 ++++++++++++++++++++++---------- pipeline/generate_responses.py | 182 ++++++++++++++++++---- pipeline/generate_test_cases.py | 16 +- pipeline/resource_policy.py | 121 ++++++++++++++ pyproject.toml | 54 ++++++- tests/test_core.py | 70 +++++++-- 49 files changed, 1474 insertions(+), 727 deletions(-) create mode 100644 config/plugins.yaml delete mode 100644 core/registry_init.py delete mode 100644 defenses/base_defense.py delete mode 100644 evaluators/base_evaluator.py create mode 100644 pipeline/resource_policy.py diff --git a/README.md b/README.md index edde67b..458447f 100644 --- a/README.md +++ b/README.md @@ -42,8 +42,18 @@ A one-stop multimodal jailbreak/defense evaluation framework for beginners, cove 1. **Requirements**: Python 3.10+, recommended to use `uv` or `pip`. 2. **Install dependencies** (choose one): - - `uv sync` - - `pip install -e .` + - **Recommended (uv, uses `uv.lock`)**: + - `uv sync` + - **Editable install (uv)**: + - `uv pip install -e .` + - **Editable install (pip)**: + - Note: this project pins PyTorch as `torch==2.6.0+cu118` / `torchvision==0.21.0+cu118`. + These CUDA wheels are **NOT** on PyPI, so plain `pip install -e .` will fail unless you add the PyTorch CUDA index. + - CUDA 11.8: + - `pip install -e . --extra-index-url https://download.pytorch.org/whl/cu118` + - If you prefer installing PyTorch first (often more reliable on servers): + - `pip install torch==2.6.0+cu118 torchvision==0.21.0+cu118 --index-url https://download.pytorch.org/whl/cu118` + - `pip install -e . --no-deps` 3. **Prepare configuration**: By default, uses `config/general_config.yaml`, `config/model_config.yaml`, and `config/attacks/*.yaml`, `config/defenses/*.yaml`. 4. **Run the full pipeline** (generate test cases → responses → evaluation): - `python run_pipeline.py --config config/general_config.yaml --full` @@ -144,10 +154,10 @@ Depending on the attack/defense methods used, the following additional configura - **Generate test cases only**: - `python run_pipeline.py --config config/general_config.yaml --stage test_case_generation` - **Generate model responses from specified test_case JSON**: - - `python run_pipeline.py --config config/general_config.yaml --stage response_generation --test-cases-file /abs/path/to/test_cases.json` + - `python run_pipeline.py --config config/general_config.yaml --stage response_generation --test-cases-file /abs/path/to/test_cases.jsonl` - Equivalent: Fill in the file path in `config.response_generation.input_test_cases`. - **Generate evaluation results from specified response JSON**: - - `python run_pipeline.py --config config/general_config.yaml --stage evaluation --input-file /abs/path/to/responses.json` + - `python run_pipeline.py --config config/general_config.yaml --stage evaluation --input-file /abs/path/to/responses.jsonl` - Equivalent: Fill in the file path in `config.evaluation.input_responses`. - **Run combined stages**: - `python run_pipeline.py --config config/general_config.yaml --stages test_case_generation,response_generation,evaluation` @@ -156,7 +166,7 @@ Depending on the attack/defense methods used, the following additional configura - These white-box attack methods use MiniGPT-4 as the target model when generating test cases, but **do not automatically generate responses**. - After generating test cases, you need to **run MiniGPT-4 separately** to process the attack-generated test case data and generate responses. - Specific steps: - 1. First run the test case generation stage to obtain the attack-generated test case files (located at `output/test_cases///...json`) + 1. First run the test case generation stage to obtain the attack-generated test case files (located at `output/test_cases///...jsonl`) 2. Use the MiniGPT-4 model separately to process these test cases and generate responses (you need to write your own script or use the framework's response_generation stage, but ensure the MiniGPT-4 model is properly configured) 3. After saving the generated responses as a JSON file, run the evaluation stage @@ -173,8 +183,8 @@ Depending on the attack/defense methods used, the following additional configura - **Attack/Defense configuration**: `config/attacks/*.yaml`, `config/defenses/*.yaml` - Fill in the parameters required by each method in `parameters`, which can be overridden in `general_config.yaml` using `attack_params` / `defense_params`. - **Directly specify input files**: - - Response stage: `response_generation.input_test_cases: /abs/path/to/test_cases.json` - - Evaluation stage: `evaluation.input_responses: /abs/path/to/responses.json` + - Response stage: `response_generation.input_test_cases: /abs/path/to/test_cases.jsonl` + - Evaluation stage: `evaluation.input_responses: /abs/path/to/responses.jsonl` ## 📊 Input/Output Formats for Each Stage @@ -221,7 +231,7 @@ Depending on the attack/defense methods used, the following additional configura **Input Format** (Test Cases JSON): - Same format as the output of Stage 1 -- Can be loaded from `output/test_cases///...json` or specified via `response_generation.input_test_cases` +- Can be loaded from `output/test_cases///...jsonl` or specified via `response_generation.input_test_cases` **Output Format** (Model Responses JSON): ```json @@ -251,7 +261,7 @@ Depending on the attack/defense methods used, the following additional configura **Input Format** (Model Responses JSON): - Same format as the output of Stage 2 -- Can be loaded from `output/responses//attack__model_.json` or specified via `evaluation.input_responses` +- Can be loaded from `output/responses//attack__model_.jsonl` or specified via `evaluation.input_responses` **Output Format** (Evaluation Results JSON): ```json @@ -288,9 +298,9 @@ Depending on the attack/defense methods used, the following additional configura ## 📁 Output Directory Structure (Default `output/`) -- `test_cases///...json`: Test cases -- `responses//attack__model_.json`: Model responses -- `evaluations/attack__model__defense_.json`: Evaluation results +- `test_cases///...jsonl`: Test cases +- `responses//attack__model_.jsonl`: Model responses +- `evaluations/attack__model__defense__evaluator_.jsonl`: Evaluation results ## 🔧 Custom Components (Requires Registry Mapping Declaration) @@ -299,8 +309,8 @@ When adding new components, please: - Attack: `core.base_classes.BaseAttack` - Defense: `core.base_classes.BaseDefense` - Model: `models.base_model.BaseModel` - - Evaluator: `evaluators.base_evaluator.BaseEvaluator` -2. **Register**: Edit the mapping table in `core/registry_init.py`, add `name: ("module.path", "ClassName")`. + - Evaluator: `core.base_classes.BaseEvaluator` +2. **Register**: Edit `config/plugins.yaml`, add `name: [ "module.path", "ClassName" ]` under the appropriate section. 3. **Configure**: Add configuration files for the new component - Attack: `config/attacks/.yaml`, and enable it in the `test_case_generation.attacks` list in `general_config.yaml`. - Defense: `config/defenses/.yaml`, and enable it in the `response_generation.defenses` list in `general_config.yaml`. @@ -357,9 +367,9 @@ When adding new components, please: ## ❓ FAQ - **How to re-run evaluation only?** - Directly use existing responses file: `--stage evaluation --input-file /abs/path/to/responses.json`. + Directly use existing responses file: `--stage evaluation --input-file /abs/path/to/responses.jsonl`. - **How to skip defense?** Keep only `"None"` in `response_generation.defenses`. - **API key reading?** Fill directly in `config/model_config.yaml`; -- **Pre-development check**: Ensure new components are mapped in `registry_init.py` and have corresponding configuration files. +- **Pre-development check**: Ensure new components are mapped in `config/plugins.yaml` and have corresponding configuration files. diff --git a/attacks/bap/attack.py b/attacks/bap/attack.py index f2ad6fd..f51bbce 100644 --- a/attacks/bap/attack.py +++ b/attacks/bap/attack.py @@ -14,7 +14,7 @@ import importlib import logging from core.data_formats import TestCase -from core.unified_registry import BaseAttack +from core.base_classes import BaseAttack from core.unified_registry import UNIFIED_REGISTRY from config.config_loader import get_model_config @@ -339,14 +339,10 @@ Feel free to use these strategies or come up with your own. At the end of your r # CoT optimize prompt optimized_prompt, cot_log = self._optimize_prompt_with_cot(original_prompt) - return TestCase( - test_case_id=case_id, - prompt=optimized_prompt, - image_path=vap_save_path, - metadata={ - "attack_method": "bap", - "original_prompt": original_prompt, - "jailbreak_prompt": optimized_prompt, - "jailbreak_image_path": vap_save_path, - }, + return self.create_test_case( + case_id=case_id, + jailbreak_prompt=optimized_prompt, + jailbreak_image_path=str(vap_save_path), + original_prompt=original_prompt, + original_image_path=str(image_path), ) diff --git a/attacks/cs_dj/attack.py b/attacks/cs_dj/attack.py index 409204a..61f7d57 100644 --- a/attacks/cs_dj/attack.py +++ b/attacks/cs_dj/attack.py @@ -20,7 +20,8 @@ from sentence_transformers import SentenceTransformer, util from core.data_formats import TestCase -from core.unified_registry import BaseAttack, UNIFIED_REGISTRY +from core.base_classes import BaseAttack +from core.unified_registry import UNIFIED_REGISTRY from config.config_loader import get_model_config @@ -500,15 +501,10 @@ class CSDJAttack(BaseAttack): img_path = str(result_root / f"{key}.jpg") grid.save(img_path) - test_case = TestCase( - test_case_id=str(case_id), - image_path=str(img_path), - prompt=self.adv_prompt, - metadata={ - "attack_method": "cs_dj", - "original_prompt": original_prompt, - "jailbreak_prompt": self.adv_prompt, - "jailbreak_image_path": str(img_path), - }, + return self.create_test_case( + case_id=case_id, + jailbreak_prompt=self.adv_prompt, + jailbreak_image_path=str(img_path), + original_prompt=original_prompt, + original_image_path=str(image_path), ) - return test_case diff --git a/attacks/figstep/attack.py b/attacks/figstep/attack.py index 6fb05c1..da4610f 100644 --- a/attacks/figstep/attack.py +++ b/attacks/figstep/attack.py @@ -9,7 +9,7 @@ import textwrap from PIL import Image, ImageFont, ImageDraw from core.data_formats import TestCase -from core.unified_registry import BaseAttack +from core.base_classes import BaseAttack # ===================== Font and Drawing Tools ===================== @@ -127,15 +127,10 @@ class FigStepAttack(BaseAttack): img_path = self.output_image_dir / f"{case_id}.png" img.save(img_path) ## Save image - test_case = TestCase( - test_case_id=str(case_id), - image_path=str(img_path), - prompt=self.cfg.attack_prompt, - metadata={ - "attack_method": "figstep", - "original_prompt": original_prompt, - "jailbreak_prompt": self.cfg.attack_prompt, - "jailbreak_image_path": str(img_path), - }, + return self.create_test_case( + case_id=case_id, + jailbreak_prompt=self.cfg.attack_prompt, + jailbreak_image_path=str(img_path), + original_prompt=original_prompt, + original_image_path=str(image_path), ) - return test_case diff --git a/attacks/hades/attack.py b/attacks/hades/attack.py index fa0a23c..9e991c6 100644 --- a/attacks/hades/attack.py +++ b/attacks/hades/attack.py @@ -11,7 +11,8 @@ import glob from PIL import Image, ImageFont, ImageDraw from core.data_formats import TestCase -from core.unified_registry import BaseAttack, BaseModel, UNIFIED_REGISTRY +from core.base_classes import BaseAttack, BaseModel +from core.unified_registry import UNIFIED_REGISTRY from config.config_loader import get_model_config from .utils import ( generate_caption, @@ -182,17 +183,10 @@ class HadesAttack(BaseAttack): image_save_path = output_dir / image_filename new_image.save(image_save_path) - # Create test case - test_case = TestCase( - test_case_id=str(case_id), - image_path=str(image_save_path), - prompt=adv_prompt, - metadata={ - "attack_method": "hades", - "original_prompt": original_prompt, - "jailbreak_prompt": adv_prompt, - "jailbreak_image_path": str(image_save_path), - }, + return self.create_test_case( + case_id=case_id, + jailbreak_prompt=adv_prompt, + jailbreak_image_path=str(image_save_path), + original_prompt=original_prompt, + original_image_path=str(image_path), ) - - return test_case diff --git a/attacks/himrd/attack.py b/attacks/himrd/attack.py index bb0a784..638f3e2 100644 --- a/attacks/himrd/attack.py +++ b/attacks/himrd/attack.py @@ -25,7 +25,8 @@ from .attack_utils import ( get_gpt_response, ) from core.data_formats import TestCase -from core.unified_registry import BaseAttack, BaseModel, UNIFIED_REGISTRY +from core.base_classes import BaseAttack, BaseModel +from core.unified_registry import UNIFIED_REGISTRY from config.config_loader import get_model_config @@ -1408,16 +1409,10 @@ class HIMRDAttack(BaseAttack): else [] ) - test_case = TestCase( - test_case_id=case_id, - image_path=images[0], - prompt=prompts_list[0], - metadata={ - "attack_method": "himrd", - "original_prompt": original_prompt, - "jailbreak_prompt": prompts_list[0], - "jailbreak_image_path": images[0], - }, + return self.create_test_case( + case_id=case_id, + jailbreak_prompt=prompts_list[0], + jailbreak_image_path=str(images[0]), + original_prompt=original_prompt, + original_image_path=str(image_path), ) - - return test_case diff --git a/attacks/imgjp/attack.py b/attacks/imgjp/attack.py index ee63856..0e2d71d 100644 --- a/attacks/imgjp/attack.py +++ b/attacks/imgjp/attack.py @@ -17,7 +17,7 @@ from .utils.torchattacks.attacks.pgdl2 import * import numpy as np from core.data_formats import TestCase -from core.unified_registry import BaseAttack +from core.base_classes import BaseAttack from dataclasses import dataclass @@ -96,14 +96,10 @@ class ImgJPAttack(BaseAttack): self.adv_image_path = self.output_image_dir / f"imgjp_adv.jpg" self.adv_image.save(self.adv_image_path) - return TestCase( - test_case_id=case_id, - prompt=original_prompt, - image_path=self.adv_image_path, - metadata={ - "attack_method": "imgjp", - "original_prompt": original_prompt, - "jailbreak_prompt": original_prompt, - "jailbreak_image_path": self.adv_image_path, - }, + return self.create_test_case( + case_id=case_id, + jailbreak_prompt=original_prompt, + jailbreak_image_path=str(self.adv_image_path), + original_prompt=original_prompt, + original_image_path=str(image_path), ) diff --git a/attacks/jood/attack.py b/attacks/jood/attack.py index e9e08e9..e04a544 100644 --- a/attacks/jood/attack.py +++ b/attacks/jood/attack.py @@ -7,7 +7,7 @@ from typing import Any, Dict, Iterable, List, Optional, Tuple, Union from pathlib import Path from core.data_formats import TestCase -from core.unified_registry import BaseAttack +from core.base_classes import BaseAttack from .utils.mixaug import ( mixup_images, @@ -142,15 +142,10 @@ class JOODAttack(BaseAttack): mixed_image_pil.save(img_path) ## Save image adv_prompt = _parse_scenario2prompts(prompt, cfg.aug) - test_case = TestCase( - test_case_id=str(case_id), - image_path=str(img_path), - prompt=adv_prompt, - metadata={ - "attack_method": "jood", - "original_prompt": original_prompt, - "jailbreak_prompt": adv_prompt, - "jailbreak_image_path": str(img_path), - }, + return self.create_test_case( + case_id=case_id, + jailbreak_prompt=adv_prompt, + jailbreak_image_path=str(img_path), + original_prompt=original_prompt, + original_image_path=str(image_path), ) - return test_case diff --git a/attacks/mml/attack.py b/attacks/mml/attack.py index 4b02419..24ad6cc 100644 --- a/attacks/mml/attack.py +++ b/attacks/mml/attack.py @@ -13,7 +13,7 @@ from PIL import Image, ImageFont, ImageDraw import textwrap from core.data_formats import TestCase -from core.unified_registry import BaseAttack +from core.base_classes import BaseAttack from . import utils @@ -236,17 +236,11 @@ class MMLAttack(BaseAttack): "images_rotate", "images_rotate" ).format(scrambled_prompt) - test_case = TestCase( - test_case_id=case_id, - image_path=img_path, - prompt=adv_prompt, - metadata={ - "attack_method": "mml", - "original_prompt": original_prompt, - "jailbreak_prompt": adv_prompt, - "jailbreak_image_path": img_path, - "aug_type": self.cfg.aug_type, - }, + return self.create_test_case( + case_id=case_id, + jailbreak_prompt=adv_prompt, + jailbreak_image_path=str(img_path), + original_prompt=original_prompt, + original_image_path=str(image_path), + metadata={"aug_type": self.cfg.aug_type}, ) - - return test_case diff --git a/attacks/pba/attack.py b/attacks/pba/attack.py index 72d710b..62d2c97 100644 --- a/attacks/pba/attack.py +++ b/attacks/pba/attack.py @@ -14,7 +14,7 @@ from tqdm import tqdm from .image_random import Image_Attacker as Image_Attacker_Random from .text_random import Text_Attacker as Text_Attacker_Random from .minigpt_utils import prompt_wrapper, generator -from core.unified_registry import BaseAttack +from core.base_classes import BaseAttack from core.data_formats import TestCase from .utils import get_perspective_score @@ -439,14 +439,10 @@ class PBAAttack(BaseAttack): ## Joint optimization self._optimization() - return TestCase( - test_case_id=case_id, - image_path=self.save_path, - prompt=original_prompt, - metadata={ - "attack_method": "pba", - "original_prompt": original_prompt, - "jailbreak_prompt": original_prompt, - "jailbreak_image_path": self.save_path, - }, + return self.create_test_case( + case_id=case_id, + jailbreak_prompt=original_prompt, + jailbreak_image_path=str(self.save_path), + original_prompt=original_prompt, + original_image_path=str(image_path), ) diff --git a/attacks/qr/attack.py b/attacks/qr/attack.py index 37c2fb8..3e76c93 100644 --- a/attacks/qr/attack.py +++ b/attacks/qr/attack.py @@ -16,7 +16,8 @@ import threading import logging from core.data_formats import TestCase -from core.unified_registry import BaseAttack, UNIFIED_REGISTRY +from core.base_classes import BaseAttack +from core.unified_registry import UNIFIED_REGISTRY from config.config_loader import get_model_config from diffusers import StableDiffusion3Pipeline @@ -475,18 +476,11 @@ class QRAttack(BaseAttack): img_path = self.output_image_dir / img_name canvas.save(img_path) - # Create test case - test_case = TestCase( - test_case_id=str(case_id), - image_path=str(img_path), - prompt=rephrased, - metadata={ - "attack_method": "qr", - "original_prompt": original_prompt, - "key_phrase": key_phrase, - "jailbreak_prompt": rephrased, - "jailbreak_image_path": str(img_path), - }, + return self.create_test_case( + case_id=case_id, + jailbreak_prompt=rephrased, + jailbreak_image_path=str(img_path), + original_prompt=original_prompt, + original_image_path=str(image_path), + metadata={"key_phrase": key_phrase}, ) - - return test_case diff --git a/attacks/si/attack.py b/attacks/si/attack.py index 16244b2..bcf57bb 100644 --- a/attacks/si/attack.py +++ b/attacks/si/attack.py @@ -19,7 +19,8 @@ import io import base64 from core.data_formats import TestCase -from core.unified_registry import BaseAttack, UNIFIED_REGISTRY +from core.base_classes import BaseAttack +from core.unified_registry import UNIFIED_REGISTRY from config.config_loader import get_model_config from .judge import judger @@ -277,13 +278,13 @@ class SIAttack(BaseAttack): if target_model is None or processor is None: self.logger.error("Target model loading failed, cannot generate test case") # Create error test case - return TestCase( - test_case_id=str(case_id), - image_path=image_path, - prompt=original_prompt, + return self.create_test_case( + case_id=case_id, + jailbreak_prompt=original_prompt, + jailbreak_image_path=str(image_path), + original_prompt=original_prompt, + original_image_path=str(image_path), metadata={ - "attack_method": "si", - "original_prompt": original_prompt, "error": "Target model loading failed", "target_model": self.cfg.target_model_name, }, @@ -341,17 +342,11 @@ class SIAttack(BaseAttack): image_per.convert("RGB").save(save_img_path, format="PNG") # Create test case - test_case = TestCase( - test_case_id=str(case_id), - image_path=str(save_img_path), - prompt=final_question, - metadata={ - "attack_method": "si", - "original_prompt": original_prompt, - "jailbreak_prompt": final_question, - "jailbreak_image_path": str(save_img_path), - "target_model": self.cfg.target_model_name, - }, + return self.create_test_case( + case_id=case_id, + jailbreak_prompt=final_question, + jailbreak_image_path=str(save_img_path), + original_prompt=original_prompt, + original_image_path=str(image_path), + metadata={"target_model": self.cfg.target_model_name}, ) - - return test_case diff --git a/attacks/umk/attack.py b/attacks/umk/attack.py index 7f432d5..0a2db6d 100644 --- a/attacks/umk/attack.py +++ b/attacks/umk/attack.py @@ -91,14 +91,11 @@ class UMKAttack(BaseAttack): save_path = images_dir / f"umk_{case_id}.jpg" save_image(adv_t_save, str(save_path)) - return TestCase( - test_case_id=case_id, - image_path=save_path, - prompt=original_prompt + text_suffix, - metadata={ - "attack_method": "umk", - "original_prompt": original_prompt, - "jailbreak_prompt": original_prompt + text_suffix, - "jailbreak_image_path": save_path, - }, + jailbreak_prompt = original_prompt + text_suffix + return self.create_test_case( + case_id=case_id, + jailbreak_prompt=jailbreak_prompt, + jailbreak_image_path=str(save_path), + original_prompt=original_prompt, + original_image_path=str(image_path), ) diff --git a/attacks/viscra/attack.py b/attacks/viscra/attack.py index 9c2999c..8d94161 100644 --- a/attacks/viscra/attack.py +++ b/attacks/viscra/attack.py @@ -13,7 +13,7 @@ import matplotlib.pyplot as plt import torch from transformers import AutoProcessor from transformers import Qwen2_5_VLForConditionalGeneration -from core.unified_registry import BaseAttack +from core.base_classes import BaseAttack from core.data_formats import TestCase from dataclasses import dataclass @@ -435,14 +435,10 @@ class VisCRAAttack(BaseAttack): # Modify query modified_query = self._modify_query(query, self.attack_type) - return TestCase( - test_case_id=case_id, - image_path=attack_image, - prompt=modified_query, - metadata={ - "attack_method": "viscra", - "original_prompt": query, - "jailbreak_prompt": modified_query, - "jailbreak_image_path": attack_image, - }, + return self.create_test_case( + case_id=case_id, + jailbreak_prompt=modified_query, + jailbreak_image_path=str(attack_image), + original_prompt=query, + original_image_path=str(image_path), ) diff --git a/attacks/visual_adv/attack.py b/attacks/visual_adv/attack.py index ce3b0f8..7a3ba6b 100644 --- a/attacks/visual_adv/attack.py +++ b/attacks/visual_adv/attack.py @@ -12,7 +12,7 @@ from torchvision.utils import save_image from PIL import Image from tqdm import tqdm -from core.unified_registry import BaseAttack +from core.base_classes import BaseAttack from core.data_formats import TestCase from dataclasses import dataclass @@ -273,14 +273,10 @@ class VisualAdvAttack(BaseAttack): self.save_path = self.output_image_dir / "advimg.jpg" save_image(adv_t_save, str(self.save_path)) - return TestCase( - test_case_id=case_id, - image_path=self.save_path, - prompt=original_prompt, - metadata={ - "attack_method": "visual_adv", - "original_prompt": original_prompt, - "jailbreak_prompt": original_prompt, - "jailbreak_image_path": self.save_path, - }, + return self.create_test_case( + case_id=case_id, + jailbreak_prompt=original_prompt, + jailbreak_image_path=str(self.save_path), + original_prompt=original_prompt, + original_image_path=str(image_path), ) diff --git a/config/plugins.yaml b/config/plugins.yaml new file mode 100644 index 0000000..f00fb7d --- /dev/null +++ b/config/plugins.yaml @@ -0,0 +1,55 @@ +version: 1 + +# Canonical plugin registry for OmniSafeBench-MM. +# Each entry maps a registry name -> [python_module_path, class_name] +plugins: + attacks: + figstep: [attacks.figstep.attack, FigStepAttack] + cs_dj: [attacks.cs_dj.attack, CSDJAttack] + hades: [attacks.hades.attack, HadesAttack] + qr: [attacks.qr.attack, QRAttack] + mml: [attacks.mml.attack, MMLAttack] + si: [attacks.si.attack, SIAttack] + jood: [attacks.jood.attack, JOODAttack] + himrd: [attacks.himrd.attack, HIMRDAttack] + bap: [attacks.bap.attack, BAPAttack] + visual_adv: [attacks.visual_adv.attack, VisualAdvAttack] + viscra: [attacks.viscra.attack, VisCRAAttack] + umk: [attacks.umk.attack, UMKAttack] + pba: [attacks.pba.attack, PBAAttack] + imgjp: [attacks.imgjp.attack, ImgJPAttack] + + models: + openai: [models.openai_model, OpenAIModel] + google: [models.google_model, GoogleModel] + anthropic: [models.anthropic_model, AnthropicModel] + qwen: [models.qwen_model, QwenModel] + doubao: [models.doubao_model, DoubaoModel] + vllm: [models.vllm_model, VLLMModel] + mistral: [models.mistral_model, MistralModel] + any: [models.openai_model, OpenAIModel] + + defenses: + adashield: [defenses.adashield, AdaShieldDefense] + dps: [defenses.dps, DPSDefense] + ecso: [defenses.ecso, ECSODefense] + jailguard: [defenses.jailguard, JailGuardDefense] + llavaguard: [defenses.llavaguard, LlavaGuardDefense] + qguard: [defenses.qguard, QGuardDefense] + shieldlm: [defenses.shieldlm, ShieldLMDefense] + uniguard: [defenses.uniguard, UniguardDefense] + cider: [defenses.cider, CIDERDefense] + mllm_protector: [defenses.mllm_protector, MLLMProtectorDefense] + llama_guard_3: [defenses.llama_guard_3, LlamaGuard3Defense] + hiddendetect: [defenses.hiddendetect, HiddenDetectDefense] + guardreasoner_vl: [defenses.guardreasoner_vl, GuardReasonerVLDefense] + llama_guard_4: [defenses.llama_guard_4, LlamaGuard4Defense] + vlguard: [defenses.vlguard, VLGuardDefense] + spa_vl: [defenses.spa_vl, SPAVLDefense] + saferlhf_v: [defenses.saferlhf_v, SafeRLHFVDefense] + coca: [defenses.coca, CoCADefense] + + evaluators: + default_judge: [evaluators.default_judge, DefaultJudge] + + diff --git a/core/base_classes.py b/core/base_classes.py index a1ee1a6..28880b3 100644 --- a/core/base_classes.py +++ b/core/base_classes.py @@ -27,6 +27,10 @@ class BaseComponent(ABC): """ self.config = config or {} self.name = self.__class__.__name__ + # Registry name injected by UnifiedRegistry (canonical component id, e.g. "llama_guard_3") + self.registry_name = ( + config.get("_registry_name") if isinstance(config, dict) else None + ) # Initialize logger (using the module name where component is located) self.logger = logging.getLogger(self.__class__.__module__) @@ -77,6 +81,13 @@ class BaseComponent(ABC): class BaseAttack(BaseComponent, ABC): """Attack method base class (enhanced version)""" + # Standard Metadata Keys + META_KEY_ATTACK_METHOD = "attack_method" + META_KEY_ORIGINAL_PROMPT = "original_prompt" + META_KEY_JAILBREAK_PROMPT = "jailbreak_prompt" + META_KEY_JAILBREAK_IMAGE = "jailbreak_image_path" + META_KEY_TARGET_MODEL = "target_model" + def __init__( self, config: Dict[str, Any] = None, output_image_dir: Optional[str] = None ): @@ -135,6 +146,61 @@ class BaseAttack(BaseComponent, ABC): """ pass + def create_test_case( + self, + case_id: str, + jailbreak_prompt: str, + jailbreak_image_path: str, + original_prompt: str, + original_image_path: str = None, + metadata: Dict[str, Any] = None, + **kwargs, + ) -> TestCase: + """ + Create a standardized `TestCase` with populated metadata. + + Args: + case_id: Unique test case ID + jailbreak_prompt: The generated adversarial prompt + jailbreak_image_path: Path to the generated adversarial image + original_prompt: The original harmful prompt + original_image_path: Path to the original image (optional) + metadata: Additional metadata + **kwargs: Other parameters to store in metadata + + Returns: + TestCase object + """ + attack_method = ( + getattr(self, "registry_name", None) + or getattr(self, "name", None) + or self.__class__.__name__.replace("Attack", "").lower() + ) + + base_metadata = { + self.META_KEY_ATTACK_METHOD: attack_method, + self.META_KEY_ORIGINAL_PROMPT: original_prompt, + self.META_KEY_JAILBREAK_PROMPT: jailbreak_prompt, + self.META_KEY_JAILBREAK_IMAGE: jailbreak_image_path, + } + + if hasattr(self, "cfg") and hasattr(self.cfg, "target_model_name"): + if self.cfg.target_model_name: + base_metadata[self.META_KEY_TARGET_MODEL] = self.cfg.target_model_name + + if original_image_path: + base_metadata["original_image_path"] = original_image_path + + # Merge additional metadata + final_metadata = {**base_metadata, **(metadata or {}), **kwargs} + + return TestCase( + test_case_id=str(case_id), + prompt=jailbreak_prompt, + image_path=jailbreak_image_path, + metadata=final_metadata, + ) + def __str__(self): return f"{self.name}(config={self.config})" @@ -173,6 +239,13 @@ class BaseModel(BaseComponent, ABC): class BaseDefense(BaseComponent, ABC): """Defense method base class (enhanced version)""" + # Metadata keys + META_KEY_SHOULD_BLOCK = "should_return_default" + META_KEY_BLOCKED = "defense_blocked" + META_KEY_DEFAULT_RESP = "default_response" + META_KEY_GENERATED_RESP = "defense_generated_response" + META_KEY_GENERATED_RESP_USED = "defense_generated_response_used" + def __init__(self, config: Dict[str, Any] = None): """ Initialize defense method @@ -219,6 +292,90 @@ class BaseDefense(BaseComponent, ABC): """ pass + def create_defended_case( + self, + test_case: TestCase, + defended_prompt: str = None, + defended_image_path: str = None, + metadata: Dict[str, Any] = None, + ) -> TestCase: + """ + Create a new `TestCase` after applying defense (common helper). + + This keeps downstream pipeline logic and metadata conventions consistent. + """ + defense_method = ( + getattr(self, "registry_name", None) + or getattr(self, "name", None) + or self.__class__.__name__.replace("Defense", "").lower() + ) + + return TestCase( + test_case_id=test_case.test_case_id, + prompt=defended_prompt, + image_path=defended_image_path, + metadata={ + **(test_case.metadata or {}), + "defense_method": defense_method, + "defended_prompt": defended_prompt, + "defended_image_path": defended_image_path, + **(metadata or {}), + }, + ) + + def block_input( + self, + test_case: TestCase, + default_response: str = "I'm sorry, I can't assist with that.", + defended_prompt: str = None, + defended_image_path: str = None, + metadata: Dict[str, Any] = None, + ) -> TestCase: + """ + Mark the test case as 'blocked' by the defense. + Pipeline will intercept this and return `default_response` immediately. + """ + defended_case = self.create_defended_case( + test_case=test_case, + defended_prompt=defended_prompt if defended_prompt is not None else test_case.prompt, + defended_image_path=defended_image_path + if defended_image_path is not None + else test_case.image_path, + metadata={ + **(metadata or {}), + self.META_KEY_SHOULD_BLOCK: True, + self.META_KEY_BLOCKED: True, + self.META_KEY_DEFAULT_RESP: default_response, + }, + ) + return defended_case + + def reply_directly( + self, + test_case: TestCase, + response_text: str, + defended_prompt: str = None, + defended_image_path: str = None, + metadata: Dict[str, Any] = None, + ) -> TestCase: + """ + Defense generates the response directly (bypassing the target model). + Pipeline will use `response_text` as the final output. + """ + defended_case = self.create_defended_case( + test_case=test_case, + defended_prompt=defended_prompt if defended_prompt is not None else test_case.prompt, + defended_image_path=defended_image_path + if defended_image_path is not None + else test_case.image_path, + metadata={ + **(metadata or {}), + self.META_KEY_GENERATED_RESP: response_text, + self.META_KEY_GENERATED_RESP_USED: True, + }, + ) + return defended_case + def __str__(self): return f"{self.name}(config={self.config})" @@ -237,7 +394,9 @@ class BaseEvaluator(BaseComponent, ABC): super().__init__(config) @abstractmethod - def evaluate(self, model_response: ModelResponse, **kwargs) -> EvaluationResult: + def evaluate_response( + self, model_response: ModelResponse, **kwargs + ) -> EvaluationResult: """ Evaluate model response diff --git a/core/registry_init.py b/core/registry_init.py deleted file mode 100644 index 6d8c68e..0000000 --- a/core/registry_init.py +++ /dev/null @@ -1,69 +0,0 @@ -def initialize_registry_lazy_imports(): - """Initialize lazy loading mappings to unified registry""" - - # Attack method lazy loading mappings - attack_lazy_imports = { - "figstep": ("attacks.figstep.attack", "FigStepAttack"), - "cs_dj": ("attacks.cs_dj.attack", "CSDJAttack"), - "hades": ("attacks.hades.attack", "HadesAttack"), - "qr": ("attacks.qr.attack", "QRAttack"), - "mml": ("attacks.mml.attack", "MMLAttack"), - "si": ("attacks.si.attack", "SIAttack"), - "jood": ("attacks.jood.attack", "JOODAttack"), - "himrd": ("attacks.himrd.attack", "HIMRDAttack"), - "bap": ("attacks.bap.attack", "BAPAttack"), - "visual_adv": ("attacks.visual_adv.attack", "VisualAdvAttack"), - "viscra": ("attacks.viscra.attack", "VisCRAAttack"), - "umk": ("attacks.umk.attack", "UMKAttack"), - "pba": ("attacks.pba.attack", "PBAAttack"), - "imgjp": ("attacks.imgjp.attack", "ImgJPAttack"), - } - - # Model lazy loading mappings - model_lazy_imports = { - "openai": ("models.openai_model", "OpenAIModel"), - "google": ("models.google_model", "GoogleModel"), - "anthropic": ("models.anthropic_model", "AnthropicModel"), - "qwen": ("models.qwen_model", "QwenModel"), - "doubao": ("models.doubao_model", "DoubaoModel"), - "vllm": ("models.vllm_model", "VLLMModel"), - "mistral": ("models.mistral_model", "MistralModel"), - "any": ( - "models.openai_model", - "OpenAIModel", - ), # "any" uses OpenAI-compatible API - } - - # Defense method lazy loading mappings - defense_lazy_imports = { - "adashield": ("defenses.adashield", "AdaShieldDefense"), - "dps": ("defenses.dps", "DPSDefense"), - "ecso": ("defenses.ecso", "ECSODefense"), - "jailguard": ("defenses.jailguard", "JailGuardDefense"), - "llavaguard": ("defenses.llavaguard", "LlavaGuardDefense"), - "qguard": ("defenses.qguard", "QGuardDefense"), - "shieldlm": ("defenses.shieldlm", "ShieldLMDefense"), - "uniguard": ("defenses.uniguard", "UniguardDefense"), - "cider": ("defenses.cider", "CIDERDefense"), - "mllm_protector": ("defenses.mllm_protector", "MLLMProtectorDefense"), - "llama_guard_3": ("defenses.llama_guard_3", "LlamaGuard3Defense"), - "hiddendetect": ("defenses.hiddendetect", "HiddenDetectDefense"), - "guardreasoner_vl": ("defenses.guardreasoner_vl", "GuardReasonerVLDefense"), - "llama_guard_4": ("defenses.llama_guard_4", "LlamaGuard4Defense"), - "vlguard": ("defenses.vlguard", "VLGuardDefense"), - "spa_vl": ("defenses.spa_vl", "SPAVLDefense"), - "saferlhf_v": ("defenses.saferlhf_v", "SafeRLHFVDefense"), - "coca": ("defenses.coca", "CoCADefense"), - } - - # Evaluator lazy loading mappings - evaluator_lazy_imports = { - "default_judge": ("evaluators.default_judge", "DefaultJudge"), - } - - return { - "attacks": attack_lazy_imports, - "models": model_lazy_imports, - "defenses": defense_lazy_imports, - "evaluators": evaluator_lazy_imports, - } diff --git a/core/unified_registry.py b/core/unified_registry.py index d541f57..2eb980e 100644 --- a/core/unified_registry.py +++ b/core/unified_registry.py @@ -2,12 +2,19 @@ Unified registry system - merged version, removed redundant code """ -from typing import Dict, Any, Optional, Type, List +from typing import Dict, Any, Optional, Type, List, TYPE_CHECKING import logging import importlib from typing import Set +from pathlib import Path +import yaml -from .base_classes import BaseAttack, BaseModel, BaseDefense, BaseEvaluator +# NOTE: +# - `core.base_classes` is the canonical place for component contracts (BaseAttack/BaseModel/BaseDefense/BaseEvaluator). +# - This registry module should NOT re-export Base* at runtime; keep imports here type-checking only. + +if TYPE_CHECKING: + from .base_classes import BaseAttack, BaseModel, BaseDefense, BaseEvaluator class UnifiedRegistry: @@ -22,6 +29,49 @@ class UnifiedRegistry: self.logger = logging.getLogger(__name__) # Imported module cache to avoid duplicate imports self._imported_modules: Set[str] = set() + # Lazy mapping cache (names only). This should not trigger heavy imports. + self._lazy_mappings_cache: Optional[Dict[str, Dict[str, Any]]] = None + + def _get_lazy_mappings(self) -> Dict[str, Dict[str, Any]]: + """Get lazy import mappings from config/plugins.yaml (cached).""" + if self._lazy_mappings_cache is not None: + return self._lazy_mappings_cache + self._lazy_mappings_cache = self._load_plugins_yaml() + return self._lazy_mappings_cache + + def _load_plugins_yaml(self) -> Dict[str, Dict[str, Any]]: + """Load plugin mappings from `config/plugins.yaml`.""" + # Repo root = parent of `core/` + repo_root = Path(__file__).resolve().parent.parent + plugins_file = repo_root / "config" / "plugins.yaml" + if not plugins_file.exists(): + raise FileNotFoundError( + f"config/plugins.yaml not found at {plugins_file}. This project requires config/plugins.yaml." + ) + + with open(plugins_file, "r", encoding="utf-8") as f: + data = yaml.safe_load(f) or {} + + plugins = data.get("plugins") or {} + mappings: Dict[str, Dict[str, Any]] = {} + for section in ("attacks", "models", "defenses", "evaluators"): + section_map = plugins.get(section) or {} + converted: Dict[str, Any] = {} + for name, spec in section_map.items(): + if not isinstance(spec, (list, tuple)) or len(spec) != 2: + raise ValueError( + f"Invalid plugin spec for {section}.{name}: expected [module, class], got {spec!r}" + ) + module_path, class_name = spec + converted[name] = (module_path, class_name) + mappings[section] = converted + + # Ensure all sections exist + mappings.setdefault("attacks", {}) + mappings.setdefault("models", {}) + mappings.setdefault("defenses", {}) + mappings.setdefault("evaluators", {}) + return mappings def register_attack(self, name: str, attack_class: Type["BaseAttack"]) -> None: """Register attack method""" @@ -64,11 +114,9 @@ class UnifiedRegistry: if name in self.attack_registry: return self.attack_registry[name] - # Get mapping information from registry_init + # Get mapping information from config/plugins.yaml try: - from .registry_init import initialize_registry_lazy_imports - - mappings = initialize_registry_lazy_imports() + mappings = self._get_lazy_mappings() if name in mappings["attacks"]: module_path, class_name = mappings["attacks"][name] module = importlib.import_module(module_path) @@ -98,11 +146,9 @@ class UnifiedRegistry: if name in self.model_registry: return self.model_registry[name] - # Get mapping information from registry_init + # Get mapping information from config/plugins.yaml try: - from .registry_init import initialize_registry_lazy_imports - - mappings = initialize_registry_lazy_imports() + mappings = self._get_lazy_mappings() if name in mappings["models"]: module_path, class_name = mappings["models"][name] module = importlib.import_module(module_path) @@ -127,11 +173,9 @@ class UnifiedRegistry: if name in self.defense_registry: return self.defense_registry[name] - # Get mapping information from registry_init + # Get mapping information from config/plugins.yaml try: - from .registry_init import initialize_registry_lazy_imports - - mappings = initialize_registry_lazy_imports() + mappings = self._get_lazy_mappings() if name in mappings["defenses"]: module_path, class_name = mappings["defenses"][name] module = importlib.import_module(module_path) @@ -160,11 +204,9 @@ class UnifiedRegistry: if name in self.evaluator_registry: return self.evaluator_registry[name] - # Get mapping information from registry_init + # Get mapping information from config/plugins.yaml try: - from .registry_init import initialize_registry_lazy_imports - - mappings = initialize_registry_lazy_imports() + mappings = self._get_lazy_mappings() if name in mappings["evaluators"]: module_path, class_name = mappings["evaluators"][name] module = importlib.import_module(module_path) @@ -199,7 +241,9 @@ class UnifiedRegistry: attack_class = self.get_attack(name) if attack_class: try: - return attack_class(config=config, output_image_dir=output_image_dir) + cfg = dict(config or {}) + cfg.setdefault("_registry_name", name) + return attack_class(config=cfg, output_image_dir=output_image_dir) except Exception as e: self.logger.error(f"Failed to create attack method '{name}': {e}") return None @@ -217,22 +261,29 @@ class UnifiedRegistry: if model_class: try: - # Extract parameters from configuration - model_name = config.get("model_name", name) - api_key = config.get("api_key", "") - base_url = config.get("base_url", None) + # Prefer factory method if provided by the model class. + # This removes provider-specific branching from the registry. + if hasattr(model_class, "from_config") and callable( + getattr(model_class, "from_config") + ): + # NOTE: models are instantiated via `models.*` which has its own config parsing. + # We keep `_registry_name` injection for symmetry, but it is optional for models. + cfg = dict(config or {}) + cfg.setdefault("_registry_name", name) + return model_class.from_config(name=name, config=cfg) - # Pass appropriate parameters based on provider type - if provider in ["openai", "qwen", "vllm", "any"]: - # OpenAI-compatible models need model_name, api_key, base_url + # ---- Backward-compatible fallback ---- + cfg = dict(config or {}) + cfg.setdefault("_registry_name", name) + model_name = cfg.get("model_name", name) + api_key = cfg.get("api_key", "") + base_url = cfg.get("base_url", None) + try: return model_class( model_name=model_name, api_key=api_key, base_url=base_url ) - elif provider in ["google", "anthropic", "doubao", "mistral"]: - # Other models only need model_name and api_key - return model_class(model_name=model_name, api_key=api_key) - else: - # Default case + except TypeError: + # Some providers don't take base_url return model_class(model_name=model_name, api_key=api_key) except Exception as e: self.logger.error( @@ -249,7 +300,9 @@ class UnifiedRegistry: defense_class = self.get_defense(name) if defense_class: try: - return defense_class(config=config) + cfg = dict(config or {}) + cfg.setdefault("_registry_name", name) + return defense_class(config=cfg) except Exception as e: self.logger.error(f"Failed to create defense method '{name}': {e}") return None @@ -261,42 +314,60 @@ class UnifiedRegistry: evaluator_class = self.get_evaluator(name) if evaluator_class: try: - return evaluator_class(config=config) + cfg = dict(config or {}) + cfg.setdefault("_registry_name", name) + return evaluator_class(config=cfg) except Exception as e: self.logger.error(f"Failed to create evaluator '{name}': {e}") return None def list_attacks(self) -> List[str]: """List all registered attack methods""" - return sorted(self.attack_registry.keys()) + names = set(self.attack_registry.keys()) + names.update(self._get_lazy_mappings().get("attacks", {}).keys()) + return sorted(names) def list_models(self) -> List[str]: """List all registered models""" - return sorted(self.model_registry.keys()) + names = set(self.model_registry.keys()) + names.update(self._get_lazy_mappings().get("models", {}).keys()) + return sorted(names) def list_defenses(self) -> List[str]: """List all registered defense methods""" - return sorted(self.defense_registry.keys()) + names = set(self.defense_registry.keys()) + names.update(self._get_lazy_mappings().get("defenses", {}).keys()) + return sorted(names) def list_evaluators(self) -> List[str]: """List all registered evaluators""" - return sorted(self.evaluator_registry.keys()) + names = set(self.evaluator_registry.keys()) + names.update(self._get_lazy_mappings().get("evaluators", {}).keys()) + return sorted(names) def validate_attack(self, name: str) -> bool: """Validate if attack method exists""" - return name in self.attack_registry + if name in self.attack_registry: + return True + return name in self._get_lazy_mappings().get("attacks", {}) def validate_model(self, name: str) -> bool: """Validate if model exists""" - return name in self.model_registry + if name in self.model_registry: + return True + return name in self._get_lazy_mappings().get("models", {}) def validate_defense(self, name: str) -> bool: """Validate if defense method exists""" - return name in self.defense_registry + if name in self.defense_registry: + return True + return name in self._get_lazy_mappings().get("defenses", {}) def validate_evaluator(self, name: str) -> bool: """Validate if evaluator exists""" - return name in self.evaluator_registry + if name in self.evaluator_registry: + return True + return name in self._get_lazy_mappings().get("evaluators", {}) def get_component_summary(self) -> Dict[str, List[str]]: """Get component summary""" diff --git a/defenses/__init__.py b/defenses/__init__.py index 2e65603..668fc5b 100644 --- a/defenses/__init__.py +++ b/defenses/__init__.py @@ -3,7 +3,7 @@ New architecture defense methods package Supports dynamic import, missing defense methods are set to None """ -from .base_defense import BaseDefense +from core.base_classes import BaseDefense # Dynamically import all defense methods, set missing ones to None try: diff --git a/defenses/adashield.py b/defenses/adashield.py index 7253cc2..9f68257 100644 --- a/defenses/adashield.py +++ b/defenses/adashield.py @@ -1,4 +1,4 @@ -from .base_defense import BaseDefense +from core.base_classes import BaseDefense from core.data_formats import TestCase diff --git a/defenses/base_defense.py b/defenses/base_defense.py deleted file mode 100644 index 9820455..0000000 --- a/defenses/base_defense.py +++ /dev/null @@ -1,57 +0,0 @@ -""" -Defense method base class - new architecture implementation -Supports two defense modes: -1. Pre-processing defense: modify input before model generation (apply_defense) -2. Post-processing defense: modify output after model generation (post_process_response) -""" - -import logging -from abc import ABC, abstractmethod -from typing import Dict, Any, Optional, Tuple - -from core.data_formats import TestCase - - -class BaseDefense(ABC): - """Abstract base class for defense methods""" - - def __init__(self, config: Dict[str, Any] = None): - self.config = config or {} - self.logger = logging.getLogger(self.__class__.__name__) - self.name = self.__class__.__name__.replace("Defense", "").lower() - - @abstractmethod - def apply_defense(self, test_case: TestCase, **kwargs) -> TestCase: - """ - Apply defense method to test case (pre-processing) - Modify input before model generates response - - Args: - test_case: Original test case - **kwargs: Additional parameters - - Returns: - Test case after applying defense - """ - pass - - def create_defended_case( - self, - test_case: TestCase, - defended_prompt: str = None, - defended_image_path: str = None, - metadata: Dict = None, - ) -> TestCase: - """Create test case after applying defense""" - return TestCase( - test_case_id=test_case.test_case_id, - prompt=defended_prompt, - image_path=defended_image_path, - metadata={ - **test_case.metadata, - "defense_method": self.name, - "defended_prompt": defended_prompt, - "defended_image_path": defended_image_path, - **(metadata or {}), - }, - ) diff --git a/defenses/cider.py b/defenses/cider.py index b6ea455..92001e4 100644 --- a/defenses/cider.py +++ b/defenses/cider.py @@ -6,7 +6,7 @@ CIDER (Content Integrity Defense with Enhanced Robustness) defense method from typing import Dict, Any, Tuple, Optional from pathlib import Path -from .base_defense import BaseDefense +from core.base_classes import BaseDefense from core.data_formats import TestCase import argparse from .cider_utils import defence, QApair @@ -49,10 +49,15 @@ class CIDERDefense(BaseDefense): behav="CIDER_defense", ) - defended_case = self.create_defended_case( + if pair.refuse == True: + return self.block_input( + test_case, + defended_prompt=test_case.prompt, + defended_image_path=test_case.image_path, + ) + + return self.create_defended_case( test_case=test_case, defended_prompt=test_case.prompt, - defended_image_path=(test_case.image_path), - metadata={"should_return_default": pair.refuse == True}, + defended_image_path=test_case.image_path, ) - return defended_case diff --git a/defenses/coca.py b/defenses/coca.py index 84f428e..3c71f4e 100644 --- a/defenses/coca.py +++ b/defenses/coca.py @@ -4,7 +4,7 @@ CoCA defense method - new architecture implementation from typing import Dict, Any -from .base_defense import BaseDefense +from core.base_classes import BaseDefense from core.data_formats import TestCase import torch @@ -45,12 +45,4 @@ class CoCADefense(BaseDefense): alpha=alpha, max_new_tokens=1000, ) - - return TestCase( - test_case_id=test_case.test_case_id, - image_path=test_case.image_path, - prompt=test_case.prompt, - metadata={ - "defense_generated_response": output, - }, - ) + return self.reply_directly(test_case, response_text=output) diff --git a/defenses/dps.py b/defenses/dps.py index f1dabf0..a418fe8 100644 --- a/defenses/dps.py +++ b/defenses/dps.py @@ -1,5 +1,5 @@ from typing import Dict, Any -from .base_defense import BaseDefense +from core.base_classes import BaseDefense from core.data_formats import TestCase from core.unified_registry import UNIFIED_REGISTRY from .dps_utils import ( diff --git a/defenses/ecso.py b/defenses/ecso.py index 67417b5..40ed7e1 100644 --- a/defenses/ecso.py +++ b/defenses/ecso.py @@ -1,6 +1,6 @@ import logging from typing import Dict, Any, Tuple, Optional -from .base_defense import BaseDefense +from core.base_classes import BaseDefense from core.data_formats import TestCase from core.unified_registry import UNIFIED_REGISTRY from .utils import generate_output @@ -54,10 +54,4 @@ class ECSODefense(BaseDefense): output = generate_output(None, gen_qs, target_model) else: output = original_output - defended_case = self.create_defended_case( - test_case=test_case, - defended_prompt=test_case.prompt, - defended_image_path=test_case.image_path, - metadata={"defense_generated_response": output}, - ) - return defended_case + return self.reply_directly(test_case, response_text=output) diff --git a/defenses/guardreasoner_vl.py b/defenses/guardreasoner_vl.py index 57fc2c6..77942bb 100644 --- a/defenses/guardreasoner_vl.py +++ b/defenses/guardreasoner_vl.py @@ -3,7 +3,7 @@ import re from typing import Dict, Any, Tuple, Optional import re -from .base_defense import BaseDefense +from core.base_classes import BaseDefense from core.data_formats import TestCase from core.unified_registry import UNIFIED_REGISTRY @@ -74,10 +74,15 @@ None print(f"Error during extraction: {e}") request_value = "unharmful" - defended_case = self.create_defended_case( + if request_value == "harmful": + return self.block_input( + test_case, + defended_prompt=test_case.metadata.get("jailbreak_prompt", test_case.prompt), + defended_image_path=test_case.image_path, + ) + + return self.create_defended_case( test_case=test_case, - defended_prompt=test_case.metadata["jailbreak_prompt"], + defended_prompt=test_case.metadata.get("jailbreak_prompt", test_case.prompt), defended_image_path=test_case.image_path, - metadata={"should_return_default": request_value == "harmful"}, ) - return defended_case diff --git a/defenses/hiddendetect.py b/defenses/hiddendetect.py index 6027fa1..44c39e8 100644 --- a/defenses/hiddendetect.py +++ b/defenses/hiddendetect.py @@ -6,7 +6,7 @@ Implementation based on original HiddenDetect logic: uses specialized model to d import logging from typing import Dict, Any, Tuple, Optional -from .base_defense import BaseDefense +from core.base_classes import BaseDefense from core.data_formats import TestCase from .utils import score_and_generate_hiddendetect from llava.model.builder import load_pretrained_model @@ -84,14 +84,19 @@ class HiddenDetectDefense(BaseDefense): threshold=self.HIDDENDETECT_THRESHOLD, max_new_tokens=1000, ) + if output == "I'm sorry, I can't assist with that.": + return self.block_input( + test_case, + default_response=output, + defended_prompt=test_case.prompt, + defended_image_path=test_case.image_path, + ) - defended_case = self.create_defended_case( - test_case=test_case, + # If not blocked, we can directly return the generated safe output + return self.reply_directly( + test_case, + response_text=output, defended_prompt=test_case.prompt, defended_image_path=test_case.image_path, - metadata={ - "should_return_default": output - == "I'm sorry, I can't assist with that." - }, + metadata={"hiddendetect_score": score}, ) - return defended_case diff --git a/defenses/jailguard.py b/defenses/jailguard.py index 0fe5fb2..8a46869 100644 --- a/defenses/jailguard.py +++ b/defenses/jailguard.py @@ -10,7 +10,7 @@ from typing import Dict, Any, Tuple, Optional, List from PIL import Image from io import BytesIO from core.unified_registry import UNIFIED_REGISTRY -from .base_defense import BaseDefense +from core.base_classes import BaseDefense from core.data_formats import TestCase from .JailGuard_utils.utils import update_divergence, detect_attack from .JailGuard_utils.augmentations import ( @@ -116,11 +116,11 @@ class JailGuardDefense(BaseDefense): pass detection_result = detect_attack(max_div, jailbreak_keywords, threshold) - defended_case = self.create_defended_case( + if detection_result: + return self.block_input(test_case) + + return self.create_defended_case( test_case=test_case, defended_prompt=test_case.prompt, defended_image_path=test_case.image_path, - metadata={"should_return_default": detection_result == True}, ) - - return defended_case diff --git a/defenses/llama_guard_3.py b/defenses/llama_guard_3.py index 55f7b5c..eac2424 100644 --- a/defenses/llama_guard_3.py +++ b/defenses/llama_guard_3.py @@ -1,4 +1,4 @@ -from .base_defense import BaseDefense +from core.base_classes import BaseDefense from core.data_formats import TestCase from core.unified_registry import UNIFIED_REGISTRY from .utils import generate_output @@ -182,18 +182,6 @@ class LlamaGuard3Defense(BaseDefense): # Step 4: Make decision based on judgment result self.logger.debug(f"[{idx}] Step 4: Make decision based on judgment result") if "unsafe" in judge_response.lower(): - defended_case = self.create_defended_case( - test_case=test_case, - defended_prompt=test_case.prompt, - defended_image_path=test_case.image_path, - metadata={"should_return_default": True}, - ) + return self.block_input(test_case) else: - defended_case = self.create_defended_case( - test_case=test_case, - defended_prompt=test_case.prompt, - defended_image_path=test_case.image_path, - metadata={"defense_generated_response": original_output}, - ) - - return defended_case + return self.reply_directly(test_case, response_text=original_output) diff --git a/defenses/llama_guard_4.py b/defenses/llama_guard_4.py index d414e19..82d279c 100644 --- a/defenses/llama_guard_4.py +++ b/defenses/llama_guard_4.py @@ -7,7 +7,7 @@ import os import logging from typing import Dict, Any, Optional, Tuple, List -from .base_defense import BaseDefense +from core.base_classes import BaseDefense from core.data_formats import TestCase from core.unified_registry import UNIFIED_REGISTRY from .utils import generate_output @@ -37,11 +37,15 @@ class LlamaGuard4Defense(BaseDefense): max_new_tokens=10, ) # print("judge_response:", judge_response) + if "unsafe" in judge_response.lower(): + return self.block_input( + test_case, + defended_prompt=attack_prompt, + defended_image_path=attack_image_path, + ) - defended_case = self.create_defended_case( + return self.create_defended_case( test_case=test_case, defended_prompt=attack_prompt, defended_image_path=attack_image_path, - metadata={"should_return_default": "unsafe" in judge_response.lower()}, ) - return defended_case diff --git a/defenses/llavaguard.py b/defenses/llavaguard.py index aeb8954..acaf1a1 100644 --- a/defenses/llavaguard.py +++ b/defenses/llavaguard.py @@ -1,6 +1,6 @@ from typing import Dict, Any, Tuple, Optional, List -from .base_defense import BaseDefense +from core.base_classes import BaseDefense from core.data_formats import TestCase from core.unified_registry import UNIFIED_REGISTRY @@ -39,10 +39,15 @@ class LlavaGuardDefense(BaseDefense): max_new_tokens=500, top_p=0.95, ) - defended_case = self.create_defended_case( + if "unsafe" in judge_response.lower(): + return self.block_input( + test_case, + defended_prompt=test_case.metadata.get("jailbreak_prompt", test_case.prompt), + defended_image_path=test_case.image_path, + ) + + return self.create_defended_case( test_case=test_case, - defended_prompt=test_case.metadata["jailbreak_prompt"], + defended_prompt=test_case.metadata.get("jailbreak_prompt", test_case.prompt), defended_image_path=test_case.image_path, - metadata={"should_return_default": "unsafe" in judge_response.lower()}, ) - return defended_case diff --git a/defenses/mllm_protector.py b/defenses/mllm_protector.py index 393471e..be4ef00 100644 --- a/defenses/mllm_protector.py +++ b/defenses/mllm_protector.py @@ -9,7 +9,7 @@ import threading import time from typing import Dict, Any, Optional, Tuple -from .base_defense import BaseDefense +from core.base_classes import BaseDefense from core.data_formats import TestCase from .utils import generate_output from core.unified_registry import UNIFIED_REGISTRY @@ -119,7 +119,7 @@ class MLLMProtectorDefense(BaseDefense): defended_prompt=attack_prompt, defended_image_path=attack_image_path, metadata={ - "defense_generated_response": output, + self.META_KEY_GENERATED_RESP: output, "harm_score": harm_score, "detoxified": harm_score < 0.5, }, @@ -134,14 +134,13 @@ class MLLMProtectorDefense(BaseDefense): def _create_error_case(self, test_case: TestCase, error_msg: str) -> TestCase: """Create defense case for error situations""" - return self.create_defended_case( - test_case=test_case, + return self.block_input( + test_case, + default_response="I'm sorry, I can't assist with that.", defended_prompt=test_case.prompt, defended_image_path=test_case.image_path, metadata={ - "should_return_default": True, "defense_error": error_msg, - "default_response": "I'm sorry, I can't assist with that.", "error_time": time.time(), }, ) diff --git a/defenses/qguard.py b/defenses/qguard.py index ee7bdfa..ee9942f 100644 --- a/defenses/qguard.py +++ b/defenses/qguard.py @@ -5,7 +5,7 @@ QGuard defense method, safety protection specifically designed for Qwen models from typing import Dict, Any, Tuple, Optional, List -from .base_defense import BaseDefense +from core.base_classes import BaseDefense from core.data_formats import TestCase from .qguard_utils.vision import load_images_batch @@ -71,10 +71,15 @@ class QGuardDefense(BaseDefense): # print("judge_result:", judge_result["prediction"]) is_harmful = judge_result["prediction"] - defended_case = self.create_defended_case( + if is_harmful.strip().lower() == "harmful": + return self.block_input( + test_case, + defended_prompt=test_case.metadata.get("jailbreak_prompt", test_case.prompt), + defended_image_path=test_case.image_path, + ) + + return self.create_defended_case( test_case=test_case, - defended_prompt=test_case.metadata["jailbreak_prompt"], + defended_prompt=test_case.metadata.get("jailbreak_prompt", test_case.prompt), defended_image_path=test_case.image_path, - metadata={"should_return_default": is_harmful.strip().lower() == "harmful"}, ) - return defended_case diff --git a/defenses/shieldlm.py b/defenses/shieldlm.py index df7c927..1e98a6e 100644 --- a/defenses/shieldlm.py +++ b/defenses/shieldlm.py @@ -6,7 +6,7 @@ Migrated from original logic in generate_outputs.py import logging from typing import Dict, Any, Optional -from .base_defense import BaseDefense +from core.base_classes import BaseDefense from core.data_formats import TestCase from .utils import generate_output @@ -71,17 +71,6 @@ class ShieldLMDefense(BaseDefense): ) if "unsafe" in judge_result_2.lower(): - defended_case = self.create_defended_case( - test_case=test_case, - defended_prompt=test_case.prompt, - defended_image_path=test_case.image_path, - metadata={"should_return_default": True}, - ) + return self.block_input(test_case) else: - defended_case = self.create_defended_case( - test_case=test_case, - defended_prompt=test_case.prompt, - defended_image_path=test_case.image_path, - metadata={"defense_generated_response": original_output}, - ) - return defended_case + return self.reply_directly(test_case, response_text=original_output) diff --git a/defenses/uniguard.py b/defenses/uniguard.py index 1cb5b6b..11ee356 100644 --- a/defenses/uniguard.py +++ b/defenses/uniguard.py @@ -13,7 +13,7 @@ from typing import Dict, Any, Optional, Tuple from PIL import Image import numpy as np -from .base_defense import BaseDefense +from core.base_classes import BaseDefense from core.data_formats import TestCase from .utils import add_patch_to_image_universal diff --git a/defenses/vlguard.py b/defenses/vlguard.py index 2f89093..63e67df 100644 --- a/defenses/vlguard.py +++ b/defenses/vlguard.py @@ -3,7 +3,7 @@ import torch from typing import Dict, Any, Tuple, Optional from PIL import Image -from .base_defense import BaseDefense +from core.base_classes import BaseDefense from core.data_formats import TestCase import torch @@ -97,10 +97,4 @@ class VLGuardDefense(BaseDefense): # Decode output outputs = self.tokenizer.batch_decode(output_ids, skip_special_tokens=True)[0] - defended_case = self.create_defended_case( - test_case=test_case, - defended_prompt=test_case.prompt, - defended_image_path=test_case.image_path, - metadata={"defense_generated_response": outputs}, - ) - return defended_case + return self.reply_directly(test_case, response_text=outputs) diff --git a/evaluators/__init__.py b/evaluators/__init__.py index 1abf0cb..3dadcca 100644 --- a/evaluators/__init__.py +++ b/evaluators/__init__.py @@ -1,4 +1,4 @@ -from .base_evaluator import BaseEvaluator +from core.base_classes import BaseEvaluator from .default_judge import DefaultJudge __all__ = [ diff --git a/evaluators/base_evaluator.py b/evaluators/base_evaluator.py deleted file mode 100644 index a858eed..0000000 --- a/evaluators/base_evaluator.py +++ /dev/null @@ -1,34 +0,0 @@ -""" -Evaluator base class - new architecture implementation -""" - -import logging -from abc import ABC, abstractmethod -from typing import Dict, Any - -from core.data_formats import ModelResponse, EvaluationResult - - -class BaseEvaluator(ABC): - """Evaluator abstract base class""" - - def __init__(self, config: Dict[str, Any] = None): - self.config = config or {} - self.logger = logging.getLogger(self.__class__.__name__) - self.name = self.__class__.__name__.lower() - - @abstractmethod - def evaluate_response( - self, model_response: ModelResponse, **kwargs - ) -> EvaluationResult: - """ - Evaluate model response - - Args: - model_response: Model response - **kwargs: Additional parameters - - Returns: - Evaluation result - """ - pass diff --git a/evaluators/default_judge.py b/evaluators/default_judge.py index da57d5f..c9cc3a6 100644 --- a/evaluators/default_judge.py +++ b/evaluators/default_judge.py @@ -5,7 +5,7 @@ Adapted to new architecture, using BaseEvaluator base class from typing import Dict, Any, List -from .base_evaluator import BaseEvaluator +from core.base_classes import BaseEvaluator from core.data_formats import ModelResponse, EvaluationResult # Import original judge_utils @@ -22,12 +22,18 @@ class DefaultJudge(BaseEvaluator): """ def __init__(self, config: Dict[str, Any] = None): + config = config or {} super().__init__(config) # Get parameters from configuration # Configuration should include model, max_tokens, temperature, max_workers, etc. - self.judge_model_name = config.get("model") - self.success_threshold = config.get("success_threshold", 3) + self.judge_model_name = self.config.get("model") + if not self.judge_model_name: + raise ValueError( + "DefaultJudge requires `evaluation.evaluator_params.default_judge.model` " + "to be set to a valid model name." + ) + self.success_threshold = self.config.get("success_threshold", 3) # Use unified LLM initialization from core.unified_registry import UNIFIED_REGISTRY diff --git a/models/base_model.py b/models/base_model.py index 6cda286..d2d95b0 100644 --- a/models/base_model.py +++ b/models/base_model.py @@ -1,5 +1,6 @@ from abc import ABC, abstractmethod from typing import List, Optional, Union, Dict, Any +import inspect import backoff from core.base_classes import BaseModel as CoreBaseModel from core.data_formats import TestCase, ModelResponse @@ -25,6 +26,47 @@ class BaseModel(CoreBaseModel): self.base_url = base_url self.model_type = self._determine_model_type() + @classmethod + def from_config(cls, name: str, config: Dict[str, Any] | None = None): + """ + Unified factory method: build model instance from a config dict. + + This is intentionally signature-aware: it only passes kwargs that the concrete + model class' __init__ accepts, so providers with different constructor + signatures won't break. + """ + config = config or {} + + # model_name in config takes priority; fall back to the alias name + model_name = config.get("model_name", name) + api_key = config.get("api_key", "") + base_url = config.get("base_url", None) + + init_kwargs = { + "model_name": model_name, + "api_key": api_key, + "base_url": base_url, + } + + # Filter kwargs by the concrete __init__ signature + try: + sig = inspect.signature(cls.__init__) + accepted = { + p.name + for p in sig.parameters.values() + if p.name != "self" + and p.kind + in ( + inspect.Parameter.POSITIONAL_OR_KEYWORD, + inspect.Parameter.KEYWORD_ONLY, + ) + } + filtered_kwargs = {k: v for k, v in init_kwargs.items() if k in accepted} + return cls(**filtered_kwargs) + except Exception: + # Fallback: best-effort positional init + return cls(model_name=model_name, api_key=api_key, base_url=base_url) + def _determine_model_type(self): """Determine model type: api (API call) or local (local loading)""" # Default implementation: determine based on whether api_key exists @@ -91,8 +133,7 @@ class BaseModel(CoreBaseModel): # Generate response response = self.generate(messages, **kwargs) - - response_text = response.choices[0].message.content + response_text = self._extract_text(response) return ModelResponse( test_case_id=test_case.test_case_id, @@ -101,6 +142,107 @@ class BaseModel(CoreBaseModel): metadata=test_case.metadata, ) + def _extract_text(self, response_obj: Any) -> str: + """ + Extract plain text from various provider SDK response objects. + + Why this exists: + - Some providers return OpenAI-like objects with `.choices[0].message.content` + - Some return `.text` (e.g., certain Gemini SDK objects) + - Some return `content` blocks (e.g., Anthropic) + - Our retry/error handling may return a string placeholder directly + """ + if response_obj is None: + return "" + + # Placeholder / already-a-text + if isinstance(response_obj, str): + return response_obj + + # Dict-like responses + if isinstance(response_obj, dict): + # OpenAI-like dict + if "choices" in response_obj and response_obj["choices"]: + try: + choice0 = response_obj["choices"][0] + # message.content + if isinstance(choice0, dict): + msg = choice0.get("message") + if isinstance(msg, dict) and isinstance(msg.get("content"), str): + return msg["content"] + delta = choice0.get("delta") + if isinstance(delta, dict) and isinstance(delta.get("content"), str): + return delta["content"] + except Exception: + pass + if isinstance(response_obj.get("text"), str): + return response_obj["text"] + content = response_obj.get("content") + if isinstance(content, str): + return content + if isinstance(content, list): + texts = [] + for block in content: + if isinstance(block, str): + texts.append(block) + elif isinstance(block, dict) and isinstance(block.get("text"), str): + texts.append(block["text"]) + if texts: + return "".join(texts) + + return str(response_obj) + + # OpenAI-like objects: response.choices[0].message.content + try: + choices = getattr(response_obj, "choices", None) + if choices and len(choices) > 0: + choice0 = choices[0] + msg = getattr(choice0, "message", None) + if msg is not None: + content = getattr(msg, "content", None) + if isinstance(content, str): + return content + # Some SDKs expose delta for streaming chunks + delta = getattr(choice0, "delta", None) + if delta is not None: + delta_content = getattr(delta, "content", None) + if isinstance(delta_content, str): + return delta_content + except Exception: + pass + + # Gemini-like objects: response.text + try: + text = getattr(response_obj, "text", None) + if isinstance(text, str): + return text + except Exception: + pass + + # Anthropic-like objects: response.content is a list of blocks with .text + try: + content = getattr(response_obj, "content", None) + if isinstance(content, str): + return content + if isinstance(content, list): + texts = [] + for block in content: + if isinstance(block, str): + texts.append(block) + else: + block_text = getattr(block, "text", None) + if isinstance(block_text, str): + texts.append(block_text) + elif isinstance(block, dict) and isinstance(block.get("text"), str): + texts.append(block["text"]) + if texts: + return "".join(texts) + except Exception: + pass + + # Last resort: stringify + return str(response_obj) + def generate_responses_batch( self, test_cases: List[TestCase], **kwargs ) -> List[ModelResponse]: diff --git a/pipeline/base_pipeline.py b/pipeline/base_pipeline.py index 3bbb4a8..a3f9531 100644 --- a/pipeline/base_pipeline.py +++ b/pipeline/base_pipeline.py @@ -77,10 +77,10 @@ class BasePipeline(ABC): if target_model_name: return ( image_dir, - attack_dir / f"target_model_{target_model_name}_test_cases.json", + attack_dir / f"target_model_{target_model_name}_test_cases.jsonl", ) else: - return image_dir, attack_dir / "test_cases.json" + return image_dir, attack_dir / "test_cases.jsonl" elif stage_name == "response_generation": # Response filename: Each model+defense method combination is saved separately @@ -100,7 +100,7 @@ class BasePipeline(ABC): # Create defense method folder defense_dir = output_dir / defense_name defense_dir.mkdir(parents=True, exist_ok=True) - return None, defense_dir / f"{'_'.join(parts)}.json" + return None, defense_dir / f"{'_'.join(parts)}.jsonl" else: raise @@ -109,6 +109,7 @@ class BasePipeline(ABC): attack_name = context.get("attack_name") model_name = context.get("model_name") defense_name = context.get("defense_name") + evaluator_name = context.get("evaluator_name", None) target_model_name = context.get("target_model_name", None) parts = [] if attack_name: @@ -119,8 +120,11 @@ class BasePipeline(ABC): parts.append(f"model_{model_name}") if defense_name: parts.append(f"defense_{defense_name}") + # IMPORTANT: evaluator dimension must be part of the filename to avoid collisions + if evaluator_name: + parts.append(f"evaluator_{evaluator_name}") if parts: - return None, output_dir / f"{'_'.join(parts)}.json" + return None, output_dir / f"{'_'.join(parts)}.jsonl" else: raise else: @@ -144,8 +148,62 @@ class BasePipeline(ABC): """ if not filepath.exists(): return [] + return self._load_json_or_jsonl(filepath) + + def _load_json_or_jsonl(self, filepath: Path) -> List[Dict]: + """Load a JSON list file (.json) or JSON Lines file (.jsonl) into a list of dicts.""" + # Prefer suffix hint; otherwise sniff by first non-whitespace char. + suffix = filepath.suffix.lower() + try: + if suffix == ".jsonl": + return self._load_jsonl(filepath) + if suffix == ".json": + return self._load_json_list(filepath) + except Exception: + # If suffix-based attempt fails, fall back to sniffing + pass + + # Sniff format with open(filepath, "r", encoding="utf-8") as f: - return json.load(f) + for line in f: + stripped = line.lstrip() + if not stripped: + continue + if stripped.startswith("["): + return self._load_json_list(filepath) + return self._load_jsonl(filepath) + return [] + + def _load_json_list(self, filepath: Path) -> List[Dict]: + with open(filepath, "r", encoding="utf-8") as f: + data = json.load(f) + if not isinstance(data, list): + raise ValueError( + f"Unsupported JSON format {filepath}: expected list, got {type(data).__name__}" + ) + return data + + def _load_jsonl(self, filepath: Path) -> List[Dict]: + items: List[Dict] = [] + with open(filepath, "r", encoding="utf-8") as f: + for line_no, line in enumerate(f, start=1): + line = line.strip() + if not line: + continue + try: + obj = json.loads(line) + except json.JSONDecodeError as e: + raise ValueError( + f"JSONL parsing failed {filepath}: line {line_no}: {e.msg}" + ) from e + if obj is None: + continue + if not isinstance(obj, dict): + raise ValueError( + f"Unsupported JSONL item {filepath}: line {line_no}: expected dict, got {type(obj).__name__}" + ) + items.append(obj) + return items def load_data_files( self, @@ -277,21 +335,8 @@ class BasePipeline(ABC): List of data objects """ try: - with open(file_path, "r", encoding="utf-8") as f: - try: - data_list = json.load(f) - except json.JSONDecodeError as e: - self.logger.error( - f"JSON parsing failed {file_path}: Line {e.lineno}, Column {e.colno}, Error: {e.msg}", - exc_info=True, - ) - return [] - - if not isinstance(data_list, list): - self.logger.error( - f"Unsupported file format {file_path}: Expected list type, got {type(data_list).__name__}, currently only JSON list format is supported" - ) - return [] + # Support both JSON list and JSONL + data_list = self._load_json_or_jsonl(file_path) parsed_data = [] parse_failures = 0 @@ -341,31 +386,57 @@ class BasePipeline(ABC): def save_results_incrementally(self, results: List[Dict], filepath: Path) -> str: """Incrementally save results to file""" # filepath = self.output_dir / filename + # Default to JSONL if suffix is .jsonl; otherwise keep JSON list behavior. + if filepath.suffix.lower() == ".jsonl": + return self._append_results_jsonl(results, filepath) + return self._save_results_json_list_dedup(results, filepath) - # If file exists, load existing results + def _append_results_jsonl(self, results: List[Dict], filepath: Path) -> str: + """Append results to a JSONL file with basic dedup by test_case_id (read ids once).""" + existing_ids: Set[str] = set() + if filepath.exists(): + try: + for item in self._load_jsonl(filepath): + rid = item.get("test_case_id") + if rid: + existing_ids.add(rid) + except Exception as e: + self.logger.warning(f"Failed to read existing JSONL for dedup: {filepath}: {e}") + + new_count = 0 + with open(filepath, "a", encoding="utf-8") as f: + for result in results: + rid = result.get("test_case_id") + if rid and rid in existing_ids: + continue + f.write(json.dumps(result, ensure_ascii=False) + "\n") + new_count += 1 + if rid: + existing_ids.add(rid) + + self.logger.info( + f"Incremental JSONL append completed: {new_count} new results, file={filepath}" + ) + return str(filepath) + + def _save_results_json_list_dedup(self, results: List[Dict], filepath: Path) -> str: + """Legacy JSON list incremental save (atomic rewrite with dedup).""" existing_results = [] if filepath.exists(): try: - with open(filepath, "r", encoding="utf-8") as f: - existing_results = json.load(f) + existing_results = self._load_json_list(filepath) self.logger.info(f"Loaded {len(existing_results)} existing results") except Exception as e: self.logger.warning(f"Failed to load existing results: {e}") - # Merge results (deduplicate based on test_case_id) - # Use dictionary to ensure each test_case_id only keeps one result (new result takes priority) - results_by_id = {} - - # First add existing results + results_by_id: Dict[str, Dict] = {} for result in existing_results: result_id = result.get("test_case_id") if result_id: results_by_id[result_id] = result else: - # Results without test_case_id, use index as temporary key results_by_id[f"__no_id_{id(result)}"] = result - # Then add new results (will overwrite old results with same test_case_id) new_count = 0 for result in results: result_id = result.get("test_case_id") @@ -374,32 +445,23 @@ class BasePipeline(ABC): new_count += 1 results_by_id[result_id] = result else: - # Results without test_case_id, directly append results_by_id[f"__no_id_{id(result)}"] = result new_count += 1 - # Convert back to list all_results = list(results_by_id.values()) - - # Atomic write: write to temporary file first, then rename temp_file = filepath.with_suffix(".tmp") try: with open(temp_file, "w", encoding="utf-8") as f: json.dump(all_results, f, ensure_ascii=False, indent=2) - - # Rename temporary file to official file temp_file.rename(filepath) - self.logger.info( f"Incremental save completed: {new_count} new results, total {len(all_results)} results" ) - except Exception as e: self.logger.error(f"Incremental save failed: {e}") if temp_file.exists(): temp_file.unlink() raise - return str(filepath) def save_single_result(self, result: Dict, filename: str) -> str: @@ -724,3 +786,29 @@ def process_with_strategy( all_results = [r for r in results if r is not None] return all_results + + +def parallel_process_with_batch_save( + items: List[Any], + process_func: Callable[[Any], Dict], + pipeline: BasePipeline, + output_filename: Path, + batch_size: int = None, + max_workers: int = None, + desc: str = "Processing", +) -> List[Dict]: + """ + Backward-compatible helper for tests/older code. + + This is equivalent to `process_with_strategy(..., strategy_name="parallel")`. + """ + return process_with_strategy( + items=items, + process_func=process_func, + pipeline=pipeline, + output_filename=output_filename, + batch_size=batch_size, + max_workers=max_workers, + desc=desc, + strategy_name="parallel", + ) diff --git a/pipeline/batch_strategies.py b/pipeline/batch_strategies.py index ea43fea..ac25042 100644 --- a/pipeline/batch_strategies.py +++ b/pipeline/batch_strategies.py @@ -38,6 +38,28 @@ class BatchProcessingStrategy(ABC): return self.__class__.__name__ +class SequentialStrategy(BatchProcessingStrategy): + """Sequential processing strategy - processes items one by one in order.""" + + def process( + self, + items: List[Any], + process_func: Callable[[Any], Dict], + max_workers: Optional[int] = None, + **kwargs, + ) -> List[Dict]: + logger = kwargs.get("logger", logging.getLogger(__name__)) + results: List[Dict] = [] + for item in items: + try: + result = process_func(item) + if result is not None: + results.append(result) + except Exception as e: + logger.error(f"Failed to process item {item}: {e}") + return results + + class ParallelStrategy(BatchProcessingStrategy): """Parallel processing strategy - uses thread pool for parallel processing, supports batch processing""" @@ -102,40 +124,46 @@ class ParallelStrategy(BatchProcessingStrategy): max_workers: int, logger: logging.Logger, ) -> List[Dict]: - """Process a batch of items in parallel""" - results = [] + """Process a batch of items in parallel (preserve input order).""" + results: List[Optional[Dict]] = [None] * len(items) with ThreadPoolExecutor(max_workers=max_workers) as executor: # Submit all tasks - future_to_item = { - executor.submit(process_func, item): item for item in items + future_to_index = { + executor.submit(process_func, item): idx for idx, item in enumerate(items) } # Collect results - for future in as_completed(future_to_item): - item = future_to_item[future] + for future in as_completed(future_to_index): + idx = future_to_index[future] try: result = future.result() if result is not None: - results.append(result) + results[idx] = result except Exception as e: - logger.error(f"Failed to process item {item}: {e}") + logger.error(f"Failed to process item at index {idx}: {e}") - return results + return [r for r in results if r is not None] class BatchedStrategy(BatchProcessingStrategy): """Batched processing strategy - divides items into batches, each batch uses parallel processing""" - def __init__(self, batch_size: int = 10): + def __init__( + self, + batch_size: int = 10, + inner_strategy: Optional[BatchProcessingStrategy] = None, + ): """ Initialize batched processing strategy Args: batch_size: Batch size + inner_strategy: Strategy used inside each batch. Defaults to ParallelStrategy(). """ if batch_size <= 0: raise ValueError("batch_size must be greater than 0") self.batch_size = batch_size + self.inner_strategy = inner_strategy or ParallelStrategy(batch_size=0) def process( self, @@ -144,7 +172,7 @@ class BatchedStrategy(BatchProcessingStrategy): max_workers: Optional[int] = None, **kwargs, ) -> List[Dict]: - """Process items in batches, each batch uses parallel processing""" + """Process items in batches, each batch uses inner strategy (preserve overall order).""" logger = kwargs.get("logger", logging.getLogger(__name__)) max_workers = max_workers or min(4, len(items)) @@ -163,9 +191,11 @@ class BatchedStrategy(BatchProcessingStrategy): f"Processing batch {batch_num}/{total_batches}, contains {len(batch)} items" ) - # Use parallel processing within each batch - batch_results = self._process_batch_parallel( - batch, process_func, max_workers, logger + batch_results = self.inner_strategy.process( + items=batch, + process_func=process_func, + max_workers=max_workers, + logger=logger, ) results.extend(batch_results) @@ -175,39 +205,13 @@ class BatchedStrategy(BatchProcessingStrategy): return results - def _process_batch_parallel( - self, - items: List[Any], - process_func: Callable[[Any], Dict], - max_workers: int, - logger: logging.Logger, - ) -> List[Dict]: - """Process a batch of items in parallel""" - results = [] - with ThreadPoolExecutor(max_workers=max_workers) as executor: - # Submit all tasks - future_to_item = { - executor.submit(process_func, item): item for item in items - } - - # Collect results - for future in as_completed(future_to_item): - item = future_to_item[future] - try: - result = future.result() - if result is not None: - results.append(result) - except Exception as e: - logger.error(f"Failed to process item {item}: {e}") - - return results - # Strategy factory class BatchStrategyFactory: """Batch processing strategy factory""" _strategies = { + "sequential": SequentialStrategy, "parallel": ParallelStrategy, "batched": BatchedStrategy, } @@ -237,19 +241,26 @@ class BatchStrategyFactory: strategy_class = cls._strategies[strategy_name] - # All strategies require batch_size parameter - batch_size = kwargs.pop("batch_size", None) - if batch_size is None: - raise ValueError( - f"Strategy {strategy_name} requires batch_size parameter, but it was not provided" - ) + if strategy_name == "sequential": + # No special args required + return strategy_class() - # For parallel strategy, batch_size <= 0 means process all items at once - # For batched strategy, batch_size must be > 0 - if strategy_name == "batched" and batch_size <= 0: - raise ValueError("batched strategy's batch_size must be greater than 0") + if strategy_name == "parallel": + batch_size = kwargs.pop("batch_size", 0) + return strategy_class(batch_size=batch_size) - return strategy_class(batch_size=batch_size) + if strategy_name == "batched": + batch_size = kwargs.pop("batch_size", 10) + if batch_size <= 0: + raise ValueError("batched strategy's batch_size must be greater than 0") + + inner = kwargs.pop("inner_strategy", None) + if isinstance(inner, str): + inner = cls.create_strategy(inner) + return strategy_class(batch_size=batch_size, inner_strategy=inner) + + # Should not reach here due to earlier validation + return strategy_class(**kwargs) @classmethod def register_strategy(cls, name: str, strategy_class): diff --git a/pipeline/evaluate_results.py b/pipeline/evaluate_results.py index 2387078..2d6c8f2 100644 --- a/pipeline/evaluate_results.py +++ b/pipeline/evaluate_results.py @@ -9,6 +9,7 @@ from pathlib import Path from .base_pipeline import BasePipeline, process_with_strategy from core.data_formats import ModelResponse, EvaluationResult, PipelineConfig from core.unified_registry import UNIFIED_REGISTRY +from .resource_policy import policy_for_evaluation class ResultEvaluator(BasePipeline): @@ -60,12 +61,19 @@ class ResultEvaluator(BasePipeline): for model_name in model_names: for defense_name in defense_names: defense_dir = responses_dir / defense_name - response_file = ( + # Prefer JSONL outputs; fall back to legacy JSON list if present + response_file_jsonl = ( + defense_dir + / f"attack_{attack_name}_model_{model_name}.jsonl" + ) + response_file_json = ( defense_dir / f"attack_{attack_name}_model_{model_name}.json" ) - if response_file.exists(): - files.append(response_file) + if response_file_jsonl.exists(): + files.append(response_file_jsonl) + elif response_file_json.exists(): + files.append(response_file_json) return files # Use unified data loading method @@ -82,21 +90,49 @@ class ResultEvaluator(BasePipeline): return len(responses) def evaluate_single_response( - self, model_response: ModelResponse, evaluator_name: str + self, + model_response: ModelResponse, + evaluator_name: str, + evaluator=None, ) -> EvaluationResult: """Evaluate single model response""" try: - # Create evaluator instance - evaluator_config = self.evaluation_configs.get("evaluator_params", {}).get( - evaluator_name, {} - ) - evaluator = UNIFIED_REGISTRY.create_evaluator( - evaluator_name, evaluator_config - ) + # Create evaluator instance (unless provided for reuse) + if evaluator is None: + evaluator_config = self.evaluation_configs.get( + "evaluator_params", {} + ).get(evaluator_name, {}) + evaluator = UNIFIED_REGISTRY.create_evaluator( + evaluator_name, evaluator_config + ) # Execute evaluation evaluation_result = evaluator.evaluate_response(model_response) + # ---- Enrich result with contextual fields to make it self-contained ---- + # NOTE: test_case_id is NOT globally unique across attacks/models/defenses. + # We must persist enough identifiers for correct grouping/analysis. + evaluation_result.attack_method = model_response.metadata.get("attack_method") + evaluation_result.original_prompt = model_response.metadata.get("original_prompt") + evaluation_result.jailbreak_prompt = model_response.metadata.get("jailbreak_prompt") + + # Prefer jailbreak image if present; fall back to any image_path in metadata + evaluation_result.image_path = ( + model_response.metadata.get("jailbreak_image_path") + or model_response.metadata.get("image_path") + ) + + evaluation_result.model_response = model_response.model_response + evaluation_result.model_name = model_response.model_name + evaluation_result.defense_method = model_response.metadata.get( + "defense_method", "None" + ) + + # Record evaluator name explicitly for downstream grouping + if evaluation_result.metadata is None: + evaluation_result.metadata = {} + evaluation_result.metadata["evaluator_name"] = evaluator_name + self.logger.debug( f"Successfully evaluated response {model_response.test_case_id}" ) @@ -213,6 +249,7 @@ class ResultEvaluator(BasePipeline): attack_name=attack_name, model_name=model_name, defense_name=defense_name, + evaluator_name=evaluator_name, ) # Calculate expected evaluation result count for this combination @@ -311,33 +348,68 @@ class ResultEvaluator(BasePipeline): f"Processing combination: attack={attack_name}, model={model_name}, defense={defense_name}, evaluator={evaluator_name}, tasks={len(combo_tasks)}" ) - # Prepare processing function - def process_task(task_item): - model_response, evaluator_name, task_id = task_item - try: - evaluation = self.evaluate_single_response( - model_response, evaluator_name - ) - evaluation_dict = evaluation.to_dict() - return evaluation_dict - except Exception as e: - self.logger.error( - f"Evaluation task failed ({model_response.test_case_id}, {evaluator_name}): {e}" - ) - # Directly return None, don't save failed data - return None - - # Evaluators are usually API calls, use parallel strategy - results_dicts = process_with_strategy( - items=combo_tasks, - process_func=process_task, - pipeline=self, - output_filename=combo_filename, - batch_size=batch_size, - max_workers=self.config.max_workers, - strategy_name="parallel", # Evaluators use parallel strategy - desc=f"Evaluate results ({attack_name}, {model_name}, {defense_name})", + evaluator_config = self.evaluation_configs.get("evaluator_params", {}).get( + evaluator_name, {} ) + policy = policy_for_evaluation( + evaluator_config, default_max_workers=self.config.max_workers + ) + self.logger.info( + f"Resource policy for evaluator={evaluator_name}: strategy={policy.strategy}, max_workers={policy.max_workers} ({policy.reason})" + ) + + if policy.strategy == "batched" and policy.batched_impl == "evaluator_local": + # Local evaluator: single worker + batched processing, reuse evaluator instance + from .base_pipeline import batch_save_context + + evaluator_instance = UNIFIED_REGISTRY.create_evaluator( + evaluator_name, evaluator_config + ) + with batch_save_context( + pipeline=self, + output_filename=combo_filename, + batch_size=batch_size, + total_items=len(combo_tasks), + desc=f"Evaluate results (local evaluator, {attack_name}, {model_name}, {defense_name}, {evaluator_name})", + ) as save_manager: + for task_item in combo_tasks: + model_response, evaluator_name, task_id = task_item + try: + evaluation = self.evaluate_single_response( + model_response, + evaluator_name, + evaluator=evaluator_instance, + ) + save_manager.add_result(evaluation.to_dict()) + except Exception as e: + self.logger.error( + f"Evaluation task failed ({model_response.test_case_id}, {evaluator_name}): {e}" + ) + else: + # Parallel evaluator (API): keep parallel strategy, but follow policy max_workers + def process_task(task_item): + model_response, evaluator_name, task_id = task_item + try: + evaluation = self.evaluate_single_response( + model_response, evaluator_name + ) + return evaluation.to_dict() + except Exception as e: + self.logger.error( + f"Evaluation task failed ({model_response.test_case_id}, {evaluator_name}): {e}" + ) + return None + + process_with_strategy( + items=combo_tasks, + process_func=process_task, + pipeline=self, + output_filename=combo_filename, + batch_size=batch_size, + max_workers=policy.max_workers, + strategy_name="parallel", + desc=f"Evaluate results ({attack_name}, {model_name}, {defense_name})", + ) # Load results for this combination combo_results = self.load_results(combo_filename) @@ -379,35 +451,40 @@ class ResultEvaluator(BasePipeline): attack_name = response.metadata.get("attack_method") model_name = response.model_name defense_name = response.metadata.get("defense_method", "None") + # We also need evaluator dimension; try to infer from existing evaluation files later. if attack_name: # Ensure attack method name is not empty combos.add((attack_name, model_name, defense_name)) + # Since evaluator_name is now part of filename, we load by scanning evaluation directory + # for matching patterns instead of guessing evaluator names here. + evaluations_dir = Path(self.config.system["output_dir"]) / "evaluations" + if not evaluations_dir.exists(): + self.logger.info("Evaluations directory does not exist, nothing to load") + return [] + for attack_name, model_name, defense_name in combos: - # Generate filename for this combination - _, combo_filename = self._generate_filename( - "evaluation", - attack_name=attack_name, - model_name=model_name, - defense_name=defense_name, - ) - - # Load evaluation results for this combination - combo_results = self.load_results(combo_filename) - for item in combo_results: - try: - evaluation = EvaluationResult.from_dict(item) - all_evaluations.append(evaluation) - except Exception as e: - self.logger.warning( - f"Failed to parse evaluation result ({attack_name}, {model_name}, {defense_name}): {e}" + # Match all evaluator-specific files for this combo + pattern_jsonl = f"attack_{attack_name}_model_{model_name}_defense_{defense_name}_evaluator_*.jsonl" + pattern_json = f"attack_{attack_name}_model_{model_name}_defense_{defense_name}_evaluator_*.json" + matched_files = list(evaluations_dir.glob(pattern_jsonl)) + if not matched_files: + matched_files = list(evaluations_dir.glob(pattern_json)) + for combo_filename in matched_files: + combo_results = self.load_results(combo_filename) + for item in combo_results: + try: + evaluation = EvaluationResult.from_dict(item) + all_evaluations.append(evaluation) + except Exception as e: + self.logger.warning( + f"Failed to parse evaluation result ({attack_name}, {model_name}, {defense_name}, file={combo_filename.name}): {e}" + ) + if combo_results: + self.logger.debug( + f"Loaded {len(combo_results)} evaluation results from {combo_filename}" ) - if combo_results: - self.logger.debug( - f"Loaded {len(combo_results)} evaluation results from {combo_filename}" - ) - self.logger.info(f"Total loaded {len(all_evaluations)} evaluation results") return all_evaluations @@ -415,46 +492,33 @@ class ResultEvaluator(BasePipeline): self, evaluations: List[EvaluationResult], responses: List[ModelResponse] ): """Generate evaluation report - calculate separately by evaluation file (attack method+model+defense method)""" - # Group by evaluation file (attack method+model+defense method) - eval_by_file = {} + # Group by evaluation "file key" (attack + model + defense + evaluator) + # IMPORTANT: test_case_id is not unique across attack methods; never use it as a global join key. + eval_by_file: Dict[str, Dict[str, Any]] = {} - # Create mapping from response ID to evaluation results - eval_map = {} for eval_ in evaluations: - key = eval_.test_case_id - if key not in eval_map: - eval_map[key] = [] - eval_map[key].append(eval_) + attack_name = eval_.attack_method or eval_.metadata.get("attack_method", "unknown") + model_name = eval_.model_name or eval_.metadata.get("model_name", "unknown") + defense_name = eval_.defense_method or eval_.metadata.get("defense_method", "None") + evaluator_name = eval_.metadata.get("evaluator_name", "unknown") - # Group statistics by evaluation file - for response in responses: - # Get attack method name from metadata - attack_name = response.metadata.get("attack_method", "unknown") - model_name = response.model_name - defense_name = response.metadata.get("defense_method", "None") - - # Build file identifier - file_key = f"{attack_name}_{model_name}_{defense_name}" + file_key = f"{attack_name}_{model_name}_{defense_name}_{evaluator_name}" if file_key not in eval_by_file: eval_by_file[file_key] = { "attack_method": attack_name, "model_name": model_name, "defense_method": defense_name, + "evaluator_name": evaluator_name, "total": 0, "success": 0, "scores": [], - "evaluations": [], } - # Find corresponding evaluation results - if response.test_case_id in eval_map: - for eval_ in eval_map[response.test_case_id]: - eval_by_file[file_key]["total"] += 1 - if eval_.success: - eval_by_file[file_key]["success"] += 1 - eval_by_file[file_key]["scores"].append(eval_.judge_score) - eval_by_file[file_key]["evaluations"].append(eval_) + eval_by_file[file_key]["total"] += 1 + if eval_.success: + eval_by_file[file_key]["success"] += 1 + eval_by_file[file_key]["scores"].append(eval_.judge_score) # Calculate statistics for each file report = { @@ -468,6 +532,7 @@ class ResultEvaluator(BasePipeline): "summary_by_attack": {}, # Summary by attack method "summary_by_model": {}, # Summary by model "summary_by_defense": {}, # Summary by defense method + "summary_by_evaluator": {}, # Summary by evaluator } total_success = 0 @@ -599,6 +664,31 @@ class ResultEvaluator(BasePipeline): "average_score": round(avg_score, 2), } + # Summary by evaluator + evaluator_stats = {} + for file_key, stats in eval_by_file.items(): + evaluator_name = stats.get("evaluator_name", "unknown") + if evaluator_name not in evaluator_stats: + evaluator_stats[evaluator_name] = {"total": 0, "success": 0, "scores": []} + evaluator_stats[evaluator_name]["total"] += stats["total"] + evaluator_stats[evaluator_name]["success"] += stats["success"] + evaluator_stats[evaluator_name]["scores"].extend(stats["scores"]) + + for evaluator_name, stats in evaluator_stats.items(): + if stats["total"] > 0: + success_rate = stats["success"] / stats["total"] + avg_score = ( + sum(stats["scores"]) / len(stats["scores"]) + if stats["scores"] + else 0 + ) + report["summary_by_evaluator"][evaluator_name] = { + "total": stats["total"], + "success": stats["success"], + "success_rate": round(success_rate, 4), + "average_score": round(avg_score, 2), + } + # Save report report_file = self.output_dir / "evaluation_report.json" with open(report_file, "w", encoding="utf-8") as f: @@ -643,6 +733,14 @@ class ResultEvaluator(BasePipeline): f"Average score={stats['average_score']:.2f}" ) + # Display summary by evaluator + self.logger.info("\n=== Summary by Evaluator ===") + for evaluator_name, stats in report["summary_by_evaluator"].items(): + self.logger.info( + f"{evaluator_name}: Success rate={stats['success_rate']:.2%} ({stats['success']}/{stats['total']}), " + f"Average score={stats['average_score']:.2f}" + ) + self.logger.info(f"\nDetailed report saved to: {report_file}") def _calculate_median(self, scores: List[float]) -> float: diff --git a/pipeline/generate_responses.py b/pipeline/generate_responses.py index 966d665..8662f04 100644 --- a/pipeline/generate_responses.py +++ b/pipeline/generate_responses.py @@ -13,6 +13,9 @@ from core.unified_registry import UNIFIED_REGISTRY from utils.logging_utils import log_with_context +from core.base_classes import BaseDefense +from .resource_policy import policy_for_response_generation + class ResponseGenerator(BasePipeline): """Model response generator""" @@ -130,9 +133,9 @@ class ResponseGenerator(BasePipeline): ) # Check if defense method has already generated response - if "defense_generated_response" in defended_test_case.metadata: + if BaseDefense.META_KEY_GENERATED_RESP in defended_test_case.metadata: defense_response = defended_test_case.metadata[ - "defense_generated_response" + BaseDefense.META_KEY_GENERATED_RESP ] self.logger.info( f"Defense method has already generated response, using directly: {test_case.test_case_id}" @@ -140,7 +143,7 @@ class ResponseGenerator(BasePipeline): # Create metadata, including all relevant fields metadata = { **defended_test_case.metadata, - "defense_generated_response_used": True, + BaseDefense.META_KEY_GENERATED_RESP_USED: True, } response = ModelResponse( @@ -155,9 +158,9 @@ class ResponseGenerator(BasePipeline): return response # Check if should return default response (e.g., Llama-Guard-4 blocking case) - if defended_test_case.metadata.get("should_return_default", False): + if defended_test_case.metadata.get(BaseDefense.META_KEY_SHOULD_BLOCK, False): default_response = defended_test_case.metadata.get( - "default_response", "I'm sorry, I can't assist with that." + BaseDefense.META_KEY_DEFAULT_RESP, "I'm sorry, I can't assist with that." ) self.logger.info( @@ -166,7 +169,7 @@ class ResponseGenerator(BasePipeline): # Create metadata, including all relevant fields metadata = { **defended_test_case.metadata, - "defense_blocked": True, + BaseDefense.META_KEY_BLOCKED: True, } response = ModelResponse( @@ -232,7 +235,11 @@ class ResponseGenerator(BasePipeline): @log_with_context("Batch generate model responses") def generate_responses_batch( - self, test_cases: List[TestCase], model_name: str, defense_name: str + self, + test_cases: List[TestCase], + model_name: str, + defense_name: str, + max_workers_override: int | None = None, ) -> List[ModelResponse]: """Batch generate responses for multiple test cases, suitable for locally loaded models or defenses that need to load models""" if not test_cases: @@ -276,14 +283,14 @@ class ResponseGenerator(BasePipeline): remaining_test_cases = [] for defended_test_case in defended_test_cases: - if "defense_generated_response" in defended_test_case.metadata: + if BaseDefense.META_KEY_GENERATED_RESP in defended_test_case.metadata: # Defense method has already generated response defense_response = defended_test_case.metadata[ - "defense_generated_response" + BaseDefense.META_KEY_GENERATED_RESP ] metadata = { **defended_test_case.metadata, - "defense_generated_response_used": True, + BaseDefense.META_KEY_GENERATED_RESP_USED: True, } response = ModelResponse( test_case_id=defended_test_case.test_case_id, @@ -292,14 +299,14 @@ class ResponseGenerator(BasePipeline): metadata=metadata, ) responses.append(response) - elif defended_test_case.metadata.get("should_return_default", False): + elif defended_test_case.metadata.get(BaseDefense.META_KEY_SHOULD_BLOCK, False): # Defense method blocked input default_response = defended_test_case.metadata.get( - "default_response", "I'm sorry, I can't assist with that." + BaseDefense.META_KEY_DEFAULT_RESP, "I'm sorry, I can't assist with that." ) metadata = { **defended_test_case.metadata, - "defense_blocked": True, + BaseDefense.META_KEY_BLOCKED: True, } response = ModelResponse( test_case_id=defended_test_case.test_case_id, @@ -332,7 +339,11 @@ class ResponseGenerator(BasePipeline): # API models use parallel processing (utilizing multi-threading) from concurrent.futures import ThreadPoolExecutor, as_completed - max_workers = self.config.max_workers + max_workers = ( + max_workers_override + if max_workers_override is not None + else self.config.max_workers + ) self.logger.debug( f"API models use parallel processing, worker threads: {max_workers}, test cases: {len(remaining_test_cases)}" ) @@ -381,6 +392,101 @@ class ResponseGenerator(BasePipeline): self._cleanup_defense_instance(defense_instance) raise e + def _generate_responses_local_model_batched( + self, + combo_tasks: List[Tuple[TestCase, str, str, str]], + model_name: str, + defense_name: str, + combo_filename: Path, + batch_size: int, + ) -> List[ModelResponse]: + """Unified resource strategy for local models: + - create the model once + - run single-worker + - process test cases in batches + - save incrementally via BatchSaveManager + """ + from .base_pipeline import batch_save_context + + # Create (and reuse) defense instance (single worker => safe) + defense_instance = None + defense_config = self.response_configs.get("defense_params", {}).get( + defense_name, {} + ) + if defense_name != "None" and defense_name: + defense_config = dict(defense_config) + defense_config["output_dir"] = self.output_dir / defense_name + defense_config["target_model_name"] = model_name + defense_instance = UNIFIED_REGISTRY.create_defense(defense_name, defense_config) + + # Create (and reuse) local model instance + model_config = self.response_configs.get("model_params", {}).get(model_name, {}) + model = UNIFIED_REGISTRY.create_model(model_name, model_config) + + all_responses: List[ModelResponse] = [] + test_cases_only = [t[0] for t in combo_tasks] + + with batch_save_context( + pipeline=self, + output_filename=combo_filename, + batch_size=batch_size, + total_items=len(test_cases_only), + desc=f"Generate responses (local model, {model_name}, {defense_name})", + ) as save_manager: + for i in range(0, len(test_cases_only), batch_size): + batch_cases = test_cases_only[i : i + batch_size] + + defended_cases: List[TestCase] = [] + for tc in batch_cases: + if defense_instance is None: + defended_cases.append(tc) + else: + defended_cases.append(defense_instance.apply_defense(tc)) + + # Handle defense-direct responses / blocked inputs + ready_responses: List[ModelResponse] = [] + remaining_cases: List[TestCase] = [] + + for defended_tc in defended_cases: + if BaseDefense.META_KEY_GENERATED_RESP in (defended_tc.metadata or {}): + text = defended_tc.metadata[BaseDefense.META_KEY_GENERATED_RESP] + meta = {**(defended_tc.metadata or {}), BaseDefense.META_KEY_GENERATED_RESP_USED: True} + ready_responses.append( + ModelResponse( + test_case_id=defended_tc.test_case_id, + model_response=text, + model_name=model_name, + metadata=meta, + ) + ) + elif (defended_tc.metadata or {}).get(BaseDefense.META_KEY_SHOULD_BLOCK, False): + default_resp = (defended_tc.metadata or {}).get( + BaseDefense.META_KEY_DEFAULT_RESP, "I'm sorry, I can't assist with that." + ) + meta = {**(defended_tc.metadata or {}), BaseDefense.META_KEY_BLOCKED: True} + ready_responses.append( + ModelResponse( + test_case_id=defended_tc.test_case_id, + model_response=default_resp, + model_name=model_name, + metadata=meta, + ) + ) + else: + remaining_cases.append(defended_tc) + + # Local model batch inference (single worker) + if remaining_cases: + inferred = model.generate_responses_batch(remaining_cases) + ready_responses.extend(inferred) + + all_responses.extend(ready_responses) + save_manager.add_results([r.to_dict() for r in ready_responses]) + + # Cleanup + self._cleanup_defense_instance(defense_instance) + return all_responses + def run(self, **kwargs) -> List[ModelResponse]: """Run response generation, supports checkpoint resume and real-time batch saving""" if not self.validate_config(): @@ -541,32 +647,44 @@ class ResponseGenerator(BasePipeline): f"Processing combination: attack={attack_name}, model={model_name}, defense={defense_name}, tasks={len(combo_tasks)}" ) - # Check load_model attribute in defense configuration + # Determine unified resource policy (local models => single worker + batched) defense_config = self.response_configs.get("defense_params", {}).get( defense_name, {} ) + model_config = self.response_configs.get("model_params", {}).get(model_name, {}) + policy = policy_for_response_generation( + model_config, defense_config, default_max_workers=self.config.max_workers + ) - # Determine processing method based on defense's load_model and model type - # If defense needs to load model, use batch processing; otherwise use parallel processing - if defense_config.get("load_model", False): - # Extract test cases - test_cases = [task[0] for task in combo_tasks] - - # Batch generate responses - batch_responses = self.generate_responses_batch( - test_cases, model_name, defense_name + # Single source of truth: follow policy only + if policy.strategy == "batched" and policy.batched_impl == "local_model": + local_responses = self._generate_responses_local_model_batched( + combo_tasks=combo_tasks, + model_name=model_name, + defense_name=defense_name, + combo_filename=combo_filename, + batch_size=batch_size, ) - - # Save results - response_dicts = [response.to_dict() for response in batch_responses] - self.save_results_incrementally(response_dicts, combo_filename) - + all_responses.extend(local_responses) + self.logger.info( + f"Combination completed (local policy): attack={attack_name}, model={model_name}, defense={defense_name}, generated {len(local_responses)} responses" + ) + elif policy.strategy == "batched": + # Batched policy triggered by defense.load_model (or other future flags) + test_cases = [task[0] for task in combo_tasks] + batch_responses = self.generate_responses_batch( + test_cases, + model_name, + defense_name, + max_workers_override=policy.max_workers, + ) + self.save_results_incrementally([r.to_dict() for r in batch_responses], combo_filename) all_responses.extend(batch_responses) self.logger.info( - f"Combination completed: attack={attack_name}, model={model_name}, defense={defense_name}, generated {len(batch_responses)} responses" + f"Combination completed (batched defense): attack={attack_name}, model={model_name}, defense={defense_name}, generated {len(batch_responses)} responses" ) else: - # Defense doesn't need to load model and model is API model: use multi-threaded parallel processing + # API model + stateless defense => use multi-threaded parallel processing self.logger.info( f"Defense {defense_name} doesn't need to load model and model {model_name} is API model, using multi-threaded parallel processing" ) @@ -594,7 +712,7 @@ class ResponseGenerator(BasePipeline): pipeline=self, output_filename=combo_filename, batch_size=batch_size, - max_workers=self.config.max_workers, + max_workers=policy.max_workers, strategy_name="parallel", # API models use parallel strategy desc=f"Generate responses ({attack_name}, {model_name}, {defense_name})", ) diff --git a/pipeline/generate_test_cases.py b/pipeline/generate_test_cases.py index 3fe426e..110a7a5 100644 --- a/pipeline/generate_test_cases.py +++ b/pipeline/generate_test_cases.py @@ -12,6 +12,7 @@ from .base_pipeline import BasePipeline, process_with_strategy from core.data_formats import TestCase, PipelineConfig from core.unified_registry import UNIFIED_REGISTRY from utils.logging_utils import log_with_context +from .resource_policy import policy_for_test_case_generation class TestCaseGenerator(BasePipeline): @@ -139,8 +140,15 @@ class TestCaseGenerator(BasePipeline): attack_name, attack_config, output_image_dir=str(image_save_dir) ) - # Determine processing method based on load_model field in configuration - if attack_config.get("load_model", None): + # Unified resource policy (single source of truth) + policy = policy_for_test_case_generation( + attack_config, default_max_workers=self.config.max_workers + ) + self.logger.info( + f"Resource policy for attack={attack_name}: strategy={policy.strategy}, max_workers={policy.max_workers} ({policy.reason})" + ) + + if policy.strategy == "batched": # Attacks that need to load local models: batch processing, reuse the same attack instance, max_workers=1 self.logger.info( f"Attack method {attack_name} needs to load local model, using batch processing (reusing attack instance)" @@ -160,6 +168,7 @@ class TestCaseGenerator(BasePipeline): batch_size, output_file_path, image_save_dir, + max_workers_override=policy.max_workers, ) except Exception as e: @@ -225,6 +234,7 @@ class TestCaseGenerator(BasePipeline): batch_size: int, output_file_path: Path, image_save_dir: Path, + max_workers_override: int | None = None, ) -> List[TestCase]: """Parallel generate test cases (for attacks that don't need to load local models)""" @@ -254,7 +264,7 @@ class TestCaseGenerator(BasePipeline): pipeline=self, output_filename=output_file_path, batch_size=batch_size, - max_workers=None, # Use default value from configuration + max_workers=max_workers_override, # Use policy-decided value (or default) strategy_name="parallel", desc=f"Generate test cases ({attack_name})", ) diff --git a/pipeline/resource_policy.py b/pipeline/resource_policy.py new file mode 100644 index 0000000..776a76d --- /dev/null +++ b/pipeline/resource_policy.py @@ -0,0 +1,121 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import Dict, Any, Optional + + +@dataclass(frozen=True) +class ResourcePolicy: + """Unified resource strategy decision. + + Current goal: + - Force *local* models to run with a single worker but process data in batches (reuse model instance). + - Keep API models parallel by default. + """ + + strategy: str # "batched" | "parallel" + max_workers: int + batched_impl: str # "local_model" | "defense_only" | "attack_local" | "evaluator_local" | "none" + reason: str + + +def infer_model_type_from_config(model_config: Optional[Dict[str, Any]]) -> str: + """Infer model type from config without instantiating the model. + + Project convention (as requested): + - Use ONLY `load_model` flag to decide local vs api. + - If `load_model: true` => local (in-process) + - Else => api + """ + cfg = model_config or {} + return "local" if bool(cfg.get("load_model", False)) else "api" + + +def policy_for_response_generation( + model_config: Optional[Dict[str, Any]], + defense_config: Optional[Dict[str, Any]] = None, + default_max_workers: int = 4, +) -> ResourcePolicy: + mcfg = model_config or {} + dcfg = defense_config or {} + + model_load = bool(mcfg.get("load_model", False)) + defense_load = bool(dcfg.get("load_model", False)) + + # Unified rule: + # - If either model or defense needs local loading, force batched execution and single worker + # - Otherwise, parallel with configured max_workers + if model_load or defense_load: + trigger = "model.load_model=true" if model_load else "defense.load_model=true" + batched_impl = "local_model" if model_load else "defense_only" + return ResourcePolicy( + strategy="batched", + max_workers=1, + batched_impl=batched_impl, + reason=f"{trigger} -> batched + max_workers=1 (reuse instance, avoid repeated loads)", + ) + + return ResourcePolicy( + strategy="parallel", + max_workers=int(default_max_workers), + batched_impl="none", + reason="no local loading flags -> parallel", + ) + + +def policy_for_test_case_generation( + attack_config: Optional[Dict[str, Any]], + default_max_workers: int = 4, +) -> ResourcePolicy: + """Policy for test case generation stage. + + Project convention: + - Use ONLY `attack_config.load_model` to decide local loading. + - If `load_model: true` => batched + max_workers=1 (reuse attack instance) + - Else => parallel + configured max_workers + """ + cfg = attack_config or {} + attack_load = bool(cfg.get("load_model", False)) + if attack_load: + return ResourcePolicy( + strategy="batched", + max_workers=1, + batched_impl="attack_local", + reason="attack.load_model=true -> batched + max_workers=1 (reuse attack instance)", + ) + return ResourcePolicy( + strategy="parallel", + max_workers=int(default_max_workers), + batched_impl="none", + reason="attack.load_model!=true -> parallel", + ) + + +def policy_for_evaluation( + evaluator_config: Optional[Dict[str, Any]], + default_max_workers: int = 4, +) -> ResourcePolicy: + """Policy for evaluation stage. + + Project convention: + - Use ONLY `evaluator_config.load_model` to decide local loading. + - If `load_model: true` => batched + max_workers=1 (reuse evaluator instance) + - Else => parallel + configured max_workers + """ + cfg = evaluator_config or {} + evaluator_load = bool(cfg.get("load_model", False)) + if evaluator_load: + return ResourcePolicy( + strategy="batched", + max_workers=1, + batched_impl="evaluator_local", + reason="evaluator.load_model=true -> batched + max_workers=1 (reuse evaluator instance)", + ) + return ResourcePolicy( + strategy="parallel", + max_workers=int(default_max_workers), + batched_impl="none", + reason="evaluator.load_model!=true -> parallel", + ) + + diff --git a/pyproject.toml b/pyproject.toml index b1856a6..f85e2ee 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,3 +1,7 @@ +[build-system] +requires = ["setuptools>=61.0", "wheel"] +build-backend = "setuptools.build_meta" + [project] name = "OmniSafeBench-MM" version = "0.1.0" @@ -19,8 +23,7 @@ dependencies = [ "pandas>=2.3.3", "peft==0.17.1", "pillow>=12.0.0", - "pip>=25.3", - "protobuf>=6.33.2", + "protobuf>=6.33.2", # 移除了 pip>=25.3，这不是项目依赖 "pyyaml>=6.0.3", "scikit-image>=0.25.2", "scipy>=1.15.3", @@ -39,10 +42,53 @@ dependencies = [ "wandb>=0.23.1", "webdataset>=1.0.2", ] + +# 关键：包发现配置（解决 flat-layout 问题） +[tool.setuptools] +# 方式1：明确列出所有包（推荐） +packages = [ + "core", + "models", + "config", + "assets", + "attacks", + "dataset", + "pipeline", + "defenses", + "evaluators", + "multimodalmodels", + "dataset_generate", +] + +# 或方式2：使用 find 指令（如果上面的包名正确） +# [tool.setuptools.packages.find] +# where = ["."] +# include = [ +# "core*", +# "models*", +# "config*", +# "assets*", +# "output*", +# "attacks*", +# "dataset*", +# "pipeline*", +# "defenses*", +# "evaluators*", +# "multimodalmodels*", +# "dataset_generate*", +# ] +# exclude = [ +# "tests*", +# "docs*", +# "examples*", +# "build*", +# "dist*", +# "*.egg-info*", +# ] + [tool.pytest.ini_options] testpaths = ["tests"] - [[tool.uv.index]] name = "pytorch-cu118" url = "https://download.pytorch.org/whl/cu118" @@ -50,4 +96,4 @@ explicit = true [tool.uv.sources] torch = { index = "pytorch-cu118" } -torchvision = { index = "pytorch-cu118" } +torchvision = { index = "pytorch-cu118" } \ No newline at end of file diff --git a/tests/test_core.py b/tests/test_core.py index cf21d7a..1e4038c 100644 --- a/tests/test_core.py +++ b/tests/test_core.py @@ -228,33 +228,85 @@ class TestUnifiedRegistry: assert UNIFIED_REGISTRY.validate_attack("non_existent_attack") is False -class TestRegistryInit: - """Test registry initialization""" +class TestPluginsYaml: + """Test config/plugins.yaml based registry initialization""" - def test_registry_init_import(self): - """Test registry_init import""" - from core.registry_init import initialize_registry_lazy_imports + def test_plugins_yaml_loaded(self): + from core.unified_registry import UNIFIED_REGISTRY - # Verify initialization function - mappings = initialize_registry_lazy_imports() + mappings = UNIFIED_REGISTRY._get_lazy_mappings() assert isinstance(mappings, dict) assert "attacks" in mappings assert "models" in mappings assert "defenses" in mappings assert "evaluators" in mappings - # Verify mapping dictionaries exist assert isinstance(mappings["attacks"], dict) assert isinstance(mappings["models"], dict) assert isinstance(mappings["defenses"], dict) assert isinstance(mappings["evaluators"], dict) - # Verify mappings are not empty assert len(mappings["attacks"]) > 0 assert len(mappings["models"]) > 0 assert len(mappings["defenses"]) > 0 assert len(mappings["evaluators"]) > 0 +class TestResourcePolicy: + def test_infer_model_type_from_config(self): + from pipeline.resource_policy import infer_model_type_from_config + + assert infer_model_type_from_config({}) == "api" + assert infer_model_type_from_config({"load_model": False}) == "api" + assert infer_model_type_from_config({"load_model": True}) == "local" + + def test_policy_for_response_generation(self): + from pipeline.resource_policy import policy_for_response_generation + + # default: no load_model flags -> parallel + p = policy_for_response_generation({}, {}, default_max_workers=7) + assert p.strategy == "parallel" + assert p.max_workers == 7 + assert p.batched_impl == "none" + + # model.load_model -> batched + single worker + p = policy_for_response_generation({"load_model": True}, {}, default_max_workers=7) + assert p.strategy == "batched" + assert p.max_workers == 1 + assert p.batched_impl == "local_model" + + # defense.load_model -> batched + single worker + p = policy_for_response_generation({}, {"load_model": True}, default_max_workers=7) + assert p.strategy == "batched" + assert p.max_workers == 1 + assert p.batched_impl == "defense_only" + + def test_policy_for_test_case_generation(self): + from pipeline.resource_policy import policy_for_test_case_generation + + p = policy_for_test_case_generation({}, default_max_workers=9) + assert p.strategy == "parallel" + assert p.max_workers == 9 + assert p.batched_impl == "none" + + p = policy_for_test_case_generation({"load_model": True}, default_max_workers=9) + assert p.strategy == "batched" + assert p.max_workers == 1 + assert p.batched_impl == "attack_local" + + def test_policy_for_evaluation(self): + from pipeline.resource_policy import policy_for_evaluation + + p = policy_for_evaluation({}, default_max_workers=5) + assert p.strategy == "parallel" + assert p.max_workers == 5 + assert p.batched_impl == "none" + + p = policy_for_evaluation({"load_model": True}, default_max_workers=5) + assert p.strategy == "batched" + assert p.max_workers == 1 + assert p.batched_impl == "evaluator_local" + + if __name__ == "__main__": pytest.main([__file__, "-v"])