⚡ FlashRAG:高效检索增强生成(RAG)研究的Python工具包

引言

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

🌟 特点

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

  1. 广泛且可定制的框架:该工具包提供了检索器、重排序器、生成器和压缩器等基本组件,允许用户灵活组装复杂的管道。
  2. 全面的基准数据集:收录了32个经过预处理的RAG基准数据集,方便测试和验证RAG模型的性能。
  3. 预实现的高级RAG算法:包含15种先进的RAG算法,并附带报告结果,用户可根据需求轻松重现不同设置下的结果。
  4. 高效的预处理阶段:通过提供多种脚本(如文档检索、索引构建和预检索文档),简化RAG工作流准备过程。
  5. 优化执行:库的效率通过使用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 组件

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

🤖 支持的方法

我们实现了15种方法,统一设置如下:

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

支持的数据集

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

任务数据集名称知识来源训练集大小开发集大小测试集大小
QANQwiki79,1688,7573,610
multi-hop QAHotpotQAwiki90,4477,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}
}
0 0 投票数
Article Rating
订阅评论
提醒
0 评论
最多投票
最新 最旧
内联反馈
查看所有评论
0
希望看到您的想法,请您发表评论x