norrrrrrr-lyn/wainjectbenchБенчмаркинг обнаружения prompt-инъекций для веб-агентов.
WAInjectBench — это комплексный бенчмарк для обнаружения промпт-инъекций в веб-агентах.
Он охватывает 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 = безопасный). В экспериментах из нашей статьи используются гиперпараметры по умолчанию.