源码角度看MemoRAG框架是怎么实现的？实现逻辑及优缺点分析

Python_金钱豹

于 2024-10-10 19:51:08 发布

阅读量544

点赞数 12

文章标签：人工智能机器学习算法语言模型自然语言处理 python 开发语言

本文链接：https://blog.csdn.net/Python_cocola/article/details/142830742

版权

《MemoRAG: Moving towards Next-Gen RAG Via Memory-Inspired Knowledge Discovery》，https://arxiv.org/pdf/2409.05591，这篇技术报告介绍了一个名为MemoRAG的新型检索增强生成系统。

其思想在于，通过长期记忆启发的知识发现来解决传统RAG系统在处理长文本和复杂任务时的局限性，底层逻辑是：从全局记忆中生成精确的线索，将原始输入与答案连接起来，并从复杂数据中解锁隐藏的洞察，这个思路其实不太新。

一句话说，就是，在接收用户查询后，MemoRAG先基于记忆生成模糊答案和答案线索，然后通过检索器补充答案细节，最终生成完整答案。

今天来看具体看看MemoRAG的一些具体的逻辑，看看这个是否真的是那么回事儿。

先说一个具体的结论：

这个是专门用来做单个长文本问答的。

给定一个长文本，先对他做编码，做成KQV，然后做存储。也就是里面的memory.bin。来一个query后，为了可以跟整个文档做交互，所以模型进行query编码时，先加载之前长文本的kqv权重，然后再进行编码。也就相当于拼接了之前的记忆。接着加上记忆之后，根据下游prompt，做rewrite，线索生成任务。

此外，为了后面做召回检索，还是用的传统RAG做的切分，embedding，并建立索引，放在faiss里面，也就是index.bin，后续流程就是上面所说的问答后续流程了。

因此，我们可以做个总结：

在优点方面：

这个场景，针对的是长文本问答，比如kimi输入一个文档，做各种问题。

有了记忆模型之后，相当于看过了整个文档。所以效果会更好。所以，核心点在于，每次上传文档后，需要先把文档读一遍，然后建memory索引。

当多个用户看一篇文档时，可以做cache加速，不用重复建memory索引。

其论文中也说了，主要擅长做摘要和qa任务。

在缺点方面：

来回建，每个文档都得事先存一个，10000个，就得存10000个memory.bin，这个空间占据很大，例如，上图所示，21w字，memory索引要11gb。

此外，文档太长，也搞不定，超出了最大长文长度，效果也很难讲，并且会很慢。

此外，其对长文本模型能力的要求是很高的。

所以，综合来看，这个工作更多的是长文本的演示demo，是个trick的工作，并做不了大规模的落地。

不过，针对长文本做历史kqv-存储是发展趋势，对长文本进行压缩，query与全局文档的语义交互这些，都很有趣。‍

供大家一起参考并思考，持续前进，保持自律，踏实成长。

一、Memo-RAG索引阶段解读

这段代码块展示了如何使用MemoRAG（Memory Reasoning about Gist）流水线中的memorize函数，通过记忆和推理来改善文本检索和理解的能力。

加载文本文件：代码首先加载一个文本文件的内容，这里的例子中使用的文件是harry_potter.txt。

处理文本：通过memorize函数处理加载的文本。这个过程可能包括文本的解析、编码和存储，以便模型可以更好地理解和检索信息。

存储结果：处理后的文本被存储在指定的目录中。在这个目录中，会创建三个关键文件：

memory.bin：这个文件存储了记忆模型的键值（KV）缓存，它允许快速检索之前处理过的信息。这可以看作是一种快速访问存储在模型中的知识的机制。
index.bin：这个文件包含了文本语料库的密集嵌入（dense embeddings），这有助于高效地检索相关段落。嵌入是一种将文本转换为数学表示（通常是向量）的方法，以便机器学习模型可以更好地处理和理解文本。
chunks.json：这个文件包含了从输入上下文中提取的段落或块（chunks），这些在检索过程中使用。这些块是文本的一部分，它们被用来在需要时快速检索信息。

保存预处理数据：如果设置了save_dir参数，该方法会将预处理的数据（即记忆、嵌入和块）保存到磁盘。这样做可以在未来的操作上大大提高效率，因为从磁盘加载缓存的数据比从头开始重新处理和编码上下文要快得多。
缓存机制：这种缓存机制在处理大型文本或频繁访问的数据集时特别有用，因为它减少了重复处理相同数据的需要，从而节省了时间和计算资源。

其中：

1、memory.bin的获取，直接使用的一个长文本模型，获取模型最后的kv进行存储。

2、index.bin的获取，采用的是传统的embedding方案，使用faiss进行存储，并对文档进行切片，并对文档进行doc-embedding。

这个可以从https://github.com/qhjqhj00/MemoRAG/blob/main/memorag/retrieval.py中找到对应的实现细节。

二、Memo-RAG推理阶段解读

这里，涉及到对第一步构建好的memory.bin进行加载。

可以从https://github.com/qhjqhj00/MemoRAG/blob/main/memorag/memorag.py中找到具体的实现逻辑：

1、总体实现逻辑

def _handle_rag(self, query: str, prompt_template: str, max_new_tokens: int, use_memory_answer: bool):           ## 1、召回线索文本           text_spans = self.mem_model.recall(query)           ## 2、生成线索问题           surrogate_queries = self.mem_model.rewrite(query)           ## 3、生成召回问题（线索文本+线索问题+草稿答案并做长度过滤，<3去掉）           retrieval_query, potential_answer = self._prepare_retrieval_query(query, text_spans, surrogate_queries, use_memory_answer)          ## 4、召回相关文档           retrieval_results = self._retrieve(retrieval_query)           if potential_answer:               retrieval_results.append(f"The answer might be {potential_answer}.")           ## 5、文档+草稿答案拼接prompt           knowledge = "\n\n".join(retrieval_results)           ## 6、生成答案           return self._generate_response("qa_gen", query, knowledge, prompt_template, max_new_tokens)

可以先看一些特定的prompt，这个是核心点之一

对应prompt:https://github.com/qhjqhj00/MemoRAG/blob/main/memorag/prompt.py

其中，维护了一个prompt的字典，来存储不同的任务，比如线索生成、片段生成、答案生成等。

2、具体实现步骤

step1：召回线索文本

text_spans = self.mem_model.recall(query)

执行原理：

`def recall(           self,           query, max_new_tokens=128) -> str:           return self.generate(self.prompts["span"], query, max_new_tokens=max_new_tokens)[0]`

调用prompt:

"span": """   你会得到一个与文章相关的问题。为了有效地回答这个问题，你需要回想文章中的具体细节。你的任务是识别并提取文章中与问题相关的一个或多个具体线索文本。      ### 问题：{question}   ### 指示：   1. 你对文章有一个大致的理解。你的任务是生成一个或多个具体的线索，帮助查找文章中的支持证据。   2. 线索应以文本片段的形式呈现，这些片段将有助于回答问题。   3. 只输出线索。如果有多个线索，请用换行符分隔。   4. 请用中文回答。""",

step2：生成线索问题

 `surrogate_queries = self.mem_model.rewrite(query)`

执行原理：

def rewrite(           self,           query, max_new_tokens=128) -> str:           return self.generate(self.prompts["sur"], query, max_new_tokens=max_new_tokens)[0]

调用prompt:

"sur": """   你会得到一个与文章相关的问题。为了有效地回答这个问题，你需要回想文章中的具体细节。你的任务是生成精确的线索问题，帮助找到文章中必要的信息。         ### 问题：{question}   ### 指示：   1. 你对文章有一个大致的理解。你的任务是生成一个或多个具体的线索，帮助查找文章中的支持证据。   2. 线索应以精确的替代问题形式呈现，澄清原问题。   3. 只输出线索。如果有多个线索，请用换行符分隔。   4. 请用中文回答。""",

step3：生成召回问题（线索文本+线索问题+草稿答案并做长度过滤，<3去掉）

  `retrieval_query, potential_answer = self._prepare_retrieval_query(query, text_spans, surrogate_queries, use_memory_answer)`

执行原理：

step4：召回相关文档

        `retrieval_results = self._retrieve(retrieval_query)           if potential_answer:               retrieval_results.append(f"The answer might be {potential_answer}.")`

执行原理：

def answer(           self,           query, max_new_tokens=128) -> str:           return self.generate(self.prompts["qa"], query, max_new_tokens=max_new_tokens)[0]

调用prompt:

"qa": """   你会得到一个与文章相关的问题。你的任务是直接回答这个问题。      ### 问题：{question}   ### 指示：   基于文章的内容，直接回答问题。不要包含除答案之外的任何额外内容。""",

step5：文档+草稿答案拼接prompt

knowledge = "\n\n".join(retrieval_results)

step6：生成答案

 `return self._generate_response("qa_gen", query, knowledge, prompt_template, max_new_tokens)`

总结

本文主要介绍了MemoRAG: Moving towards Next-Gen RAG Via Memory-Inspired Knowledge Discovery》，https://arxiv.org/pdf/2409.05591的具体代码，结论就是，这个跟graphrag一样，是个特定的trick，用空间来换性能，所以会受到很大的落地挑战。

所以，无论是graphrag，还是memorag，都不用吹，没什么太落地的东西，trick思维居多，rag落地，还是要老老实实地走传统那一套，做知识库等的优化。

如何学习大模型 AI ？

由于新岗位的生产效率，要优于被取代岗位的生产效率，所以实际上整个社会的生产效率是提升的。

但是具体到个人，只能说是：

“最先掌握AI的人，将会比较晚掌握AI的人有竞争优势”。

这句话，放在计算机、互联网、移动互联网的开局时期，都是一样的道理。

我在一线互联网企业工作十余年里，指导过不少同行后辈。帮助很多人得到了学习和成长。

我意识到有很多经验和知识值得分享给大家，也可以通过我们的能力和经验解答大家在人工智能学习中的很多困惑，所以在工作繁忙的情况下还是坚持各种整理和分享。但苦于知识传播途径有限，很多互联网行业朋友无法获得正确的资料得到学习提升，故此将并将重要的AI大模型资料包括AI大模型入门学习思维导图、精品AI大模型学习书籍手册、视频教程、实战学习等录播视频免费分享出来。

在这里插入图片描述