PoisonCraft

本仓库提供了 POISONCRAFT：针对大语言模型的检索增强生成的实用投毒 的官方实现。

概述

POISONCRAFT 旨在演示恶意行为者如何将“中毒”内容植入检索增强生成（RAG）流水线所使用的语料库中，从而误导大语言模型产生幻觉或引用恶意内容。本代码库提供以下脚本：

• 准备和预处理标准数据集（例如 Natural Questions、MS MARCO、HotpotQA）。

• 将对抗性后缀（即“毒素”）注入类查询文本，以降低或操纵检索结果。

• 在不同检索器（例如 Contriever、SimCSE 和 BGE）下评估投毒效果，并衡量可迁移性。

快速使用

环境配置

以下是配置环境的高级指南。具体要求详见 requirements.txt。为避免冲突，我们建议创建一个新的虚拟环境：

# Example using conda
conda create -n PoisonCraft python=3.10
conda activate PoisonCraft

# Install dependencies
pip install torch==2.4.1 torchvision==0.19.1 torchaudio==2.4.1 --index-url https://download.pytorch.org/whl/cu118
pip install -r requirements.txt

设置 API 密钥与模型配置

本项目的某些部分（如领域分类和查询分类）依赖 API 调用（如 OpenAI GPT 模型）。你需要按如下方式配置 API 密钥和模型设置：

1. 进入 src/models_configs/ 目录。
2. 找到相应的配置文件（例如 gpt3.5_config.py）。
3. 使用你的 API 密钥更新 api_keys 字段：

   "api_key_info": {
       "api_keys": ["your_openai_api_key"],
       "api_key_use": 0
   }
   

注意：本项目支持多种 LLM 提供商（例如 OpenAI、Llama 和 PaLM2）。请确保为你的模型选择相应的配置文件。

如何运行代码

我们在 scripts 文件夹下提供了用于流水线不同阶段的多种 bash 脚本。以下是分步参考：

准备数据集

1. 下载并准备数据集

运行 scripts/run_prepare_dataset.sh 以下载并解压数据集，并创建训练/测试划分：

bash scripts/run_prepare_dataset.sh

该脚本内部会调用：

• experiments/prepare_datasets.py
• experiments/classify_queries_by_domain.py (用于领域分类，如需要)

2. 预处理数据集

要计算每个领域的 top-k 真实相似度分数，请运行：

bash scripts/run_process_data.sh

3. [可选]评估 BEIR 以获取 Ground Truth

仅当你想评估检索器模型时才执行此步骤，因为它可能会花费大量时间。请按照以下步骤操作：

1. 运行 BEIR 评估脚本以计算检索结果：

   python src/beir_eval/eval_beir.py --model_code bge-small --dataset nq --top_k 100
   

2. 使用检索到的结果生成 ground truth 分数：

   bash scripts/run_process_data.sh
   

投毒攻击训练

要生成对抗性后缀（即“毒素”），你可以使用：

bash scripts/run_multi_poisoncraft.sh

• 该脚本将并行启动多次 run_poisoncraft.sh 运行，每次探索不同的对抗性后缀长度和领域索引。
• 核心逻辑调用 experiments/poisoncraft_exp.py

你可以自定义参数，例如：

• DOMAIN_LIST（要投毒的查询的领域索引）
• ADV_LENGTH_LIST（要测试的对抗性后缀长度）

运行检索攻击评估

生成对抗性后缀后，你可以运行以下命令来评估有多少查询获得了“中毒”引用：

bash scripts/run_retrieval_attack.sh

该脚本：

• 加载预先计算好的对抗性后缀。
• 使用 experiments/retrieval_attack.py 运行攻击测试。
• 记录最终检索到中毒文本的查询数量。

运行目标生成攻击评估

或者，对于 LLM 级别的评估（即包含 top-k 检索结果的最终文本生成），你可以使用：

bash scripts/run_target_attack.sh

结果与日志

• 日志存储在 logs/ 目录中。
• 结果（例如中毒查询和攻击成功率）保存在 results/ 目录中。

致谢

• Open-Prompt-Injection: 基础模型代码基于此仓库。
• PoisonedRAG: 部分实现改编自该项目。
• Beir Benchmark: 用作评估的基准框架。
• Contriever: 用于检索增强生成（RAG）能力。
• llms-attacks: 部分对抗性训练的实现源自此仓库。