【大模型RAG安全基准】安装和使用SafaRAG框架

1 介绍

SafeRAG 是首个专注于系统化评估检索增强生成（RAG）管道安全风险的基准测试框架，特别针对中文场景。它揭示了攻击者如何通过操纵 RAG 流程中的数据流，导致模型生成不准确、有误导性甚至有害的内容。

SafeRAG 框架由以下几个核心部分组成:

1.1 检索器 (Retrievers)

检索器负责从知识库中检索相关文档。框架提供四种检索器:

BaseRetriever: 基础向量检索器,使用嵌入模型进行语义检索
CustomBM25Retriever: 基于 BM25 算法的关键词检索器
EnsembleRetriever: 混合检索器,结合向量和关键词检索
EnsembleRerankRetriever: 混合重排序检索器,在混合检索基础上增加重排序

1.2 过滤器 (Filters)

过滤器用于过滤检索到的上下文,提高内容质量。框架支持三种过滤模式:

off: 不使用过滤器
nli: 使用自然语言推理(NLI)过滤器
skr: 使用 SKR 压缩器过滤冗余信息

1.3 生成器 (Generators/LLMs)

生成器是大语言模型,负责根据检索到的上下文生成答案 7 。框架支持:

API 模型: GPT 系列(gpt-3.5-turbo, gpt-4, gpt-4o)、DeepSeek
本地模型: Qwen 7B/14B、Baichuan2 13B、ChatGLM3 6B

1.4 攻击任务 (Attack Tasks)

框架实现四种攻击任务来评估 RAG 系统的安全性 :

Silver_noise (SN): 噪声攻击,注入低质量文本
Inter_context_conflict (ICC): 上下文冲突攻击,注入矛盾信息
Soft_ad (SA): 软广告攻击,注入毒性内容
White_DoS (WDoS): 拒绝服务攻击,导致模型拒绝回答

1.5 评估系统 (Evaluation System)

评估系统由 BaseEvaluator 类协调,负责执行评估流程 :

多线程并行处理评估样本
计算多种评估指标(F1、RA、AFR、BERTScore、QuestEval)
生成 JSON 格式的评估报告

1.6 评估指标 (Metrics)

框架实现多种评估指标 :

BLEU: 衡量生成文本与参考答案的 n-gram 重叠度
ROUGE: 评估文本摘要质量
BERTScore: 基于语义嵌入的相似度评估
QuestEval: 基于 LLM 的多选题评估

1.7 知识库 (Knowledge Base)

知识库包含干净文本和攻击文本 :

每个攻击任务有独立的知识库目录(SN/ICC/SA/WDoS)
db.txt: 干净的知识库文本
add_* 目录: 包含攻击文本

2 安装

必须是Linux系统

bash 复制代码

conda create -n saferag python=3.10 -y
conda activate saferag
git clone https://githSafeRAGub.com/IAAR-Shanghai/SafeRAG.git
cd SafeRAG

pip install torch==2.6.0+cu126 torchvision==0.21.0+cu126 torchaudio==2.6.0+cu126 --index-url https://download.pytorch.org/whl/cu126 --upgrade

pip install -U pip setuptools wheel
pip install -r requirements.txt

安装Elasticsearch,SafeRAG的BM25检索器依赖Elasticsearch服务

bash 复制代码

sudo apt-get install apt-transport-https
wget https://artifacts.elastic.co/downloads/elasticsearch/elasticsearch-9.2.0-amd64.deb
wget https://artifacts.elastic.co/downloads/elasticsearch/elasticsearch-9.2.0-amd64.deb.sha512
shasum -a 512 -c elasticsearch-9.2.0-amd64.deb.sha512
sudo dpkg -i elasticsearch-9.2.0-amd64.deb

#4. 启动Elasticsearch服务  
 sudo systemctl daemon-reload
 sudo systemctl enable elasticsearch.service
 sudo systemctl start elasticsearch.service
  #5. 验证服务运行  
curl -X GET "localhost:9200/"

安装Qwen/Qwen-7B-Chat需要的文件

bash 复制代码

pip install transformers==4.32.0 accelerate tiktoken einops scipy transformers_stream_generator==0.0.4 peft deepspeed
pip install sentence-transformers==2.2.2
pip install huggingface-hub==0.25.2

从https://github.com/Dao-AILab/flash-attention/releases/tag/v2.8.3安装编译好的flash-attention，对应以上的Python、torch和cuda版本选择下载，例如"flash_attn-2.8.3+cu12torch2.6cxx11abiFALSE-cp310-cp310-linux_x86_64.whl"

bash 复制代码

pip install flash_attn-2.8.3+cu12torch2.6cxx11abiFALSE-cp310-cp310-linux_x86_64.whl

3 运行

注意该框架的nctd.json文件在nctd_datasets/nctd.json路径下

3.1 启动 Milvus-lite 服务

SafeRAG 使用 milvus-lite 作为向量数据库,需要启动该服务:

bash 复制代码

milvus-server

如果要使用BM25检索器，需要启动elasticsearch 服务

bash 复制代码

sudo systemctl start elasticsearch  
sudo systemctl enable elasticsearch

需要下载 bge-base-zh-v1.5 模型到本地目录。在运行项目之前,需要修改 configs/config.py 文件,填入 API 密钥或本地模型路径，配置文件用于初始化 RAG 系统中大语言模型的加载参数。

txt 复制代码

GPT_transit_url = ""
DeepSeek_key = ''
DeepSeek_base = 'https://api.deepseek.com'
GPT_api_key = ''
GPT_api_base = '' 
Qwen_7B_local_path = '/home/user/SafeRAG/QwenQwen-7B-Chat'
Qwen_14B_local_path = ''
Baichuan2_13b_local_path = ''
ChatGLM3_local_path = ''

3.2 使用本地模型

使用Qwen-7B-Chat模型来测试
从hugging face下载模型文件，放在根目录Qwen-7B-Chat

https://huggingface.co/Qwen/Qwen-7B-Chat/tree/main
修改版本冲突
由于Qwen-7B-Chat和SafeRAG需要的FlagEmbedding版本存在冲突，为了避免使用FlagEmbedding，将retrievers/init.py和quick_start_nctd.py中的以下代码注释掉，避免使用修改避免导入 FlagReranker，然后只使用 bm25、base 或 hybrid 检索器。

from .hybrid import EnsembleRetriever

from .hybrid_rerank import EnsembleRerankRetriever

from retrievers import BaseRetriever, CustomBM25Retriever, EnsembleRetriever, EnsembleRerankRetriever

from retrievers import BaseRetriever, CustomBM25Retriever
执行快速启动脚本
配置完成后,运行 quick_start_nctd.py 脚本:

bash 复制代码

python quick_start_nctd.py\
    --model_name qwen7b\
    --retriever_name base\
    --retrieve_top_k 6\
    --filter_module off\
    --attack_task SN\
    --attack_module indexing\
    --attack_intensity 0.5\
    --embedding_name /home/user/SafeRAG/bge-base-zh-v1.5\
    --num_threads 1\
    --show_progress_bar True
    --attack_data_path ./nctd_datasets/nctd.json

主要配置参数说明

索引设置: --clean_docs_path, --chunk_size, --chunk_overlap 用于配置知识库路径和文档分块

检索器设置: --retriever_name (支持 'base', 'bm25', 'hybrid', 'hybrid-rerank'), --retrieve_top_k, --embedding_name

过滤器设置: --filter_module (支持 'off', 'nli', 'skr')

生成器设置: --model_name (支持 'gpt-3.5-turbo', 'gpt-4', 'deepseek-chat', 'qwen7b' 等)

攻击设置: --attack_task (SN/ICC/SA/WDoS), --attack_module (indexing/retrieval/generation), --attack_intensity

3.3 使用API模型

如何使用deepseek的模型，将deepseek中的API_key添加到以下配置文件中即可

config.py中配置如下

txt 复制代码

GPT_transit_url = ""
DeepSeek_key = '自己的API_KEY'
DeepSeek_base = ''
GPT_api_key = ''
GPT_api_base = '' 
Qwen_7B_local_path = ''
Qwen_14B_local_path = ''
Baichuan2_13b_local_path = ''
ChatGLM3_local_path = ''

执行脚本

bash 复制代码

python quick_start_nctd.py \  
    --model_name deepseek-chat \  
    --retriever_name base\  
    --retrieve_top_k 6 \  
    --filter_module off \  
    --attack_task SN \  
    --attack_module indexing \  
    --attack_intensity 0.5 \  
    --embedding_name /home/user/SafeRAG/bge-base-zh-v1.5 \  
    --num_threads 1 \  
    --show_progress_bar True
    --attack_data_path ./nctd_datasets/nctd.json

3.4 自定义API模型

使用kimi模型，需要在SafeRAG中添加一个新的Kimi模型类,参考现有的DeepSeek类实现。

（1）创建Kimi类 (llms/api_model.py)

在llms/api_model.py文件中添加Kimi类,参考DeepSeek的实现模式:

python 复制代码

class Kimi(BaseLLM):
    def __init__(self, model_name='kimi-k2-turbo-preview', temperature=0.01, max_new_tokens=4096, report=False):
        super().__init__(model_name, temperature, max_new_tokens)
        self.report = report
        self.model_name = model_name
    
    def request(self, query: str) -> str:
        kimi = openai.OpenAI(api_key=conf.Kimi_key, base_url=conf.Kimi_base)
        res = kimi.chat.completions.create(
            model=self.params['model_name'],
            messages=[{"role": "system", "content": query}],
            temperature=self.params['temperature'],
            max_tokens=self.params['max_new_tokens'],
            top_p=self.params['top_p'],
            stream=False
        )
        real_res = res.choices[0].message.content
        token_consumed = res.usage.total_tokens
        logger.info(f'Kimi token consumed: {token_consumed}') if self.report else ()
        return real_res

（2）配置API密钥 (configs/config.py)

在配置文件中添加Kimi的API密钥和base URL:

python 复制代码

Kimi_key = 'your-kimi-api-key'
Kimi_base = 'https://api.moonshot.cn/v1'

（3）导入Kimi类 (quick_start_nctd.py)

在主入口文件中导入Kimi类:

python 复制代码

from llms.api_model import GPT, DeepSeek, Kimi

（4）添加模型选择逻辑 (quick_start_nctd.py)

在模型选择部分添加Kimi的条件分支

python 复制代码

elif args.model_name == "kimi-k2-turbo-preview":
    llm = Kimi(model_name=args.model_name, temperature=args.temperature, max_new_tokens=args.max_new_tokens)

（5）使用方法

运行时指定Kimi模型:

bash 复制代码

python quick_start_nctd.py\
    --model_name 'kimi-k2-turbo-preview'\
    --retriever_name base\
    --retrieve_top_k 6\
    --filter_module off\
    --attack_task SN\
    --attack_module indexing\
    --attack_intensity 0.5\
    --embedding_name /home/user/SafeRAG/bge-base-zh-v1.5\
    --num_threads 1\
    --show_progress_bar True\
    --attack_data_path ./nctd_datasets/nctd.json

4 nctd.json

nctd.json 文件是 SafeRAG 框架中的攻击数据集文件。

nctd.json 是包含四种攻击任务(Noise、Conflict、Toxicity、DoS)评估数据的 JSON 文件。该文件通过 get_task_datasets() 函数加载,根据 attack_task 参数(SN/ICC/SA/WDoS)提取相应的评估数据。

由于原文作者并没有提供该文件，需要自定义，以下是示例文件。

存放在nctd_datasets/nctd.json目录下。并在脚本中通过--attack_data_path指定路径

bash 复制代码

python quick_start_nctd.py\
    --model_name qwen7b\
    --retriever_name base\
    --retrieve_top_k 6\
    --filter_module off\
    --attack_task SN\
    --attack_module indexing\
    --attack_intensity 0.5\
    --embedding_name /home/user/SafeRAG/bge-base-zh-v1.5\
    --num_threads 1\
    --show_progress_bar True\
    --attack_data_path ./nctd_datasets/nctd.json

【大模型RAG安全基准】安装和使用SafaRAG框架

1 介绍

1.1 检索器 (Retrievers)

1.2 过滤器 (Filters)

1.3 生成器 (Generators/LLMs)

1.4 攻击任务 (Attack Tasks)

1.5 评估系统 (Evaluation System)

1.6 评估指标 (Metrics)

1.7 知识库 (Knowledge Base)

2 安装

3 运行

注意该框架的nctd.json文件在nctd_datasets/nctd.json路径下

3.1 启动 Milvus-lite 服务

3.2 使用本地模型

from .hybrid import EnsembleRetriever

from .hybrid_rerank import EnsembleRerankRetriever

from retrievers import BaseRetriever, CustomBM25Retriever, EnsembleRetriever, EnsembleRerankRetriever

3.3 使用API模型

3.4 自定义API模型

4 nctd.json