Aryan-20-04/prompt-injection-safeguard
GitHub: Aryan-20-04/prompt-injection-safeguard
基于DistilBERT的多标签提示词注入攻击检测模型训练框架，支持识别越狱、指令覆盖、角色劫持和数据窃取等攻击类型。
Stars: 1 | Forks: 0
# 使用 Transformers 进行提示词注入检测
此仓库包含用于在 Hugging Face 数据集上训练的**多标签提示词注入检测模型**的完整训练流程。
## 特性
- 提示词攻击的多标签分类
- 检测类型：
- BENIGN
- JAILBREAK
- INSTRUCTION_OVERRIDE
- ROLE_HIJACK
- DATA_EXFILTRATION
- 使用 Hugging Face `Trainer`
- 支持 GPU / Colab 训练
- 使用 Micro & Macro F1 进行评估
## 模型
- 主干网络：DistilBERT
- 任务：多标签文本分类
- 损失函数：BCEWithLogitsLoss (通过 Trainer)
## 结果
在验证集上取得了高性能表现：
- Micro F1 ≈ 0.98
- Macro F1 ≈ 0.95
## 训练
```
python train_prompt_injection_model.py
```
标签：Apex, Atomic Red Team, DistilBERT, Hugging Face, IaC 扫描, NLP, PFX证书, Python, Transformer, 二分类, 凭据扫描, 多标签分类, 大模型安全, 指令覆盖, 攻击检测, 数据泄露检测, 文本分类, 无后门, 机器学习, 模型训练, 深度学习, 系统调用监控, 网络安全, 角色劫持, 越狱检测, 逆向工具, 隐私保护