LongRAG

本代码库包含"LongRAG：利用长上下文LLM增强检索增强生成"的代码。<span style="color: red;">我们仍在完善我们的代码库。</span>

简介

在传统的RAG框架中，基本检索单元通常较短。这种设计迫使检索器在大型语料库中搜索以找到"针"单元。相比之下，阅读器只需要从短小的检索单元中提取答案。这种不平衡的重检索器轻阅读器设计可能导致次优性能。我们提出了一个新的框架LongRAG，包含一个"长检索器"和一个"长阅读器"。我们的框架使用4K令牌的检索单元，比之前长30倍。我们的研究为未来将RAG与长上下文LLM结合的路线图提供了洞见。

安装

克隆此仓库并安装所需的包：

git clone https://github.com/TIGER-AI-Lab/LongRAG.git
cd LongRAG
pip install -r requirements.txt

快速开始

请前往"长阅读器"部分并按照说明操作。这将帮助你获得100个示例的最终预测。输出将类似于我们在exp/目录中的示例文件。

语料库准备（可选）

这是一个可选步骤。你可以直接使用我们处理好的语料库。我们已在Hugging Face上发布了NQ和HotpotQA的两个版本的检索语料库。

from datasets import load_dataset
corpus_nq = load_dataset("TIGER-Lab/LongRAG", "nq_corpus")
corpus_hotpotqa = load_dataset("TIGER-Lab/LongRAG", "hotpot_qa_corpus")

如果你仍对我们如何制作语料库感兴趣，可以从这里开始阅读。

维基百科原始数据清理： 我们首先按照标准流程清理维基百科原始数据。我们使用WikiExtractor。这是一个广泛使用的Python脚本，用于从维基百科数据库备份转储中提取和清理文本。请确保使用所需的Python环境。示例脚本如下：

sh scripts/extract_and_clean_wiki_dump.sh

预处理维基百科数据 清理维基百科原始数据后，运行以下脚本以收集更多信息。

sh scripts/process_wiki_page.sh

dir_path：清理后的维基百科转储的目录路径，这是上一步的输出。
output_path_dir：输出目录将包含几个pickle文件，每个文件代表维基百科页面的字典。 degree.pickle：键是维基百科页面标题，值是超链接数量。 abs_adj.pickle：键是维基百科页面标题，值是摘要段落中的链接页面。 full_adj.pickle：键是维基百科页面标题，值是整个页面中的链接页面。 doc_size.pickle：键是维基百科页面标题，值是该页面上的令牌数量。 doc_dict.pickle：键是维基百科页面标题，值是页面的文本。
corpus_title_path：用于过滤NQ数据集的键。在原始DPR论文中，某些维基百科页面（如列表页面和消歧页面）被移除，将维基百科页面总数从500万减少到300万。为了公平比较，我们也选择排除这些页面。（对于HotpotQA，我们没有删除任何页面，所以维基百科页面数量仍然是500万。）你可以从这个链接下载DPR的标题。

我们已经在我们的huggingface仓库中提供了处理过的维基百科。请查看nq_wiki和hotpot_qa_wiki子集以获取更多信息。你可以轻松从这两个数据集中得到这些pickle文件。 检索语料库： 通过将多个相关文档组合在一起，我们可以构建超过4K个标记的长检索单元。这种设计还可以显著减少语料库的大小（语料库中检索单元的数量）。这样，检索器的任务就变得容易得多。此外，长检索单元还将改善信息的完整性，避免歧义或混淆。

sh scripts/group_documents.sh

processed_wiki_dir：上述步骤的输出目录。
mode：abs用于HotpotQA语料库，full用于NQ语料库。
output_dir：输出目录，该目录将包含几个pickle文件，每个文件代表检索语料库的一个字典。最重要的是group_text.pickle，它将语料库ID映射到语料库文本。更多详情，请参考我们在Hugging Face上发布的语料库。

长文本检索器

我们利用开源的密集检索工具包Tevatron进行所有的检索实验。我们使用的基础嵌入模型是bge-large-en-v1.5。我们提供了一个示例脚本；请确保使用你自己的数据集本地路径更新参数。此外，我们的脚本使用4个GPU来编码语料库以节省时间；请根据你自己的使用情况进行更新。

sh scripts/run_retrieve_tevatron.sh

长文本阅读器

我们选择Gemini-1.5-Pro和GPT-4o作为我们的长文本阅读器，因为它们具有处理长上下文输入的强大能力。（我们还计划在未来测试其他能够处理长上下文的LLM。）

阅读器的输入是长文本检索器所有长检索单元的连接。我们已在Huggingface仓库中提供了输入文件。

mkdir -p exp/
sh scripts/run_eval_qa.sh

test_data_name：测试集名称，nq（NQ）或hotpot_qa（HotpotQA）。
test_data_split：对于每个测试集，有三个拆分：full、subset_1000、subset_100。我们建议从subset_100开始，以快速入门或调试，并使用subset_1000获得相对稳定的结果。
output_file_path：输出文件，这里放在exp/目录中。
reader_model：我们使用的长上下文阅读器模型，目前我们的代码支持GPT-4o、GPT-4-Turbo、Gemini-1.5-Pro、Claude-3-Opus。请注意，你需要在代码中更新相关的API密钥和API配置。例如，如果你使用GPT-4系列，需要在utils/gpt_inference.py中配置代码；如果使用Gemini系列，需要在utils/gemini_inference.py中配置代码。我们将在未来继续支持更多模型。

输出文件每行包含一个测试用例。short_ans字段是我们的最终预测。

{
    "query_id": "383", 
    "question": "how many episodes of touching evil are there", 
    "answers": ["16"], 
    "long_ans": "16 episodes.", 
    "short_ans": "16", 
    "is_exact_match": 1, 
    "is_substring_match": 1, 
    "is_retrieval": 1
}

我们在exp/目录中提供了一些示例输出文件。例如，exp/nq_gpt4o_100.json包含运行以下文件的结果：

python eval/eval_qa.py \
  --test_data_name "nq" \
  --test_data_split "subset_100" \
  --output_file_path "./exp/nq_gpt4o_100.json" \
  --reader_model "GPT-4o"

top-1检索准确率为88%，精确匹配率为64%。

许可证

请查看我们在工作中使用的每个子集的许可证。

数据集名称	许可证类型
NQ	Apache License 2.0
HotpotQA	CC BY-SA 4.0 License

引用

如果您发现我们的项目有用，请kindly引用我们的论文：

@article{jiang2024longrag
  title={LongRAG: Enhancing Retrieval-Augmented Generation with Long-context LLMs},
  author={Ziyan Jiang, Xueguang Ma, Wenhu Chen},
  journal={arXiv preprint arXiv:2406.15319},
  year={2024},
  url={https://arxiv.org/abs/2406.15319}
}