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 = malicious, 0 = benign)। हमारे पेपर में प्रयोग डिफ़ॉल्ट हाइपरपैरामीटर का उपयोग करते हैं।