从零实现 RAG

一、前言

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

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

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

二、RAG 的流程

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

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

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

其步骤具体如下:

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

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

三、索引

在开始前,我们需要下载两个模型,分别是 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 相关内容。

相关推荐
算法_小学生28 分钟前
支持向量机(SVM)完整解析:原理 + 推导 + 核方法 + 实战
算法·机器学习·支持向量机
iamlujingtao1 小时前
js多边形算法:获取多边形中心点,且必定在多边形内部
javascript·算法
算法_小学生1 小时前
逻辑回归(Logistic Regression)详解:从原理到实战一站式掌握
算法·机器学习·逻辑回归
DebugKitty2 小时前
C语言14-指针4-二维数组传参、指针数组传参、viod*指针
c语言·开发语言·算法·指针传参·void指针·数组指针传参
qystca2 小时前
MC0241防火墙
算法
行然梦实4 小时前
粒子群优化算法(Particle Swarm Optimization, PSO) 求解二维 Rastrigin 函数最小值问题
算法·机器学习·数学建模
XH华4 小时前
C语言第六章函数递归
c语言·开发语言·算法
斯安5 小时前
LRU(Least Recently Used)原理及算法实现
算法
go54631584655 小时前
基于LSTM和GRU的上海空气质量预测研究
图像处理·人工智能·深度学习·神经网络·算法·gru·lstm
亮亮爱刷题5 小时前
算法提升之数论(矩阵+快速幂)
线性代数·算法·矩阵