HippoRAG

<h1 align="center"> <img src="https://yellow-cdn.veclightyear.com/835a84d5/5c7dc648-aa62-4846-b7b1-3e644ef6a0c7.png" width="30" height="30"> HippoRAG <br> 神经生物学启发的长期记忆 <br> 用于大型语言模型</h1>

HippoRAG 是一个受人类长期记忆神经生物学启发的新型检索增强生成 (RAG) 框架，使大型语言模型能够持续整合外部文档中的知识。我们的实验表明，HippoRAG 可以为 RAG 系统提供通常需要昂贵且高延迟的迭代 LLM 流程才能实现的功能，而计算成本仅为其一小部分。

欲了解更多详情，请查看我们的<a href="https://arxiv.org/abs/2405.14831">论文</a>！

HippoRAG

设置环境

创建 conda 环境并安装依赖：

conda create -n hipporag python=3.9
conda activate hipporag
pip install -r requirements.txt

GPUS=0,1,2,3 #替换为您自己的可用 GPU 设备

将 conda 环境添加到 PATH，如下所示，其中 /path/HippoRAG 是 HippoRAG 的根目录，/path/HippoRAG/hipporag 是 conda 环境的路径。考虑将此添加到您的 ~/.bashrc 中

export PATH=$PATH:/path/HippoRAG/hipporag/bin

设置 LLM API 密钥：TOGETHER_API_KEY 是可选的，当您想使用他们的开源模型（如 Llama-3）时设置它。

export OPENAI_API_KEY='在此添加您自己的 OpenAI API 密钥。'
export TOGETHER_API_KEY='在此添加您自己的 TogetherAI API 密钥。'

要使用 ColBERTv2，下载预训练的检查点并将其放在 exp/colbertv2.0 下。

cd exp
wget https://downloads.cs.stanford.edu/nlp/data/colbert/colbertv2/colbertv2.0.tar.gz
tar -xvzf colbertv2.0.tar.gz
cd .. # 返回根目录

使用 HippoRAG

注意，以下命令的工作目录为 HippoRAG 根目录。

设置您的数据

要设置您自己的检索语料库，请按照 data/sample_corpus.json 中显示的格式和命名约定（您的数据集名称后应跟 _corpus.json）。如果要运行预定义问题的实验，请根据查询文件 data/sample.json 组织您的查询语料库，务必同样遵循我们的命名约定。

语料库和可选查询 JSON 文件应具有以下格式：

检索语料库 JSON

[
  {
    "title": "第一段落标题",
    "text": "第一段落文本",
    "idx": 0
  },
  {
    "title": "第二段落标题",
    "text": "第二段落文本",
    "idx": 1
  }
]

（可选）查询 JSON


[
  {
    "id": "sample/question_1.json",
    "question": "问题",
    "answer": [
      "答案"
    ],
    "answerable": true,
    "paragraphs": [
      {
        "title": "{第一个支持段落标题}",
        "text": "{第一个支持段落文本}",
        "is_supporting": true,
        "idx": 0
      },
      {
        "title": "{第二个支持段落标题}",
        "text": "{第二个支持段落文本}",
        "is_supporting": true,
        "idx": 1
      }
    ]
  }
]

（可选）分段语料库

在准备语料库数据时，您可能需要对每个段落进行分段，因为较长的段落可能对 OpenIE 过程来说过于复杂。如果需要，请查看 src/data_process/util.py 以对整个语料库进行分段。

与 LangChain 集成

此代码库通过 LangChain 调用 LLM，这使得 HippoRAG 更容易调用不同的在线 LLM API 或离线 LLM 部署。查看 src/langchain_util.py 以了解我们如何为实验设置 OpenAI 和 TogetherAI。您还可以使用 LangChain 设置自己的 LLM 选择，例如，Ollama 支持本地的 Llama、Gemma 和 Mistral 模型。

索引

一旦创建了语料库，将其添加到 data 目录下。现在我们准备使用以下命令开始索引。

我们将使用论文中定义的最佳超参数，并假设您的数据集名称为 sample。

对于以下命令，您不必同时使用 ColBERTv2 和 Contriever 进行索引。根据您的偏好选择其中一个。

使用 ColBERTv2 为同义边索引

DATA=sample
LLM=gpt-3.5-turbo-1106
SYNONYM_THRESH=0.8
GPUS=0,1,2,3
LLM_API=openai # LLM API 提供商，例如 'openai'、'together'，见 'src/langchain_util.py'

bash src/setup_hipporag_colbert.sh $DATA $LLM $GPUS $SYNONYM_THRESH $LLM_API

使用 HuggingFace 检索编码器为同义边索引（即 Contriever）

DATA=sample
HF_RETRIEVER=facebook/contriever
LLM=gpt-3.5-turbo-1106
SYNONYM_THRESH=0.8
GPUS=0,1,2,3
LLM_API=openai # LLM API 提供商，例如 'openai'、'together'，见 'src/langchain_util.py'

bash src/setup_hipporag.sh $DATA $HF_RETRIEVER $LLM $GPUS $SYNONYM_THRESH $LLM_API

检索

索引后，HippoRAG 已准备好协助在线检索。以下我们提供了两种使用 HippoRAG 的策略：

对一组预定义查询进行检索。
直接与我们的 API 集成以接收用户的查询。

预定义查询

要对上述格式的特定预定义查询集进行检索，我们根据您想使用的检索主干运行以下命令之一：

ColBERTv2

RETRIEVER=colbertv2

python3 src/ircot_hipporag.py --dataset $DATA --retriever $RETRIEVER --llm $LLM_API --llm_model $LLM --max_steps 1 --doc_ensemble f --top_k 10  --sim_threshold $SYNONYM_THRESH --damping 0.5

Huggingface 模型（即 Contriever）

RETRIEVER=$HF_RETRIEVER

python3 src/ircot_hipporag.py --dataset $DATA --retriever $RETRIEVER --llm $LLM_API --llm_model $LLM --max_steps 1 --doc_ensemble f --top_k 10 --sim_threshold $SYNONYM_THRESH --damping 0.5

注意： 在此设置中，您可以将 HippoRAG 与 IRCoT 结合以获得互补的改进。要运行此操作，只需将上面的 --max_steps 参数更改为所需的最大 LLM 推理步骤数。此外，请确保在 data/ircot_prompt/ 下创建一个以您数据集名称命名的目录，并添加一个名为 gold_with_3_distractors_context_cot_qa_codex.txt 的文件，其中包含适用于您数据集的 IRCoT 提示。查看其他数据集的 IRCoT 提示以获取格式和内容灵感。

HippoRAG 集成

我们在 src/test_hipporag.py 中提供了一个示例，以指导希望将 HippoRAG API 直接集成到他们代码库中的用户。


from src.langchain_util import LangChainModel
from src.qa.qa_reader import qa_read

import argparse

from src.hipporag import HippoRAG

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--dataset', type=str, required=True, default='sample')
    parser.add_argument('--extraction_model', type=str, default='gpt-3.5-turbo-1106')
    parser.add_argument('--retrieval_model', type=str, required=True, help='例如 "facebook/contriever", "colbertv2"')
    parser.add_argument('--doc_ensemble', action='store_true')
    parser.add_argument('--dpr_only', action='store_true')
    args = parser.parse_args()

    assert not (args.doc_ensemble and args.dpr_only)
    hipporag = HippoRAG(args.dataset, 'openai', args.extraction_model, args.retrieval_model, doc_ensemble=args.doc_ensemble, dpr_only=args.dpr_only,
                        qa_model=LangChainModel('openai', 'gpt-3.5-turbo'))

    queries = ["哪位斯坦福大学教授在研究阿尔茨海默症"]
    # qa_few_shot_samples = [{'document': '', 'question': '', 'thought': '', 'answer': ''}]
    # 准备一个用于少样本 QA 的列表，其中每个元素是一个包含 'document'、'question'、'thought'、'answer' 键的字典（'document' 和 'thought' 是可选的）
    qa_few_shot_samples = None
对于每个查询：
    排名、分数、日志 = hipporag.给文档排名(查询, 前k个=10)
    检索到的段落 = [hipporag.通过索引获取段落(排名) for 排名 in 排名列表]

    回复 = qa_阅读(查询, 检索到的段落, qa少量样本, hipporag.qa模型)
    打印(排名)
    打印(分数)
    打印(回复)

要初始化HippoRAG类的实例，只需选择一个**LLM**和一个**检索编码器模型**，这些模型你之前已用于为检索数据集建立索引。

### 自定义数据集演示

要在自定义数据集上运行HippoRAG，请按照上述步骤进行索引和检索，同时还要添加数据处理和评估。以下是[BEIR](https://arxiv.org/abs/2104.08663)数据集的示例，请参见`src/data_process`和`src/demo`。

## 论文复现

在本节中，你将找到复现[我们论文](https://arxiv.org/abs/2405.14831)中所示结果所需的所有代码。

### 数据

我们提供了复现我们实验所需的所有必要数据。

为了在复现过程中节省成本和时间，我们还包括了通过GPT-3.5 Turbo (1106)、两个Llama-3模型和REBEL在所有三个子集和超参数调优数据集上生成的知识图谱。我们还包括了通过GPT-3.5 Turbo (1106)在所有数据集上获得的NER结果。

### 基线

更多详情请查看`src/baselines/README.md`。

### 运行HippoRAG

使用我们的HippoRAG框架需要两个步骤：索引和检索。

#### 索引

要为我们的主要实验和消融实验运行索引，请运行以下bash脚本。如果此步骤未成功，检索将失败。

```shell
bash src/setup_hipporag_main_exps.sh $GPUS

HippoRAG检索

运行索引后，运行以下bash脚本以使用HippoRAG和Contriever及ColBERTv2测试单步和多步检索。

仅HippoRAG

bash src/run_hipporag_main_exps.sh

HippoRAG与IRCoT

bash src/run_hipporag_ircot_main_exps.sh

消融实验

要运行所有消融实验，请运行以下bash脚本：

bash src/setup_hipporag_ablations.sh $GPUS
bash src/run_hipporag_ablations.sh

超参数调优

要复现我们的超参数调优，我们必须首先通过运行以下脚本在MuSiQue训练子集上运行索引：

bash src/setup_hipporag_hyperparameter_tune.sh $GPUS

索引完成后，运行以下脚本并记录测试的每个超参数组合的性能。

bash src/run_hipporag_hyperparameter_tune.sh

问答

更多详情请查看src/qa/README.md。运行QA只能在运行基线和HippoRAG的检索之后进行，因为它使用检索的输出。

路径查找多跳QA案例研究

要运行我们论文中展示的案例研究示例（我们也包括在数据目录中），请运行以下脚本。请注意，要运行这些示例，需要设置你自己的OpenAI API密钥。

索引

bash src/setup_hipporag_case_study.sh $GPUS

检索

bash src/run_hipporag_case_study.sh

运行这些之后，你可以在output/ircot/目录中探索输出。

待办事项

本地部署的LLM
提示灵活性
支持图形数据库，如Neo4j
图的读/写API

联系

有问题或问题？提交一个issue或联系 Bernal Jiménez Gutiérrez, Yiheng Shu, Yu Su, 俄亥俄州立大学

引用

如果你觉得这项工作有用，请考虑引用我们的论文：

@article{gutiérrez2024hipporag,
      title={HippoRAG: Neurobiologically Inspired Long-Term Memory for Large Language Models}, 
      author={Bernal Jiménez Gutiérrez and Yiheng Shu and Yu Gu and Michihiro Yasunaga and Yu Su},
      journal={arXiv preprint arXiv:2405.14831},
      year={2024},
      url={https://arxiv.org/abs/2405.14831}
}