从零实现 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 相关内容。

相关推荐
FirstFrost --sy1 小时前
数据结构之二叉树
c语言·数据结构·c++·算法·链表·深度优先·广度优先
森焱森1 小时前
垂起固定翼无人机介绍
c语言·单片机·算法·架构·无人机
搂鱼1145141 小时前
(倍增)洛谷 P1613 跑路/P4155 国旗计划
算法
Yingye Zhu(HPXXZYY)1 小时前
Codeforces 2021 C Those Who Are With Us
数据结构·c++·算法
无聊的小坏坏3 小时前
三种方法详解最长回文子串问题
c++·算法·回文串
长路 ㅤ   3 小时前
Java后端技术博客汇总文档
分布式·算法·技术分享·编程学习·java后端
秋说3 小时前
【PTA数据结构 | C语言版】两枚硬币
c语言·数据结构·算法
qq_513970443 小时前
力扣 hot100 Day37
算法·leetcode
不見星空4 小时前
leetcode 每日一题 1865. 找出和为指定值的下标对
算法·leetcode
我爱Jack4 小时前
时间与空间复杂度详解:算法效率的度量衡
java·开发语言·算法