MemoRAG：记忆增强的长文本 RAG 方案

用通过记忆增强的数据接口为通用应用赋能 RAG！

快速开始 | 路线图 | 使用 | 演示 | 数据集 | 常见问题

MemoRAG 是一个建立在高效超长记忆模型基础之上的创新型 RAG 框架。与主要处理具有显式信息需求查询的标准 RAG 不同，MemoRAG 利用其记忆模型来实现对整个数据库的全局理解。通过从记忆中召回特定查询的线索，MemoRAG 增强了证据检索，从而实现更准确和上下文丰富的响应生成。MemoRAG 的论文已被 TheWebConf 2025 接收。

[23/04/25] 我们发布了 MemoRAG 的训练脚本和训练数据集。我们将在不久的将来继续改进本仓库。

[21/09/24] MemoRAG 推出 Lite 模式，仅需几行代码即可处理数百万令牌的内存增强 RAG。更多详情请参考示例笔记本。

[13/09/24] MemoRAG 添加 Meta-Llama-3.1-8B-Instruct 和 Llama3.1-8B-Chinese-Chat 作为记忆模型，见示例。

[10/09/24] 我们发布了 MemoRAG 的技术报告。

[09/09/24] 你可以在 Google Colab 上免费试用 MemoRAG。

[05/09/24] 基于 Qwen2 的记忆模型已在 TommyChien/memorag-qwen2-7b-inst 中提供。

[03/09/24] 基于 Mistral 的记忆模型已在 TommyChien/memorag-mistral-7b-inst 中提供。

[01/09/24] 项目启动！

全局记忆：在单个上下文中处理最多 100 万个令牌，提供跨大规模数据集的全面理解。

可优化且灵活：轻松适应新任务，仅需数小时额外训练即可实现优化性能。

上下文线索：从全局记忆生成精确线索，联系原始输入与答案，从复杂数据中释放隐藏的见解。

高效缓存：将上下文预填充速度提升最多 30 倍，支持缓存分块、索引和编码。

上下文重用：对长上下文进行一次编码，支持重复使用，提升需要反复数据访问的任务效率。

MemoRAG 目前处于积极开发中，资源和原型不断在本仓库中发布。

代码 / 模型 / 数据集发布

支持 OpenAI/Azure 模型

技术报告发布

内存模型训练代码发布

集成任何检索方法

丰富记忆能力

注：MemoRAG 的近期目标是通过工程改进实现轻量级优化，并增强其记忆能力，使其能适应更广泛的应用并支持更长的上下文（例如超过 100 万个令牌）。

🆓 在 Google Colab 上免费试用 MemoRAG

在本笔记本中，我们在 Google Colab 提供的单个 T4 GPU 上运行完整的 MemoRAG 管道（记忆模型 + 检索器 + 生成模型），该 GPU 拥有 15GiB 内存。尽管资源有限，MemoRAG 仍可处理示例书籍的一半内容（约 68K 令牌）并执行其所有功能。

安装

要使用 Memorizer 和 MemoRAG，你需要安装 Python 以及必要的库。可以使用以下命令安装必要的依赖：

pip install torch==2.3.1
conda install -c pytorch -c nvidia faiss-gpu=1.8.0

# clone this repo first
cd MemoRAG
pip install -e .

pip install memorag

如需快速开始，我们在这里提供了一个笔记本来展示 MemoRAG 的所有功能。

MemoRAG Lite 模式 🆕 🆕 🆕

我们引入了 MemoRAG 的 Lite 模式，旨在提供快速且用户友好的 MemoRAG 管道体验。仅需几行代码，你就可以轻松尝试 MemoRAG。虽然我们推荐使用拥有 24GiB 内存的 GPU，但在大多数情况下，拥有 16GiB 内存的 GPU 也能在默认设置下处理该管道。

from memorag import MemoRAGLite
pipe = MemoRAGLite()
context = open("examples/harry_potter.txt").read()
pipe.memorize(context, save_dir="harry_potter", print_stats=True)

query = "What's the book's main theme?"
print(pipe(query))

MemoRAG Lite 使用简便，支持最多数百万令牌的英文或中文上下文。虽然它可能适用于其他语言，但由于默认提示词是英文，性能可能会下降。有关 MemoRAG Lite 的更多详情，请参考示例笔记本。

MemoRAG 的基础使用

MemoRAG 易于使用，可直接用 HuggingFace 模型初始化。通过使用 MemoRAG.memorize() 方法，记忆模型可在长输入上下文上构建全局记忆。根据经验，在默认参数设置下，TommyChien/memorag-qwen2-7b-inst 可处理最多 400K 令牌的上下文，而 TommyChien/memorag-mistral-7b-inst 可管理最多 128K 令牌的上下文。通过增加 beacon_ratio 参数，可以扩展模型处理更长上下文的能力。例如，TommyChien/memorag-qwen2-7b-inst 在 beacon_ratio=16 时可处理最多 100 万个令牌。

from memorag import MemoRAG

# Initialize MemoRAG pipeline
pipe = MemoRAG(
    mem_model_name_or_path="TommyChien/memorag-mistral-7b-inst",
    ret_model_name_or_path="BAAI/bge-m3", 
    gen_model_name_or_path="mistralai/Mistral-7B-Instruct-v0.2", # Optional: if not specify, use memery model as the generator
    cache_dir="path_to_model_cache",  # Optional: specify local model cache directory
    access_token="hugging_face_access_token",  # Optional: Hugging Face access token
    beacon_ratio=4
)

context = open("examples/harry_potter.txt").read()
query = "How many times is the Chamber of Secrets opened in the book?"

# Memorize the context and save to cache
pipe.memorize(context, save_dir="cache/harry_potter/", print_stats=True)

# Generate response using the memorized context
res = pipe(context=context, query=query, task_type="memorag", max_new_tokens=256)
print(f"MemoRAG generated answer: \n{res}")

运行上述代码时，编码的键值 (KV) 缓存、Faiss 索引和分块段落存储在指定的 save_dir 中。之后，如果再次使用相同的上下文，数据可从磁盘快速加载：

pipe.load("cache/harry_potter/", print_stats=True)

通常，加载缓存权重效率很高。例如，使用 TommyChien/memorag-qwen2-7b-inst 作为记忆模型对 200K 令牌上下文进行编码、分块和索引需要大约 35 秒，但从缓存文件加载仅需 1.5 秒。

使用长上下文 LLM 作为记忆模型 🆕 🆕 🆕

最近的 LLM 因上下文窗口扩展而成为有效的记忆模型。MemoRAG 现在支持利用这些长上下文 LLM 作为记忆模型，利用 MInference 优化上下文预填充。我们已测试 Meta-Llama-3.1-8B-Instruct 和 Llama3.1-8B-Chinese-Chat 作为记忆模型，两者原生支持 128K 上下文长度。我们目前正在探索其他合适的 LLM 并优化策略，以进一步增强记忆机制和上下文长度。详细使用说明请参考提供的脚本和笔记本：

from memorag import MemoRAG
model = MemoRAG(
    mem_model_name_or_path="shenzhi-wang/Llama3.1-8B-Chinese-Chat",    # For Chinese
    # mem_model_name_or_path="meta-llama/Meta-Llama-3.1-8B-Instruct",  # For English
    ret_model_name_or_path="BAAI/bge-m3",
    # cache_dir="path_to_model_cache",  # to specify local model cache directory (optional)
    # access_token="hugging_face_access_token"  # to specify local model cache directory (optional)
    )

之后，你可以像往常一样使用 MemoRAG 的功能。

要执行总结任务，请使用以下脚本：

res = pipe(context=context, task_type="summarize", max_new_tokens=512)
print(f"MemoRAG summary of the full book:\n {res}")

使用 API 作为生成器

如果你想使用 API 作为生成器，请参考下面的脚本：

from memorag import Agent, MemoRAG

# API configuration
api_dict = {
    "endpoint": "",
    "api_version": "2024-02-15-preview",
    "api_key": ""
}
model = "gpt-35-turbo-16k"
source = "azure"

# Initialize Agent with the API
agent = Agent(model, source, api_dict)
print(agent.generate("hi!"))  # Test the API

# Initialize MemoRAG pipeline with a customized generator model
pipe = MemoRAG(
    mem_model_name_or_path="TommyChien/memorag-qwen2-7b-inst",
    ret_model_name_or_path="BAAI/bge-m3",
    cache_dir="path_to_model_cache",  # Optional: specify local model cache directory
    customized_gen_model=agent,
)

# Load previously cached context
pipe.load("cache/harry_potter_qwen/", print_stats=True)

# Use the loaded context for question answering
query = "How are the mutual relationships between the main characters?"
context = open("harry_potter.txt").read()

res = pipe(context=context, query=query, task_type="memorag", max_new_tokens=256)
print(f"MemoRAG with GPT-3.5 generated answer: \n{res}")

支持的生成器 API

内置 Agent 对象支持来自 openai 和 deepseek 的模型。以下是初始化这些模型的配置：

# Using deepseek models
model = ""
source = "deepseek"
api_dict = {
    "base_url": "",
    "api_key": ""
}

# Using openai models
model = ""
source = "openai"
api_dict = {
    "api_key": ""
}

记忆模型使用

记忆模型可独立使用来存储、召回和交互上下文。这是一个示例：

from memorag import Memory

# Initialize the Memory model
memo_model = Memory(
    "TommyChien/memorag-qwen2-7b-inst",
    cache_dir="path_to_model_cache",  # Optional: specify local model cache directory
    beacon_ratio=4  # Adjust beacon ratio for handling longer contexts
)

# Load and memorize the context
context = open("harry_potter.txt").read()
memo_model.memorize(context)

# Save the memorized context to disk
memo_model.save("cache/harry_potter/memory.bin")

# Query the model for answers
query = "How are the mutual relationships between the main characters?"

res = memo_model.answer(query)
print("Using memory to answer the query:\n", res)

# Recall text clues for evidence retrieval
res = memo_model.recall(query)
print("Using memory to recall text clues to support evidence retrieval:\n", res)

# Rewrite the query into more specific surrogate queries
res = memo_model.rewrite(query)
print("Using memory to rewrite the input query into more specific surrogate queries:\n", res)

内存增强检索使用

除了独立的记忆模型，MemoRAG 还提供内存增强检索功能。这允许基于从记忆中召回的线索来改进证据检索。

from memorag import MemoRAG

# Initialize MemoRAG pipeline
pipe = MemoRAG(
    mem_model_name_or_path="TommyChien/memorag-qwen2-7b-inst",
    ret_model_name_or_path="BAAI/bge-m3",
    cache_dir="path_to_model_cache",  # Optional: specify local model cache directory
    access_token="hugging_face_access_token"  # Optional: Hugging Face access token
)

# Load and memorize the context
test_txt = open("harry_potter.txt").read()
pipe.memorize(test_txt, save_dir="cache/harry_potter/", print_stats=True)

# Define the query
query = "How are the mutual relationships between the main characters?"

# Recall clues from memory
clues = pipe.mem_model.recall(query).split("\n")
clues = [q for q in clues if len(q.split()) > 3]  # Filter out short or irrelevant clues
print("Clues generated from memory:\n", clues)

# Retrieve relevant passages based on the recalled clues
retrieved_passages = pipe._retrieve(clues)
print("\n======\n".join(retrieved_passages[:3]))

基准评估

以下是记忆模型与三个生成模型相结合的实验结果。

要评估 MemoRAG，请使用以下脚本：

cd examples
bash longbench/eval.sh

我们将很快更新其他评估脚本。

UltraDomain 基准：此仓库。

其他评估数据：此仓库。

许可证

MemoRAG 采用 Apache 2.0 许可证发布。

引用

如果你在研究中使用了 MemoRAG，请引用我们的论文：

@inproceedings{qian2025memorag,
  title     = {MemoRAG: Boosting Long Context Processing with Global Memory-Enhanced Retrieval Augmentation},
  author    = {Hongjin Qian and Zheng Liu and Peitian Zhang and Kelong Mao and Defu Lian and Zhicheng Dou and Tiejun Huang},
  booktitle = {Proceedings of the ACM Web Conference 2025 (TheWebConf 2025)},
  year      = {2025},
  address   = {Sydney, Australia},
  publisher = {ACM},
  url       = {https://arxiv.org/abs/2409.05591},
}