论文分享|Arxiv2024‘人大|FlashRAG—快速开发和公平评估RAG算法

BrownSearch

已于 2024-07-19 23:33:00 修改

阅读量1k

点赞数 21

分类专栏： NLP/IR知识分享文章标签：算法 python nlp

于 2024-07-10 10:43:31 首次发布

本文链接：https://blog.csdn.net/weixin_45783724/article/details/140317352

版权

NLP/IR知识分享专栏收录该内容

23 篇文章 6 订阅

订阅专栏

构建RAG系统的一个比较难的问题是，如何快速且公平地对比各种RAG方法？本文介绍一个最近比较火的项目FlashRAG，讲解如何快速实现并公平评估自己的RAG算法。

论文题目：FlashRAG: A Modular Toolkit for Efficient Retrieval-Augmented Generation Research

来源：中国人民大学

开源地址：https://github.com/RUC-NLPIR/FlashRAG

概述

FlashRAG 是一个 Python 工具包，用于复现和开发检索增强生成（RAG）算法。此工具包包括 32 个经过预处理的基准 RAG 数据集和 12 种最先进的 RAG 算法。以下是架构图：

最底层为组件层，包含RAG中的各种常用组件，包含检索器，重排器，压缩器，生成器，打分器等。

第二层为流程层，包含顺序，分支，迭代，循环，条件等流程，这部分主要是合理使用各种组件来实现具体的算法

第三层为数据层，包含用于检索的语料数据和用于评估的各种任务数据。

使用教程

环境安装

执行如下命令，需要python>=3.9。注意，faiss-cpu需要用conda重新安装从而适配特定设备

git clone https://github.com/RUC-NLPIR/FlashRAG.git
cd FlashRAG
pip install -e . 
conda install -c pytorch faiss-cpu==1.8.0

数据准备

我们使用如下的python脚本，利用hf镜像快速下载 FlashRAG_datasets 数据集到本地的 FlashRAG_datasets 目录。

import os
os.environ['HF_ENDPOINT'] = 'https://hf-mirror.com'
from pathlib import Path
import fire

def download(
        repo: str = "RUC-NLPIR/FlashRAG_datasets",
        output: str = "FlashRAG_datasets",
):
    Path(output).mkdir(parents=True, exist_ok=True)
    os.system(
        f'huggingface-cli download --repo-type dataset --resume-download {repo} --local-dir {output} --local-dir-use-symlinks False')

if __name__ == "__main__":
    fire.Fire(download)

下载后我们需要将其中的 FlashRAG_datasets/retrieval-corpus/wiki-18.jsonl.gz 使用gzip命令解压，从而得到jsonl文件： FlashRAG_datasets/retrieval-corpus/wiki-18.jsonl

索引建立

FlashRAG默认使用统一的WikiPedia-2018作为检索语料库，使用统一的 e5-base-v2 作为检索器。我们使用如下命令建立相应的索引库，A100 40G上大概需要运行2小时左右。

CUDA_VISIBLE_DEVICES=0 python3 -m flashrag.retriever.index_builder \
    --retrieval_method e5 \
    --model_path intfloat/e5-base-v2 \
    --corpus_path FlashRAG_datasets/retrieval-corpus/wiki-18.jsonl \
    --save_dir index/ \
    --use_fp16 \
    --max_length 256 \
    --batch_size 512 \
    --pooling_method mean \
    --faiss_type Flat

算法复现

我们先命令行来到 examples/methods 目录，目前修改 my_config.yaml 中以下参数：

model2path.llama2-7B-chat 改为 meta-llama/Meta-Llama-3-8B-Instruct ，从而使用默认的llama3
index_path 改为 …/…/index/e5_Flat.index
data_dir 改为 …/…/FlashRAG_datasets/
corpus_path 改为 …/…/FlashRAG_datasets/retrieval-corpus/wiki-18.jsonl
framework 改为 fschat
generation_params：do_sample 改为false，max_tokens 改为32

接下来就可以复现算法了，我们以naive和replug为例，在NQ上做测试，运行如下的命令。大概需要20min左右跑完一个方法。其他方法参考 run_exp.py 中的方法字典替换 method_name 即可

# replug
python3 run_exp.py \
--method_name=replug \
--dataset_name=nq \
--gpu_id=0 \
--split=test

replug会得到如下结果，EM结果和论文中报告的基本一致。

算法实现

如果想自己实现和测试一个新的算法，只需要继承BasicPipeline，并合理使用库中的各种组件即可，这里我们以一个简单的naive rag作为例子：

class NaivePipeline(BasicPipeline):
    def __init__(self, config):
        """
        inference stage:
            query -> retriever -> generator
        """
        # 初始化检索器和生成器
        super().__init__(config, prompt_template)
        self.retriever = get_retriever(config)
        self.generator = get_generator(config)

    def run(self, dataset, do_eval=True, pred_process_fun=None):
        # 获得输入问题列表
        input_query = dataset.question

        # 检索
        retrieval_results = self.retriever.batch_search(input_query)
        dataset.update_output('retrieval_result', retrieval_results)
        
        # 用提示模板拼接检索结果与问题，得到上下文
        input_prompts = [
            self.prompt_template.get_string(question=q, retrieval_result=r)
            for q, r in zip(dataset.question, dataset.retrieval_result)
        ] 
        dataset.update_output('prompt', input_prompts)
        
        # 利用上下文生成结果
        pred_answer_list = self.generator.generate(input_prompts)
        dataset.update_output("pred",pred_answer_list)
        
        # 评估
        dataset = self.evaluate(dataset, do_eval=do_eval, pred_process_fun=pred_process_fun)

        return dataset

config_dict = {'save_note': 'naive',
                'gpu_id':'0',
                'dataset_name':'nq'} #在此替换数据集

# 创建config
config = Config('my_config.yaml',config_dict)
# 加载数据集
all_split = get_dataset(config)
test_data = all_split['test']
# 创建RAG算法对象
pipeline = NaivePipeline(config)
# 测试
result = pipeline.run(test_data)