norrrrrrr-lyn/wainjectbenchEvaluación comparativa de detecciones de inyección de prompts para agentes web.
WAInjectBench es un benchmark integral para la detección de inyección de prompts en agentes web.
Cubre 6 tipos de ataques, en dos modalidades: texto e imagen.
data/
text/
benign/ → 4 categorías, almacenadas como archivos JSONLmalicious/ → 8 tipos de ataques, almacenados como archivos JSONLimage/
benign/ → 2 categorías, almacenadas en subcarpetasmalicious/ → 7 tipos de ataques, almacenados en subcarpetasClona el repositorio y crea el entorno:
git clone https://github.com/Norrrrrrr-lyn/WAInjectBench.git
cd WAInjectBench
conda env create -f environment.yml
conda activate wainjectbench
WAInjectBench admite dos pipelines de evaluación: detección basada en texto y detección basada en imagen.
python main_text.py \
--data_dir [path to text dataset] \
--detector [detector name] \
--result_dir [output path] \
--gpu [gpu id]
Detectores disponibles: ["kad", "promptarmor", "embedding-t", "promptguard", "datasentinel", "ensemble"]
PromptArmor → requiere OPENAI_API_KEY como variable de entorno.
DataSentinel →
git clone https://github.com/liu00222/Open-Prompt-Injection.git
Descarga el modelo preentrenado en: WAInjectBench/Open-Prompt-Injection/DataSentinel_Models Establece el directorio y la ruta del modelo en detector_text/datasentinel.py.
python main_image.py \
--data_dir [path to image dataset] \
--detector [detector name] \
--result_dir [output path] \
--gpu [gpu id]
Detectores disponibles: ["gpt-4o-prompt", "llava-1.5-7b-prompt", "jailguard", "embedding-i", "llava-1.5-7b-ft", "ensemble"]
GPT-4o-Prompt → requiere OPENAI_API_KEY como variable de entorno.
JailGuard →
git clone https://github.com/shiningrain/JailGuard.git
Sigue su README para configurar MiniGPT4.
LLaVA-1.5-7B-FT → requiere descargar nuestro modelo ajustado y establecer su ruta en detector_image/llava.py.
También proporcionamos versiones entrenadas en dominio de los modelos Embedding-T y Embedding-I, disponibles en model/embedding-t/in-domain y model/embedding-i/in-domain. Para usarlos, sigue el mismo procedimiento de evaluación que en los experimentos principales, pero actualiza la ruta del modelo en detector_text/embedding-t.py y detector_image/embedding-i.py.
Proporcionamos código para entrenar clasificadores binarios basados en embeddings tanto para texto como para imagen.
Clasificador de embeddings de texto
python train/embedding-t.py \
--input_dir [dir with training text jsonl files] \
--output_dir [model output path]
Formato JSONL:
{"text": "example", "label": 1} # 1 for malicious, 0 for benign
Clasificador de embeddings de imagen
python train/embedding-i.py \
--input_dir [dir with training image jsonl files] \
--output_dir [model output path]
Formato JSONL:
{"path": "path/to/image.png", "label": 1}
Ajuste fino 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
Los archivos JSONL deben contener rutas de imágenes y etiquetas (1 = malicioso, 0 = benigno). Los experimentos en nuestro artículo utilizan los hiperparámetros predeterminados.