git clone https://github.com/Norrrrrrr-lyn/WAInjectBench.git
cd WAInjectBench
conda env create -f environment.yml
conda activate wainjectbench
WAInjectBenchは、テキストベースの検出と画像ベースの検出という2つの評価パイプラインをサポートしています。
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
MiniGPT4を設定するには、そのREADMEに従ってください。
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 = 良性）を含める必要があります。 本論文の実験ではデフォルトのハイパーパラメータを使用しています。