Al-Amin95/PromptInjectionDetectionSystem
GitHub: Al-Amin95/PromptInjectionDetectionSystem
基于微调 Transformer 模型与 SHAP 可解释 AI 的 Prompt Injection 二元分类检测系统,用于在 LLM 处理前识别恶意 prompt。
Stars: 1 | Forks: 0
# 基于微调 Transformer 模型与可解释 AI 白盒方法的 Prompt Injection 检测
一个基于机器学习的大型语言模型 (LLM) 应用 Prompt Injection 检测系统。该系统旨在使用基于 Transformer 的文本分类模型,将输入 prompt 分类为 **SAFE** 或 **INJECTION**。
主要目标是开发一个准确且可解释的二元文本分类器,用于检测 Prompt Injection 攻击,并利用基于 SHAP 的 token 级解释提供可解释的输出。
## 目标
本项目的目标是通过微调基于 Transformer 的语言模型,并在带标签的 Prompt Injection 数据集上评估其性能,从而构建一个可解释的 Prompt Injection 检测系统。
该系统旨在通过在 LLM 处理之前识别潜在的恶意 prompt,来支持更安全的 LLM 集成应用程序。
## 具体任务
* 准备并分析带标签的 Prompt Injection 数据集,以进行二元文本分类。
* 微调基于 Transformer 的模型,进行 **SAFE/INJECTION** 分类。
* 使用标准分类指标比较所选 Transformer 模型的性能。
* 应用 SHAP 可解释性来识别影响模型预测的重要 token。
* 评估预测性能和解释质量。
* 开发一个简单的 Prompt Injection 检测原型界面。
## 分类任务
本项目采用二元分类方法。
| 标签 | 类别名称 | 含义 |
| ----- | ---------- | ----------------------------------------- |
| `0` | SAFE | 正常或良性 prompt |
| `1` | INJECTION | Prompt Injection 或恶意指令 |
## 数据集
本项目使用 Parquet 格式的 Prompt Injection 数据集。
当前原始数据集文件存储在:
```
data/raw/
```
预期文件:
```
train-00000-of-00001-9564e8b05b4757ab.parquet
test-00000-of-00001-701d16158af87368.parquet
```
初始数据集摘要:
| 数据集划分 | 行数 | 列 |
| ----- | ---: | ----------- |
| Train | 546 | text, label |
| Test | 116 | text, label |
| 总计 | 662 | text, label |
## 模型
本项目计划使用并比较基于 Transformer 的模型,包括:
* **DistilBERT**
* **RoBERTa**
* **SecBERT / 网络安全领域 BERT 变体**
这些模型将针对二元文本分类进行微调。
## 评估指标
分类模型将使用以下指标进行评估:
* Accuracy
* Precision
* Recall
* F1-score
* 混淆矩阵
* AUC-ROC(如适用)
对于 Prompt Injection 检测,**recall** 尤为重要,因为假阴性可能会让恶意 prompt 绕过系统。**Precision** 同样重要,因为过多的假阳性可能会错误地拦截安全的用户 prompt。
## 可解释 AI
本项目使用 **SHAP** 进行可解释性分析。
可解释性阶段旨在展示哪些单词或 token 对模型预测的贡献最大。这通过使检测器的决策更具可解释性,支持了本项目的白盒目标。
计划输出的解释包括:
* Token 级重要性分数
* SHAP 高亮显示的 prompt 示例
* SAFE 与 INJECTION 预测之间的解释对比
* 基于充分性和全面性的解释质量分析
## 项目结构
```
PromptInjectionDetectionSystem/
│
├── data/
│ ├── raw/ # Original dataset files
│ └── processed/ # Cleaned and prepared dataset files
│
├── docs/ # Project documentation and methodology notes
│
├── models/ # Saved model information or checkpoint notes
│
├── notebooks/
│ ├── 00_environment_check.ipynb
│ └── 01_dataset_preparation.ipynb
│
├── results/
│ ├── figures/ # Graphs and visualisations
│ └── tables/ # Dataset summaries and evaluation tables
│
├── screenshots/ # Project screenshots and evidence
│
├── src/ # Reusable Python scripts
│
├── webapp/ # Prototype web application files
│
├── requirements.txt
├── .gitignore
└── README.md
```
## 环境设置
本项目使用 Python 开发。
推荐环境:
* Python 3.10+
* Google Colab 或本地 Python 环境
* 建议使用 GPU 支持进行 Transformer 微调
使用以下命令安装所需的库:
```
pip install -r requirements.txt
```
主要 Python 库:
```
pandas
numpy
pyarrow
openpyxl
matplotlib
seaborn
scikit-learn
datasets
transformers
torch
evaluate
accelerate
shap
tqdm
```
## 当前项目阶段
当前阶段:
```
Data Analysis and Dataset Preparation
```
本阶段已完成或计划的任务:
* [x] 项目文件夹结构已创建
* [x] 原始数据集文件夹已准备就绪
* [ ] 数据集文件已添加
* [x] 数据集加载 notebook 已准备就绪
* [ ] 数据质量检查已完成
* [ ] 类别分布分析已完成
* [ ] 文本长度分析已完成
* [ ] 清洗后的数据集已准备就绪
## 计划工作流
```
1. Load raw dataset
2. Inspect dataset structure
3. Verify labels
4. Check missing values
5. Check duplicate records
6. Analyse class distribution
7. Analyse text length
8. Clean and standardise text
9. Create train, validation, and test splits
10. Save processed dataset
11. Fine-tune transformer models
12. Evaluate model performance
13. Apply SHAP explainability
14. Build prototype detection interface
```
## 项目输出
最终系统预计将提供:
* Prompt 分类结果:**SAFE** 或 **INJECTION**
* 置信度分数
* 模型评估结果
* 使用 SHAP 对重要 token 进行的解释
* 用于演示的简单基于 Web 的原型
## 学术背景
此代码库是网络安全理学硕士 (MSc) 论文项目的一部分。本工作专注于 Prompt Injection 检测、基于 Transformer 的 NLP 分类,以及面向更安全的 LLM 集成系统的可解释 AI。
## 作者
**Al-Amin**
网络安全理学硕士学生
GitHub: [Al-Amin95](https://github.com/Al-Amin95)
## 注意
此代码库属于学术论文项目的一部分。结果将在实验完成并验证后才会进行报告。
标签:Apex, DLL 劫持, 人工智能, 凭据扫描, 可解释AI, 大语言模型, 安全检测工具, 提示词注入检测, 机器学习, 用户模式Hook绕过, 系统调用监控, 逆向工具