Paul-ans-typ/Prompt-Injection-Detection-Supervised-Learning
GitHub: Paul-ans-typ/Prompt-Injection-Detection-Supervised-Learning
基于监督学习的 Prompt 注入攻击检测方案，对比传统机器学习与微调 BERT 模型在恶意 Prompt 分类任务上的效果差异。
Stars: 0 | Forks: 0
# LLM 软件服务中的 Prompt 注入攻击检测
## 概述
嵌入在软件服务（如客户支持聊天机器人和代码生成工具）中的 Large Language Models (LLMs) 引入了新的安全风险，其中最显著的是 **Prompt 注入攻击**。在这些攻击中，恶意用户精心构造输入，旨在覆盖系统指令、绕过安全控制或提取敏感信息。
本项目实现了一个离线的监督 Machine Learning Pipeline，用于检测这些基于文本的 Prompt 注入攻击。它建立了一个传统的 Machine Learning 基线，并将其与微调后的 Deep Learning 方法进行对比，以将 Prompt 分类为良性或恶意。
## 数据集
本项目使用 **Malicious Prompt Detection Dataset (MPDD)**。
* **大小：** 39,234 条文本 Prompt。
* **类别平衡：** 完美平衡（50% 良性，50% 恶意）。
* **格式：** 二分类（`isMalicious` 标签）。
## 仓库结构
```
prompt-injection-detection/
├── .gitignore
├── .env.example # Template for environment variables (Kaggle API key)
├── README.md # Project documentation
├── requirements.txt # Python dependencies
│
├── data/ # Ignored by Git
│ ├── raw/ # Raw MPDD.csv dataset
│ └── processed/ # 80/10/10 Train/Val/Test splits
│
├── models/ # Ignored by Git (Saved model weights & vectorizers)
│
├── notebooks/ # Jupyter notebooks for EDA and evaluation visuals
│
└── src/ # Core Python scripts
├── download_data.py # Script to securely download the dataset
├── preprocess.py # Data cleaning and train/test splitting
├── train_baseline.py # TF-IDF + Logistic Regression training
├── train_bert.py # BERT fine-tuning using PyTorch
└── evaluate.py # Model comparison and confusion matrix generation
```
## 已实现的模型
1. **基线模型 (TF-IDF + Logistic Regression)：** 因其轻量级占用、可解释性和速度而被选用。使用最大特征数为 10,000 的 TF-IDF Vectorizer。
2. **Deep Learning 模型 (BERT)：** 一个针对序列分类进行微调的 `bert-base-uncased` Transformer 模型。使用 PyTorch 和 AdamW Optimizer 进行训练。
## 结果与评估
两个模型均在包含 3,924 条 Prompt 的未见测试集上进行了评估。
* **Logistic Regression 基线：**
* 达到了约 95.2% 的总体准确率。
* **False Negatives：** 145
* **False Positives：** 44
* **微调后的 BERT：**
* 达到了约 97.5% 的总体准确率。
* **False Negatives：** 37
* **False Positives：** 62
**关键结论：** 虽然 Logistic Regression 基线表现得出奇地好，但 BERT 模型明显更适合安全场景。与基线相比，BERT 将 False Negatives（即未被检测到的实际恶意 Prompt）减少了近 75%。在安全环境中，最大限度地减少 False Negatives 对于防止成功的 Prompt 注入攻击至关重要。
## 如何运行
1. **环境设置：**
```
python -m venv .venv
source .venv/bin/activate # On Windows use: .venv\Scripts\activate
pip install -r requirements.txt
```
2. **下载数据：**
你需要自己的 Kaggle API 密钥来下载数据集。使用你的凭据设置 `.env` 文件（参考 `.env.example`），然后运行：
```
python src/download_data.py
```
3. **数据预处理：**
处理原始数据并生成 train/val/test Split：
```
python src/preprocess.py
```
4. **训练模型：**
```
python src/train_baseline.py
python src/train_bert.py
```
5. **评估：**
```
python src/evaluate.py
```
## 未来工作
随着项目的扩展，计划的未来改进包括：
* **实时 API 集成：** 将训练好的 BERT 模型封装在 FastAPI 或 Flask 应用程序中，作为实际 LLM 应用程序的中间件安全层。
* **在未见攻击向量上测试：** 针对较新的、分布外的越狱技术（例如 base64 编码攻击、多语言绕过）评估模型。
* **模型量化：** 使用 ONNX 或 TensorRT 等技术减少 BERT 模型的内存占用，以降低生产环境中的推理延迟。
* **探索更小的 LLM 进行检测：** 测试本地的、无审查的 7B/8B 参数模型作为 Few-shot 分类器，观察它们是否优于微调后的 BERT 表示。
标签：AI安全, Apex, API密钥检测, BERT, Chat Copilot, DLL 劫持, LLM, MPDD, NLP, NoSQL, PyTorch, TF-IDF, Unmanaged PE, URL发现, 二分类, 人工智能安全, 凭据扫描, 合规性, 大语言模型, 安全检测, 对抗攻击, 恶意提示检测, 恶意检测, 提示注入, 敏感信息检测, 数据科学, 文本分类, 机器学习, 模型微调, 深度学习, 监督学习, 策略即代码, 网络安全, 聊天机器人安全, 资源验证, 输入验证, 逆向工具, 逻辑回归, 隐私保护, 集群管理