引言
在当今的信息时代,随着数据量的激增,如何有效检索和利用这些数据成为了一个重要课题。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}
}