⚡ FlashRAG：高效检索增强生成（RAG）研究的Python工具包

引言

在当今的信息时代，随着数据量的激增，如何有效检索和利用这些数据成为了一个重要课题。FlashRAG，作为一个强大的Python工具包，旨在为检索增强生成（RAG）研究提供一个灵活且高效的框架。本工具包不仅集成了32个经过预处理的基准RAG数据集，还提供了15种最前沿的RAG算法，让研究者能够轻松重现现有的最优成果，或实现自定义的RAG流程。

🌟 特点

FlashRAG具备了以下几个显著特点：

广泛且可定制的框架：该工具包提供了检索器、重排序器、生成器和压缩器等基本组件，允许用户灵活组装复杂的管道。
全面的基准数据集：收录了32个经过预处理的RAG基准数据集，方便测试和验证RAG模型的性能。
预实现的高级RAG算法：包含15种先进的RAG算法，并附带报告结果，用户可根据需求轻松重现不同设置下的结果。
高效的预处理阶段：通过提供多种脚本（如文档检索、索引构建和预检索文档），简化RAG工作流准备过程。
优化执行：库的效率通过使用vLLM、FastChat等工具加速LLM推理，并利用Faiss进行向量索引管理。

graph TD; A[FlashRAG Toolkit] --> B[Retrievers] A --> C[Rerankers] A --> D[Generators] A --> E[Compressors]

🛠️ 安装

要开始使用FlashRAG，您可以通过pip命令轻松安装：

pip install flashrag-dev[core]

或者从GitHub克隆并安装（需要Python 3.9+）：

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

如果您想使用vllm、sentence-transformers或pyserini，可以安装可选依赖项：

pip install flashrag[full]
pip install vllm>=0.4.1
pip install sentence-transformers
pip install pyserini

由于某些系统安装faiss时的不兼容性，您需要使用以下conda命令进行安装：

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

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

🚀 快速启动

玩具示例

为了帮助初学者熟悉我们的工具包，我们提供了一个新手入门指南，您可以直接参考下面的代码实现简单的RAG过程。

演示

我们提供了一个玩具演示来实现简单的RAG过程。您可以随意更换要使用的语料库和模型。下面是如何运行演示的步骤：

cd examples/quick_start
cp ../methods/my_config.yaml .
streamlit run demo_en.py  # 运行英文演示
streamlit run demo_zh.py  # 运行中文演示

管道执行

我们还提供了使用框架执行管道的示例。运行以下代码，使用提供的玩具数据集实现简单的RAG管道。

cd examples/quick_start
python simple_pipeline.py \
    --model_path <Llama-3-8B-instruct-PATH> \
    --retriever_path <E5-PATH>

⚙️ 组件

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

RAG 组件

类型	模块	描述
Judger	SKR Judger	判断是否使用SKR方法进行检索
Retriever	Dense Retriever	使用FAISS进行搜索的双编码器模型
Generator	Encoder-Decoder Generator	支持Fusion-in-Decoder (FiD)

🤖 支持的方法

我们实现了15种方法，统一设置如下：

生成器：LLAMA3-8B-instruct，输入长度2048
检索器：e5-base-v2作为嵌入模型，每个查询检索5个文档
提示：统一的默认提示，模板详见方法细节。

支持的数据集

我们收集并处理了35个在RAG研究中广泛使用的数据集，确保其格式一致，方便使用。这些数据集可在Huggingface datasets获取。

任务	数据集名称	知识来源	训练集大小	开发集大小	测试集大小
QA	NQ	wiki	79,168	8,757	3,610
multi-hop QA	HotpotQA	wiki	90,447	7,405	/

📚 参考文献

如需引用我们的工作，请使用以下BibTex格式：

@article{FlashRAG,
    author={Jiajie Jin and
            Yutao Zhu and
            Xinyu Yang and
            Chenghao Zhang and
            Zhicheng Dou},
    title={FlashRAG: A Modular Toolkit for Efficient Retrieval-Augmented Generation Research},
    journal={CoRR},
    volume={abs/2405.13576},
    year={2024},
    url={https://arxiv.org/abs/2405.13576},
    eprinttype={arXiv},
    eprint={2405.13576}
}