norrrrrrr-lyn/wainjectbenchBenchmarking des détections d'injection de prompt pour les agents web.
WAInjectBench est un benchmark complet pour la détection d'injection de prompt dans les agents web.
Il couvre 6 types d'attaques, à travers deux modalités : texte et image.
data/
text/
benign/ → 4 catégories, stockées sous forme de fichiers JSONLmalicious/ → 8 types d'attaques, stockés sous forme de fichiers JSONLimage/
benign/ → 2 catégories, stockées dans des sous-dossiersmalicious/ → 7 types d'attaques, stockés dans des sous-dossiersClonez le dépôt et créez l'environnement :
git clone https://github.com/Norrrrrrr-lyn/WAInjectBench.git
cd WAInjectBench
conda env create -f environment.yml
conda activate wainjectbench
WAInjectBench prend en charge deux pipelines d'évaluation : la détection basée sur le texte et la détection basée sur l'image.
python main_text.py \
--data_dir [path to text dataset] \
--detector [detector name] \
--result_dir [output path] \
--gpu [gpu id]
Détecteurs disponibles : ["kad", "promptarmor", "embedding-t", "promptguard", "datasentinel", "ensemble"]
PromptArmor → nécessite la variable d'environnement OPENAI_API_KEY.
DataSentinel →
git clone https://github.com/liu00222/Open-Prompt-Injection.git
Téléchargez le modèle pré-entraîné dans : WAInjectBench/Open-Prompt-Injection/DataSentinel_Models Définissez le répertoire et le chemin du modèle dans detector_text/datasentinel.py.
python main_image.py \
--data_dir [path to image dataset] \
--detector [detector name] \
--result_dir [output path] \
--gpu [gpu id]
Détecteurs disponibles : ["gpt-4o-prompt", "llava-1.5-7b-prompt", "jailguard", "embedding-i", "llava-1.5-7b-ft", "ensemble"]
GPT-4o-Prompt → nécessite la variable d'environnement OPENAI_API_KEY.
JailGuard →
git clone https://github.com/shiningrain/JailGuard.git
Suivez son README pour configurer MiniGPT4.
LLaVA-1.5-7B-FT → nécessite le téléchargement de notre modèle affiné et la définition de son chemin dans detector_image/llava.py.
Nous fournissons également des versions entraînées in-domain des modèles Embedding-T et Embedding-I, disponibles dans model/embedding-t/in-domain et model/embedding-i/in-domain. Pour les utiliser, suivez la même procédure d'évaluation que dans les expériences principales, mais mettez à jour le chemin du modèle dans detector_text/embedding-t.py et detector_image/embedding-i.py.
Nous fournissons le code pour entraîner des classifieurs binaires basés sur les embeddings, pour le texte et l'image.
Classifieur d'embeddings de texte
python train/embedding-t.py \
--input_dir [dir with training text jsonl files] \
--output_dir [model output path]
Format JSONL :
{"text": "example", "label": 1} # 1 for malicious, 0 for benign
Classifieur d'embeddings d'image
python train/embedding-i.py \
--input_dir [dir with training image jsonl files] \
--output_dir [model output path]
Format JSONL :
{"path": "path/to/image.png", "label": 1}
Affinage de 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
Les fichiers JSONL doivent contenir les chemins d'images et les étiquettes (1 = malveillant, 0 = bénin). Les expériences de notre article utilisent les hyperparamètres par défaut.