WAInjectBench

WAInjectBench é um benchmark abrangente para detecção de prompt injection em agentes web.
Ele cobre 6 tipos de ataques, em duas modalidades: texto e imagem.

---

📂 Estrutura do Dataset

• data/
  • text/
    • benign/ → 4 categorias, armazenadas como arquivos JSONL
    • malicious/ → 8 tipos de ataque, armazenados como arquivos JSONL
  • image/
    • benign/ → 2 categorias, armazenadas em subpastas
    • malicious/ → 7 tipos de ataque, armazenados em subpastas

---

⚙️ Instalação

Clone o repositório e crie o ambiente:

git clone https://github.com/Norrrrrrr-lyn/WAInjectBench.git
cd WAInjectBench
conda env create -f environment.yml
conda activate wainjectbench

🚀 Avaliação

O WAInjectBench suporta dois pipelines de avaliação: detecção baseada em texto e detecção baseada em imagem.

🔹 Detecção baseada em texto

python main_text.py \
  --data_dir [path to text dataset] \
  --detector [detector name] \
  --result_dir [output path] \
  --gpu [gpu id]

Detectores disponíveis: ["kad", "promptarmor", "embedding-t", "promptguard", "datasentinel", "ensemble"]

PromptArmor → requer OPENAI_API_KEY como variável de ambiente.

DataSentinel →

git clone https://github.com/liu00222/Open-Prompt-Injection.git

Baixe o modelo pré-treinado para: WAInjectBench/Open-Prompt-Injection/DataSentinel_Models Defina o diretório e o caminho do modelo em detector_text/datasentinel.py.

🔹 Detecção baseada em imagem

python main_image.py \
  --data_dir [path to image dataset] \
  --detector [detector name] \
  --result_dir [output path] \
  --gpu [gpu id]

Detectores disponíveis: ["gpt-4o-prompt", "llava-1.5-7b-prompt", "jailguard", "embedding-i", "llava-1.5-7b-ft", "ensemble"]

GPT-4o-Prompt → requer OPENAI_API_KEY como variável de ambiente.

JailGuard →

git clone https://github.com/shiningrain/JailGuard.git

Siga o README dele para configurar o MiniGPT4.

LLaVA-1.5-7B-FT → requer o download do nosso modelo ajustado e a definição do seu caminho em detector_image/llava.py.

🔹 Generalização in-domain

Também fornecemos versões treinadas in-domain dos modelos Embedding-T e Embedding-I, disponíveis em model/embedding-t/in-domain e model/embedding-i/in-domain. Para usá-los, siga o mesmo procedimento de avaliação dos experimentos principais, mas atualize o caminho do modelo em detector_text/embedding-t.py e detector_image/embedding-i.py.

🏋️ Treinamento

Fornecemos código para treinar classificadores binários baseados em embeddings tanto para texto quanto para imagem.

Classificador de embedding 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

Classificador de embedding de imagem

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}

Finetuning do 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

Os arquivos JSONL devem conter caminhos de imagens e rótulos (1 = malicioso, 0 = benigno). Os experimentos em nosso artigo usam os hiperparâmetros padrão.