norrrrrrr-lyn/wainjectbench웹 에이전트를 위한 프롬프트 인젝션 탐지 벤치마킹.
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
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 = 정상)이 포함되어야 합니다. 본 논문의 실험에서는 기본 하이퍼파라미터를 사용합니다.