WAInjectBench 是一个面向 Web 智能体中提示注入检测的综合基准。 它涵盖6 种攻击类型，跨越两种模态：文本和图像。
data/
text/
benign/ → 4 个类别，以 JSONL 文件存储malicious/ → 8 种攻击类型，以 JSONL 文件存储image/
benign/ → 2 个类别，存储在子文件夹中malicious/ → 7 种攻击类型，存储在子文件夹中克隆仓库并创建环境：
git clone https://github.com/Norrrrrrr-lyn/WAInjectBench.git
cd WAInjectBench
conda env create -f environment.yml
conda activate wainjectbench
WAInjectBench 支持两种评估流程：基于文本的检测和基于图像的检测。
python main_text.py \
--data_dir [path to text dataset] \
--detector [detector name] \
--result_dir [output path] \
--gpu [gpu id]
可用检测器： ["kad", "promptarmor", "embedding-t", "promptguard", "datasentinel", "ensemble"]
PromptArmor → 需要将 OPENAI_API_KEY 设置为环境变量。
DataSentinel →
git clone https://github.com/liu00222/Open-Prompt-Injection.git
将预训练模型下载到：WAInjectBench/Open-Prompt-Injection/DataSentinel_Models 在 detector_text/datasentinel.py 中设置目录和模型路径。
python main_image.py \
--data_dir [path to image dataset] \
--detector [detector name] \
--result_dir [output path] \
--gpu [gpu id]
可用检测器： ["gpt-4o-prompt", "llava-1.5-7b-prompt", "jailguard", "embedding-i", "llava-1.5-7b-ft", "ensemble"]
GPT-4o-Prompt → 需要将 OPENAI_API_KEY 设置为环境变量。
JailGuard →
git clone https://github.com/shiningrain/JailGuard.git
按照其 README 配置 MiniGPT4。
LLaVA-1.5-7B-FT → 需要下载我们微调后的模型，并在 detector_image/llava.py 中设置其路径。
我们还提供了 Embedding-T 和 Embedding-I 模型的域内训练版本，位于 model/embedding-t/in-domain 和 model/embedding-i/in-domain。要使用它们，请遵循与主实验相同的评估流程，但需更新 detector_text/embedding-t.py 和 detector_image/embedding-i.py 中的模型路径。
我们提供了用于训练基于嵌入的文本和图像二分类器的代码。
文本嵌入分类器
python train/embedding-t.py \
--input_dir [dir with training text jsonl files] \
--output_dir [model output path]
JSONL 格式：
{"text": "example", "label": 1} # 1 for malicious, 0 for benign
图像嵌入分类器
python train/embedding-i.py \
--input_dir [dir with training image jsonl files] \
--output_dir [model output path]
JSONL 格式：
{"path": "path/to/image.png", "label": 1}
微调 LLaVA-1.5-7B
python train.py \
--train_jsonl train.jsonl \
--val_jsonl val.jsonl \
--use_lora \
--amp_dtype bf16 \
--device_mode single \
--gpu_id 0
JSONL 文件应包含图像路径和标签（1 = 恶意，0 = 良性）。 我们论文中的实验使用默认超参数。