从零实现 RAG - 技术栈

一、前言

RAG是大语言模型中最重要的技术之一，在早期 LLM 依赖自身记忆回答，而 LLM 自己数据往往有滞后性。同时由于LLM 的上下文长度限制，由此就出现了 RAG 技术。

虽然如今模型上下文的长度达到了 20 万甚至更多的 token 数，RAG 技术依旧活跃。从此可以看出 RAG 技术的强大。

今天我们将使用最基础的 Python 从 0 实现 RAG 的大致功能。

二、RAG 的流程

RAG 有两个步骤，分别是索引和检索。

索引指的是将文档转换成向量，存储到向量数据库。而检索则是指从用户提问到生成回答的过程。

其检索流程流程如图所示：

其步骤具体如下：

输入 prompt+用户问题
得到 query，从知识库检索
返回相关信息用于上下文增强
prompt+用户问题+增强上下文传递给 LLM
生成响应

今天我们来实现整个过程。

三、索引

在开始前，我们需要下载两个模型，分别是 qwen2.5 和 bge-m3 模型。安装 Ollama 后，执行如下命令：

bash 复制代码

ollama pull qwen2.5
ollama pull bge-m3

下面我们要指定索引要做哪些事情。第一步肯定是读取文档内容，因为 embedding 模型只能处理文本。然后由于文本内容较多，且 embedding 模型有上下文长度限制，所以需要对文档分段。最后就是使用 embedding 模型提取向量，并存储到向量数据库。

具体代码如下：

python 复制代码

import numpy as np
from ollama import embeddings, Message, chat


class Kb:
    def __init__(self, filepath):
        with open(filepath, encoding="utf-8") as f:
            content = f.read()
        self.docs = self.split_content(content)
        self.embeds = self.encode(self.docs)
 
    @staticmethod
    def split_content(content, max_length=256):
        chunks = []
        for i in range(0, len(content), max_length):
            chunks.append(content[i:i + max_length])
        return chunks
 
    @staticmethod
    def encode(texts):
        embeds = []
        for text in texts:
            response = embeddings(model='bge-m3', prompt=text)
            embeds.append(response['embedding'])
        return np.array(embeds)

这里我们实现了一个 Kb 类，里面有 split_content 方法用于将文档拆分成长度不超过max_length 的字符串列表。还有一个 encode 方法，用于将文本转换成向量。

为了方便后续操作，这里我们再添加计算相似度和检索的代码：

python 复制代码

class Kb:
    
    ...
 
    @staticmethod
    def similarity(A, B):
        dot_product = np.dot(A, B)
        norm_A = np.linalg.norm(A)
        norm_B = np.linalg.norm(B)
        cosine_sim = dot_product / (norm_A * norm_B)
        return cosine_sim
 
    def search(self, text):
        max_similarity = 0
        max_similarity_index = 0
        e = self.encode([text])[0]
        for idx, te in enumerate(self.embeds):
            similarity = self.similarity(e, te)
            if similarity > max_similarity:
                max_similarity = similarity
                max_similarity_index = idx
        return self.docs[max_similarity_index]

这里相似度使用余弦相似度，而检索方法则使用线性检索。线性检索的性能较低，实际向量数据库对此有专门的优化，这里为了演示则不考虑性能优化问题。

search 函数会返回与 text 最相关的一个文档，而实际我们可能会需要多个文档。

四、检索

下面我们来根据第二节的步骤来实现一个 RAG 类：

python 复制代码

class Rag:
    def __init__(self, model, kb):
        self.model = model
        self.kb = kb
        self.prompt_template = """
        基于：%s
        回答：%s
        """
 
    def chat(self, text):
        # 1. 输入 prompt+用户问题
        prompt = self.prompt_template % (context, text)
        # 2. 得到 query，从知识库检索
        # 3. 返回相关信息用于上下文增强
        context = self.kb.search(text)
        # 4. prompt+用户问题+增强上下文传递给 LLM
        # 5. 生成响应
        response = chat(self.model, Message(role='system', content=prompt))
        return response['message']
 
 
if __name__ == '__main__':
    kb = Kb('data/kb.txt')
    rag = Rag('qwen2.5', kb)
    while True:
        q = input("Q:")
        r = rag.chat(q)
        print(f"A:{r['content']}")

这里我们使用了一个非常简单的 Prompt：

python 复制代码

"""
基于：%s
回答：%s
"""

在填入模板后，其样式大致如下：

ini 复制代码

基于：阿尔伯特·爱因斯坦（德语：Albert Einstein，1879年3月14日---1955年4月18日），德国犹太裔理论物理学家，拥有瑞士和美国国籍。其创立了现代物理学的两大支柱的相对论及量子力学[36]:274[33]，也是质能等价的发现者[37]。他在科学哲学领域颇具影响力[38][39]。因为"对理论物理的贡献，特别是发现了光电效应的原理"，他荣获1921年度的诺贝尔物理学奖（1922年颁发）。这一发现为量子理论的建立踏出了关键性的一步。[40]

回答：爱因斯坦什么时候获得诺贝尔奖的？

这样大模型就能基于外部文档来回答我们的问题了。

五、总结

本文实现了一个 demo 级别的 RAG 应用，实际的 RAG 每个阶段都要考虑很多事情。比如检索返回结果需要有多个、检索的方式应该采用复杂的更低的算法、文档分割方式需要更加细致等。而这些都不是本文要讨论的内容。本文的目的在于有一个通俗易懂的例子实现 RAG，并解释 RAG 的基本原理。感兴趣的读者可以阅读更多 RAG 相关内容。