⚡FlashRAG: 一个高效的RAG研究Python工具包

[ English | 中文 ]

安装 | 特性 | 快速开始 | 组件 | FlashRAG-UI | 支持的基线方法 | 支持的数据集 | 常见问题

FlashRAG是一个用于复现和开发检索增强生成（RAG）研究的Python工具包。我们的工具包包括36个预处理的基准RAG数据集和17个最先进的RAG算法。

通过FlashRAG和提供的资源，您可以轻松复现现有的SOTA工作，或实现自定义的RAG流程和组件。我们也提供了一个简介易用的可视化界面：

https://github.com/user-attachments/assets/e34a90c0-5e0f-4f76-886e-e30de8a6f8d1

:link: 导航

• 特性
• 开发路线图
• 更新日志
• 安装
• 快速开始
• 组件
• FlashRAG-UI
• 基线方法
• 支持的数据集和文档语料库
• 其他常见问题
• 许可证
• 引用

:sparkles: 特性

• 广泛且可定制的框架：包括RAG场景的基本组件，如检索器、重排序器、生成器和压缩器，允许灵活组装复杂的管道。

• 全面的基准数据集：收集了36个预处理的RAG基准数据集，用于测试和验证RAG模型的性能。

• 预实现的先进RAG算法：基于我们的框架，提供17个先进的RAG算法及其报告结果。轻松在不同设置下复现结果。

• 高效的预处理过程：通过提供各种脚本，如语料处理、索引构建和文档预检索，简化RAG工作流准备。

• 优化的执行效率：库的效率通过vLLM、FastChat用于LLM推理加速和Faiss用于向量索引管理得到增强。

• 方便易用的UI界面：我们开发了一个十分方便易用的UI界面，以供轻松、快捷地配置和体验我们已经实现的RAG方法，以及在可视化界面上运行评估脚本。

:mag_right: 开发路线图

FlashRAG仍在开发中，存在许多问题和改进空间。我们将继续更新，并诚挚欢迎对这个开源工具包的贡献。

• [x] 支持OpenAI模型
• [x] 提供每个组件的使用说明
• [x] 集成sentence-transformers
• [ ] 包含更多RAG方法
• [ ] 添加更多评估指标（如Unieval、命名实体F1）和基准（如RGB基准）
• [ ] 增强代码的适应性和可读性

:page_with_curl: 更新日志

[25/03/21] 我们新增了对推理流程的支持，这是一种结合推理能力和检索的新范式，代表工作包括Search-o1、R1-Searcher、ReSearch。我们在各种RAG基准上评估了流程的性能，在多跳推理数据集（如Hotpotqe）上可以达到接近60的F1分数。详见结果表。

[25/02/24] 我们新增了对多模态RAG的支持，包括Llava、Qwen、InternVL等MLLM，以及基于Clip结构的各类检索器。

[25/01/21] 我们的论文 FlashRAG: A Python Toolkit for Efficient RAG Research 已经被 Resource Track of the 2025 ACM Web Conference (WWW 2025) 接收!

[25/01/12] 介绍 FlashRAG-UI，一个易于使用的界面。您可以通过它轻松快捷地配置和体验支持的RAG方法，并在基准上评估这些方法。

[25/01/11] 我们添加了对新方法RQRAG方法的支持，详见复现实验。

[25/01/07] 我们目前支持多个检索器的聚合，详见多检索器使用。

[25/01/07] 我们集成了一个非常灵活且轻量级的语料库分块库Chunkie，支持各种自定义分块方法（tokens、句子、语义等）。在分块文档语料库中使用它。

[24/10/21] 我们发布了一个基于Paddle框架的版本，支持中国硬件平台。详情请参阅FlashRAG Paddle。

[24/10/13] 添加了一个新的领域内数据集和语料库 - DomainRAG，基于中国人民大学的内部招生数据，涵盖七种任务类型，可用于进行领域特定的RAG测试。

[24/09/24] 我们发布了一个基于MindSpore框架的版本，支持中国硬件平台。详情请参阅FlashRAG MindSpore。

显示更多

[24/09/18] 由于在某些环境中安装Pyserini的复杂性和限制，我们引入了一个轻量级的BM25s包作为替代（更快且更易使用）。基于Pyserini的检索器将在未来版本中被弃用。要使用bm25s作为检索器，只需在配置中将bm25_backend设置为bm25s。

[24/09/09] 我们添加了对新方法Adaptive-RAG的支持，该方法可以根据查询类型自动选择执行的RAG流程。请参阅结果表。

[24/08/02] 我们添加了对新方法Spring的支持，通过仅添加少量的token嵌入显著提高LLM的性能。请参阅结果表。

[24/07/17] 由于HuggingFace的一些未知问题，我们的原始数据集链接已失效。我们已更新链接。如果遇到任何问题，请检查新链接。

[24/07/06] 我们添加了对新方法Trace的支持，通过构建知识图谱来优化文本。请参阅结果和详情。

[24/06/19] 我们添加了对新方法IRCoT的支持，并更新了结果表。

[24/06/15] 我们提供了一个演示来使用我们的工具包执行RAG流程。

[24/06/11] 我们在检索器模块中集成了句子转换器。现在使用检索器时无需设置池化方法。

[24/06/05] 我们提供了详细的文档来复现现有方法（参见如何复现，基线详情），以及配置设置。

[24/06/02] 我们为初学者提供了FlashRAG的介绍，参见FlashRAG介绍（中文版 한국어）。

[24/05/31] 我们支持Openai系列模型作为生成器。

:wrench: 安装

要开始使用FlashRAG，您可以通过pip简单安装：

pip install flashrag-dev --pre

或者您可以从Github克隆并安装（需要Python 3.10+）：

git clone https://github.com/RUC-NLPIR/FlashRAG.git
cd FlashRAG
pip install -e .

如果您想使用vllm、句子转换器或pyserini，可以安装可选依赖项：

# 安装所有额外依赖项
pip install flashrag[full]

# 安装vllm以提高速度
pip install vllm>=0.4.1

# 安装sentence-transformers
pip install sentence-transformers

# 安装pyserini(使用bm25检索)
pip install pyserini

由于使用pip安装faiss时的不兼容性，请使用以下conda命令进行安装。

# 仅CPU版本
conda install -c pytorch faiss-cpu=1.8.0

# GPU(+CPU)版本
conda install -c pytorch -c nvidia faiss-gpu=1.8.0

注意：在某些系统上无法安装最新版本的faiss。

来自官方Faiss仓库（来源）：

• 仅CPU的faiss-cpu conda包目前在Linux（x86_64和arm64）、OSX（仅arm64）和Windows（x86_64）上可用
• 包含CPU和GPU索引的faiss-gpu在Linux（仅x86_64）上可用，适用于CUDA 11.4和12.1

:rocket: 快速开始

语料库构建

要构建索引，首先需要将语料库保存为如下格式的jsonl文件，每行是一个文档。

{"id": "0", "contents": "..."}
{"id": "1", "contents": "..."}

如果您想使用维基百科作为语料库，可以参考我们的文档处理维基百科将其转换为索引格式。

索引构建

您可以使用以下代码构建您自己的索引。

• 对于稠密检索方法，特别是流行的嵌入模型，我们使用faiss来构建索引。

• 对于稀疏检索方法 (BM25)，我们基于Pyserini或bm25s将语料库构建为Lucene倒排索引。构建的索引包含原始文档。

对于稠密检索方法

修改以下代码中的参数为您的参数。

python -m flashrag.retriever.index_builder \
    --retrieval_method e5 \
    --model_path /model/e5-base-v2/ \
    --corpus_path indexes/sample_corpus.jsonl \
    --save_dir indexes/ \
    --use_fp16 \
    --max_length 512 \
    --batch_size 256 \
    --pooling_method mean \
    --faiss_type Flat 

• --pooling_method: 如果未指定此参数，我们将根据模型名称和模型文件自动选择。然而，由于不同嵌入模型使用的池化方法不同，我们可能没有完全实现它们。为了确保准确性，您可以指定您使用的检索模型对应的池化方法（mean、pooler 或 cls）。

• ---instruction: 一些嵌入模型在编码前需要附加指令来连接查询，可以在此处指定。目前，我们将自动为 E5 和 BGE 模型填写指令，而其他模型需要手动补充。

如果检索模型支持 sentence transformers 库，您可以使用以下代码构建索引（无需考虑池化方法）。

python -m flashrag.retriever.index_builder \
    --retrieval_method e5 \
    --model_path /model/e5-base-v2/ \
    --corpus_path indexes/sample_corpus.jsonl \
    --save_dir indexes/ \
    --use_fp16 \
    --max_length 512 \
    --batch_size 256 \
    --pooling_method mean \
    --sentence_transformer \
    --faiss_type Flat 

对于稀疏检索方法 (BM25)

如果构建bm25索引，则无需指定model_path。

使用BM25s构建索引

python -m flashrag.retriever.index_builder \
    --retrieval_method bm25 \
    --corpus_path indexes/sample_corpus.jsonl \
    --bm25_backend bm25s \
    --save_dir indexes/ 

使用Pyserini构建索引

python -m flashrag.retriever.index_builder \
    --retrieval_method bm25 \
    --corpus_path indexes/sample_corpus.jsonl \
    --bm25_backend pyserini \
    --save_dir indexes/ 

使用现成的管道

您可以使用我们已经构建的管道类（如管道所示）来实现内部的RAG流程。在这种情况下，您只需配置配置文件并加载相应的管道。

首先，加载整个流程的配置，记录RAG流程中所需的各种超参数。您可以将yaml文件作为参数输入，也可以直接作为变量输入。

请注意，变量作为输入的优先级高于文件。

from flashrag.config import Config

# hybrid load configs
config_dict = {'data_dir': 'dataset/'}
my_config = Config(
    config_file_path = 'my_config.yaml',
    config_dict = config_dict
)

我们提供了全面的配置设置指南，您可以参阅我们的配置指南。 您还可以参考我们提供的yaml文件来设置自己的参数。

接下来，加载相应的数据集并初始化管道。管道中的组件将自动加载。

from flashrag.utils import get_dataset
from flashrag.pipeline import SequentialPipeline
from flashrag.prompt import PromptTemplate
from flashrag.config import Config

config_dict = {'data_dir': 'dataset/'}
my_config = Config(
    config_file_path = 'my_config.yaml',
    config_dict = config_dict
)
all_split = get_dataset(my_config)
test_data = all_split['test']

pipeline = SequentialPipeline(my_config)

您可以使用PromptTemplete指定自己的输入提示：

prompt_templete = PromptTemplate(
    config, 
    system_prompt = "根据给定的文档回答问题。只给我答案，不输出其他任何词。\n以下是给定的文档。\n\n{reference}",
    user_prompt = "问题: {question}\n答案:"
)
pipeline = SequentialPipeline(
    my_config,
    prompt_template = prompt_templete
)

最后，执行pipeline.run以获得最终结果。

output_dataset = pipeline.run(test_data, do_eval=True)

output_dataset包含输入数据集中每个项目的中间结果和指标分数。 同时，包含中间结果和整体评估分数的数据集也将保存为文件（如果指定了save_intermediate_data和save_metric_score）。

构建您自己的管道！

有时您可能需要实现更复杂的RAG流程，您可以构建自己的管道来实现它。 您只需继承BasicPipeline，初始化所需的组件，并完成run函数。

from flashrag.pipeline import BasicPipeline
from flashrag.utils import get_retriever, get_generator

class ToyPipeline(BasicPipeline):
    def __init__(self, config, prompt_templete=None):
        # 加载自己的组件
        pass

    def run(self, dataset, do_eval=True):
        # 完成自己的流程逻辑

        # 使用`.`获取数据集中的属性
        input_query = dataset.question
        ...
        # 使用`update_output`保存中间数据
        dataset.update_output("pred",pred_answer_list)
        dataset = self.evaluate(dataset, do_eval=do_eval)
        return dataset

请首先从我们的文档中了解您需要使用的组件的输入和输出形式。

仅使用组件

如果您已经有自己的代码，只想使用我们的组件嵌入原始代码，您可以参考组件基本介绍获取每个组件的输入和输出格式。

:gear: 组件

在FlashRAG中，我们构建了一系列常见的RAG组件，包括检索器、生成器、精炼器等。基于这些组件，我们组装了多个管道来实现RAG工作流，同时也提供了灵活性，允许用户自定义组合这些组件以创建自己的管道。

RAG组件

类型	模块	描述
判断器	SKR 判断器	使用 SKR 方法判断是否检索
检索器	稠密检索器	使用faiss进行搜索的双编码器模型，如dpr、bge、e5
BM25 检索器	基于Lucene的稀疏检索方法
基于双向编码器的重排序器	使用双向编码器计算匹配分数
基于交叉编码器的重排序器	使用交叉编码器计算匹配分数
精炼器	抽取式精炼器	通过抽取重要上下文来精炼输入
摘要式精炼器	通过seq2seq模型精炼输入
LLMLingua 精炼器	LLMLingua系列提示压缩器
选择性上下文精炼器	选择性上下文提示压缩器
知识图谱精炼器	使用 Trace 方法构建知识图谱
生成器	编码器-解码器生成器	支持 解码器融合 (FiD) 的编码器-解码器模型
仅解码器生成器	原生transformers实现
FastChat 生成器	使用 FastChat 加速
vllm 生成器	使用 vllm 加速

管道

我们参考了一篇检索增强生成的综述，根据推理路径将RAG方法分为四类：

• 顺序：顺序执行RAG过程，如Query-(预检索)-检索器-(后检索)-生成器
• 条件：针对不同类型的输入查询实现不同的路径
• 分支：并行执行多个路径，合并每个路径的响应
• 循环：迭代执行检索和生成

在每个类别中，我们实现了相应的常见管道。一些管道有对应的工作论文。

类别	模块	描述
顺序	Sequential Pipeline	线性执行查询，支持精炼器和重排序器
条件	Conditional Pipeline	使用判断器模块，根据不同的查询类型执行不同的路径
分支	REPLUG Pipeline	通过在多个生成路径中整合概率来生成答案
SuRe Pipeline	根据每个文档对生成的结果进行排序和合并
循环	Iterative Pipeline	交替检索和生成
Self-Ask Pipeline	使用 self-ask 将复杂问题分解为子问题
Self-RAG Pipeline	自适应检索、反思和