本地实现RAG知识库其实很简单（实践）

AGI老王
51 阅读6分钟

本地实现RAG知识库其实很简单（实践）

在很多大模型平台上其实都已经内置了RAG知识库的功能，通常都是上传文档之后，做好分段切分，然后在调用大模型的时候就会关联上知识库进行输出。本篇就尝试本地开发一个RAG知识库，来学习一下具体的实现逻辑。【AI大模型教程】

需求分析

通常我们与大模型交互时，大模型只能回复它所学习到的知识。 有的平台支持联网查询，可以补充部分信息。  不过如果我们希望它回答在某个领域下更垂类的问题时，效果可能就不尽人意。这个时候，就需要给大模型一些参考资料，让它根据资料进行“开卷作答”。

但是这就涉及到一个问题，我们当然可以将所需的内容比如一篇字数比较少的文章扔给大模型，并让它根据文章内容回复。但是，如果当参考资料是一本书的时候，我们不可能将所有几万字甚至几十上百万字的内容都直接在交互对话的过程中扔给大模型，根本不现实。

那么就有两个主要的思路：

  • 微调：根据资料内容制作训练集，在原模型的基础上进行在训练，得到一个更垂类的模型
  • RAG：另一种方法就是本篇文章中要讲的RAG技术。

RAG

RAG（Retrieval-augmented Generation，检索增强生成）是一种结合了语言模型的信息检索技术。从链路上来说，该技术并不改变大模型本身，而是在用户与大模型交互的过程中，增加一个步骤：从知识库检索资料的过程。

具体如下：

可以看到，相当于是在用户的输入到达大模型之前，先送给RAG这边，根据用户的输入先对知识库进行信息检索，然后拿到相关的部分，然后构造prompt，合并到用户的输入中，再送给大模型。prompt示例：

你是一位知识助手，请根据用户的问题和下列的片段生成准确的回答。用户问题：{query}
相关片段：
{separator.join(chunks)}
请基于上述内容作答，不要编造信息

这样知识库和大模型的耦合度更低，是一种可插拔的实现方式，通过更换知识库的内容就可以实现大模型知识的更替，成本更低。

技术实现

这里将通过代码段的方式对实现步骤进行讲解

制作向量数据库

首先，我们要对知识库做处理，将其向量化来可被RAG使用

from typing import List
from sentence_transformers import SentenceTransformer
from sentence_transformers import CrossEncoder
import chromadb
# step1. 知识库分片
def split_info_chunks(doc_file: str) -> List[str]:
"""
:param doc_file: 知识库的文件
:return:
"""
with open(doc_file, 'r', encoding="utf8") as f:
content = f.read()
# 按段落进行分块
return [chunk for chunk in content.split("\n\n")]
# 传入你的知识库文件
chunks = split_info_chunks("doc.txt")
# step2. 向量化 使用embedding模型进行向量化
embedding_model = SentenceTransformer("shibing624/text2vec-base-chinese")
def embed_chunk(chunk: str) -> List[float]:
embedding = embedding_model.encode(chunk)
return embedding.tolist()
embeddings = [embed_chunk(chunk) for chunk in chunks]
# step3. 存入将向量化后的数据存入向量数据库
# 内存型向量数据库
chromadb_client = chromadb.EphemeralClient()
# 创建一个数据库表
chromadb_collection = chromadb_client.get_or_create_collection(name="default")
def save_embedding(chunks: List[str], embeddings: List[List[float]]) -> None:
# 定义id chromadb需要定义好id
ids = [str(i) for i in range(len(chunks))]
chromadb_collection.add(
documents=chunks,
embeddings=embeddings,
ids=ids
)
save_embedding(chunks, embeddings)

特殊说明

  1. sentence_transformer

    库：用于加载embedding和cross-encoder模型，详细使用可以参考：

https://sbert.net/docs/quickstart.html

  1. chromadb

    库：一个流行的向量数据库

  2. 在本demo中，主要是对中文做处理，所以使用的是shibing624/text2vec-base-chinese该模型；关于模型的选择可自行上hugging face找合适的embedding模型

至此，在用户提问前的服务对于知识库的处理便做好了，后续将会对做好的向量数据库进行查阅。

知识检索召回

本质上，这里实现的就是基于用户输入来对向量数据库检索的过程，发生在用户提问后，对用户输入的处理

# 召回，即选出最匹配的若干个chunk
def retrieve(query: str, top_k: int) -> List[str]:
"""
:param query: 用户输入的问题
:param top_k: 最匹配的chunk知识块的数量
:return:
"""
query_embedding = embed_chunk(query)
results = chromadb_collection.query(
query_embeddings=[query_embedding],
n_results=top_k
)
return results['documents'][0]
# 由于召回功能不一定能把最优的知识块排序到最前面
# 所以需要重排模型 CrossEncoder，进行检索内容的进一步优化
def rerank(query: str, retrieved_chunks: List[str], top_k: int) -> List[str]:
"""
:param query: 用户输入的问题
:param retrieved_chunks: 召回的chunks列表
:param top_k: 重新去最匹配的若干
:return:
"""
cross_encoder = CrossEncoder('cross-encoder/mmarco-mMiniLMv2-L12-H384-v1')
# 将每个召回结果和问题进行关联
pairs = [(query, chunk) for chunk in retrieved_chunks]
# 重新计算找回结果和问题的匹配度，计算得分
scores = cross_encoder.predict(pairs)
chunk_with_score_list = [(chunk, score) for chunk, score in zip(retrieved_chunks, scores)]
chunk_with_score_list.sort(key=lambda pair: pair[1], reverse=True)
# print(chunk_with_score_list)
return [chunk for chunk, _ in chunk_with_score_list][:top_k]

至此，数据库检索以及优化部分也完成了。

demo演示

接下来就演示一下实际的效果，这里的知识库doc.txt我选用的是鲁迅先生的《狂人日记》来进行RAG演示

# 假设用户的提问
query = "妹妹发生什么事了"
# 开始进行知识库召回检索和优化
# 优先检索最匹配的5个知识块
retrieved_chunks = retrieve(query, 5)
# 对知识块的关联性进行重新排序，返回关联度最高的3个知识块
reranked_chunks = rerank(query, retrieved_chunks, 3)
# 将用户提问和返回的知识块拼接成一个prompt
separator = "\n\n"
prompt = f"""你是一位知识助手，请根据用户的问题和下列的片段生成准确的回答。
用户问题：{query}
相关片段：
{separator.join(reranked_chunks)}
请基于上述内容作答，不要编造信息
"""
print(prompt)
#后续就是将该prompt发送给大模型，等待大模型的返回，这里就不再实现了

运行后，最终prompt如上，相关片段就是RAG的结果，本篇主要是应用化实现，实际产品中仍需进一步优化。

avatar
文章
阅读
粉丝