Python搭建一个RAG系统(分片/检索/召回/重排序/生成)

发布时间：2025-08-28 15:43:17编辑：Run阅读（45）

RAG是构建知识库的首选方案。检索增强生成（RAG）是指对大语言模型输出进行优化 。

整体流程一共分为两部分：

第一部分：用户提问前的数据准备部分，包括分片和索引两个环节。


第二部分：用户提问后的回答部分，包括召回，重排和生成三个环节。




用户提问前，准备阶段，把相关的资料(文档，数据)做个分片，然后把所有的片段都扔给Embedding模型（嵌入模型，当前EmbeddingModel主要用于将文本转换为数值向量 ），让它给每个片段都产出一个对应的向量，最后把向量存入到向量数据库中，到这里提问前的准备流程就结束了。

用户提问后，用户的问题会给到Embedding模型转为一个向量，然后把这个向量传给向量数据库，让它找到10个与用户问题最相近的片段，找到后再把这个10个片段送给cross-encoder模型，让它做重排，从10个里面再筛选出3个相关程度最高的，然后把这三个片段外加用户问题一起发送给大模型，大模型就可以产出最终答案了。




使用Python实现一个RAG系统。

项目环境搭建：（这里使用conda创建虚拟环境，conda是python的包管理器）

创建虚拟环境，安装python=3.12版本

conda create --name rag python=3.12

激活虚拟环境

conda activate rag

安装python依赖


pip install sentence_transformers chromadb  python-dotenv




每个包的作用

sentence_transformer  加载emberdding和cross-encoder模型

chromadb  一个非常流行的向量数据库

python-dotenv 将Google API Key映射到环境变量中




为了保证测试的准确性，特地准备了一篇文章。（讲诉哆啦A梦，大雄与超级撒亚人一起合力打败黑暗势力拯救地球的故事）


完整的文章：

哆啦A梦与超级赛亚人：时空之战

在一个寻常的午后，大雄依旧坐在书桌前发呆，作业堆得像山，连第一页都没动。哆啦A梦在一旁翻着漫画，时不时叹口气，觉得这孩子还是一如既往的不靠谱。正当他们的生活照常进行时，一道强光突然从天而降，整个房间震动不已。光芒中走出一名金发少年，身披战甲、气势惊人，他就是来自未来的超级赛亚人——特兰克斯。他一出现便说出了惊人的话：未来的地球即将被黑暗势力摧毁，他来此是为了寻求哆啦A梦的帮助。




哆啦A梦与大雄听后大惊，但也从特兰克斯坚定的眼神中读出了不容拒绝的决心。特兰克斯解释说，未来的敌人并非普通反派，而是一个名叫“黑暗赛亚人”的存在，他由邪恶科学家复制了贝吉塔的基因并加以改造，实力超乎想象。这个敌人不仅拥有赛亚人战斗力，还能操纵扭曲的时间能量，几乎无人可敌。特兰克斯已经独自战斗多年，但每一次都以惨败告终。他说：“科技，是我那个时代唯一缺失的武器，而你们，正好拥有它。”




于是，哆啦A梦带着特兰克斯与大雄启动时光机，穿越到了那个即将崩溃的未来世界。眼前的景象令人震撼：城市沦为废墟，大地裂痕纵横，天空中浮动着压抑的黑雾。特兰克斯说，这正是黑暗赛亚人带来的结果，一切生命几乎都被抹杀，只剩他在苦苦支撑。大雄虽感到恐惧，但看到无辜的人类遭殃，内心逐渐燃起斗志。哆啦A梦则冷静地分析局势，决定使用他最强的三样秘密道具来对抗黑暗势力。




三件秘密道具分别是：可以临时赋予超级战力的“复制斗篷”，能暂停时间五秒的“时间停止手表”，以及可在一分钟中完成一年修行的“精神与时光屋便携版”。大雄被推进精神屋内，在其中接受密集的训练，虽然只有几分钟现实时间，他却经历了整整一年的苦修。刚开始他依旧软弱，想放弃、想逃跑，但当他想起静香、父母，还有哆啦A梦那坚定的眼神时，他终于咬牙坚持了下来。出来之后，他的身体与精神都焕然一新，眼神中多了一份成熟与自信。




最终战在黑暗赛亚人的空中要塞前爆发，特兰克斯率先出击，释放全力与敌人正面对决。哆啦A梦则用任意门和道具支援，从各个方向制造混乱，尽量压制敌人的时空能力。但黑暗赛亚人太过强大，仅凭特兰克斯一人根本无法压制，更别说击败。就在特兰克斯即将被击倒之际，大雄披上复制斗篷、冲破恐惧从高空跃下。他的拳头燃烧着金色光焰，目标直指敌人心脏。




时间停止装置在关键时刻启动，世界陷入静止，大雄用这个短短五秒接近了敌人的盲点。他集中全力，一记重拳击穿了黑暗赛亚人的能量核心，引发巨大的能量反冲。黑暗赛亚人尖叫着化为碎光，天空中的黑雾瞬间散去，阳光重新洒落大地。特兰克斯倒在地上，看着眼前这个曾经懦弱的少年，露出了欣慰的笑容。他知道，这一次，是大雄救了世界。




战后，未来世界开始恢复，植物重新生长，人类重建家园。特兰克斯告别时紧紧握住大雄的手，说：“你是我见过最特别的战士。”哆啦A梦也为大雄感到骄傲，说他终于真正成长了一次。三人站在山丘上，看着远方重新明亮的地平线，心中感受到从未有过的安宁。随后，哆啦A梦与大雄乘坐时光机返回了属于他们的那个年代，一切仿佛又恢复平静。




回到现代后，大雄仿佛变了一个人，不再轻易抱怨、不再逃避责任。他认真写完作业，帮妈妈买菜，甚至主动练习体育，哆啦A梦惊讶得说不出话来。他知道，这不是一时兴起，而是大雄真正内心成长的结果。大雄有时会望着天空出神，仿佛还能看见未来世界的那一片废墟与重生的希望。他不会说出来，但他心中永远铭记那一战。




几天后，电视新闻中突然出现一则画面：一位金发少年在街头击退了失控的机器人，引发市民围观与猜测。大雄放下手中的课本，望向哆啦A梦，两人心照不宣地笑了。也许，特兰克斯又回来了，也许，新的敌人正在逼近。冒险从未真正结束，而他们，早已准备好了。无论时空如何动荡，他们将永远并肩作战。







把上面的文章喂给RAG系统，让它基于这篇文章来回答问题。

RAG系统的第一步是分片（按行分片），分片相关的代码如下：

vim rag_split.py

from typing import List


doc_file = r"/home/sam_admin/rag/doc.txt"

def split_into_chunks(doc_file: str) -> List[str]:
    with open(doc_file, mode='r', encoding='utf8') as file:
        content = file.read()

    return [chunk for chunk in content.split("\n\n")]

chunks = split_into_chunks(doc_file)

for i, chunk in enumerate(chunks):
    print(f"[{i}] {chunk}\n")






索引相关的代码：

vim rag_index.py

from sentence_transformers import SentenceTransformer
from typing import List


doc_file = r"/home/sam_admin/rag/doc.txt"

def split_into_chunks(doc_file: str) -> List[str]:
    with open(doc_file, mode='r', encoding='utf8') as file:
        content = file.read()

    return [chunk for chunk in content.split("\n\n")]

chunks = split_into_chunks(doc_file)

embedding_model = SentenceTransformer("shibing624/text2vec-base-chinese")

def embed_chunk(chunk: str) -> List[float]:
    embedding = embedding_model.encode(chunk, normalize_embeddings=True)
    return embedding.tolist()


embeddings = [embed_chunk(chunk) for chunk in chunks]
print(len(embeddings))
print(embeddings)









把向量存到向量数据库中，这里使用的向量数据库是chromadb

vim rag_savedb.py

import chromadb
from sentence_transformers import SentenceTransformer
from typing import List


doc_file = r"/home/sam_admin/rag/doc.txt"

def split_into_chunks(doc_file: str) -> List[str]:
    with open(doc_file, mode='r', encoding='utf8') as file:
        content = file.read()

    return [chunk for chunk in content.split("\n\n")]

chunks = split_into_chunks(doc_file)

embedding_model = SentenceTransformer("shibing624/text2vec-base-chinese")

def embed_chunk(chunk: str) -> List[float]:
    embedding = embedding_model.encode(chunk, normalize_embeddings=True)
    return embedding.tolist()


embeddings = [embed_chunk(chunk) for chunk in chunks]

# EphemeralClient创建内存型向量数据库,数据不会写入磁盘，脚本运行结束之后数据就会清除
chromadb_client = chromadb.EphemeralClient()
chromadb_collection = chromadb_client.get_or_create_collection(name="default")

def save_embeddings(chunks: List[str], embeddings: List[List[float]]) -> None:
    for i, (chunk, embedding) in enumerate(zip(chunks, embeddings)):
        chromadb_collection.add(
            documents=[chunk],
            embeddings=[embedding],
            ids=[str(i)]
        )

save_embeddings(chunks, embeddings)




上面的所有操作都是用户提问前的准备，下面处理用户提问后的事情，分别是召回，重排和生成。

召回：

vim rag_retrieve.py

import chromadb
from sentence_transformers import SentenceTransformer
from typing import List


doc_file = r"/home/sam_admin/rag/doc.txt"

def split_into_chunks(doc_file: str) -> List[str]:
    with open(doc_file, mode='r', encoding='utf8') as file:
        content = file.read()

    return [chunk for chunk in content.split("\n\n")]

chunks = split_into_chunks(doc_file)

embedding_model = SentenceTransformer("shibing624/text2vec-base-chinese")

def embed_chunk(chunk: str) -> List[float]:
    embedding = embedding_model.encode(chunk, normalize_embeddings=True)
    return embedding.tolist()


embeddings = [embed_chunk(chunk) for chunk in chunks]

# EphemeralClient创建内存型向量数据库,数据不会写入磁盘，脚本运行结束之后数据就会清除
chromadb_client = chromadb.EphemeralClient()
chromadb_collection = chromadb_client.get_or_create_collection(name="default")

def save_embeddings(chunks: List[str], embeddings: List[List[float]]) -> None:
    for i, (chunk, embedding) in enumerate(zip(chunks, embeddings)):
        chromadb_collection.add(
            documents=[chunk],
            embeddings=[embedding],
            ids=[str(i)]
        )

save_embeddings(chunks, embeddings)

def retrieve(query: str, top_k: int) -> List[str]:
    query_embedding = embed_chunk(query)
    results = chromadb_collection.query(
        query_embeddings=[query_embedding],
        n_results=top_k
    )
    return results['documents'][0]

query = "哆啦A梦使用的3个秘密道具分别是什么？"
retrieved_chunks = retrieve(query, 5)

for i, chunk in enumerate(retrieved_chunks):
    print(f"[{i}] {chunk}\n")



可以看到一共召回了5个结果，其中问题的答案就是第一条结果。




重排序：返回评分最高的结果，分值代表与对应片段内容的相似程度。

vim rag_rerank.py

from sentence_transformers import CrossEncoder
import chromadb
from sentence_transformers import SentenceTransformer
from typing import List


doc_file = r"/home/sam_admin/rag/doc.txt"

def split_into_chunks(doc_file: str) -> List[str]:
    with open(doc_file, mode='r', encoding='utf8') as file:
        content = file.read()

    return [chunk for chunk in content.split("\n\n")]

chunks = split_into_chunks(doc_file)

embedding_model = SentenceTransformer("shibing624/text2vec-base-chinese")

def embed_chunk(chunk: str) -> List[float]:
    embedding = embedding_model.encode(chunk, normalize_embeddings=True)
    return embedding.tolist()


embeddings = [embed_chunk(chunk) for chunk in chunks]

# EphemeralClient创建内存型向量数据库,数据不会写入磁盘，脚本运行结束之后数据就会清除
chromadb_client = chromadb.EphemeralClient()
chromadb_collection = chromadb_client.get_or_create_collection(name="default")

def save_embeddings(chunks: List[str], embeddings: List[List[float]]) -> None:
    for i, (chunk, embedding) in enumerate(zip(chunks, embeddings)):
        chromadb_collection.add(
            documents=[chunk],
            embeddings=[embedding],
            ids=[str(i)]
        )

save_embeddings(chunks, embeddings)

def retrieve(query: str, top_k: int) -> List[str]:
    query_embedding = embed_chunk(query)
    results = chromadb_collection.query(
        query_embeddings=[query_embedding],
        n_results=top_k
    )
    return results['documents'][0]

query = "哆啦A梦使用的3个秘密道具分别是什么？"
retrieved_chunks = retrieve(query, 5)

print('召回相似度检索')
for i, chunk in enumerate(retrieved_chunks):
    print(f"[{i}] {chunk}\n")



def rerank(query: str, retrieved_chunks: List[str], top_k: int) -> List[str]:
    cross_encoder = CrossEncoder('cross-encoder/mmarco-mMiniLMv2-L12-H384-v1')
    pairs = [(query, chunk) for chunk in retrieved_chunks]
    scores = cross_encoder.predict(pairs)

    scored_chunks = list(zip(retrieved_chunks, scores))
    # 按分数倒叙返回
    scored_chunks.sort(key=lambda x: x[1], reverse=True)

    return [chunk for chunk, _ in scored_chunks][:top_k]

reranked_chunks = rerank(query, retrieved_chunks, 3)

print('重排序')
for i, chunk in enumerate(reranked_chunks):
    print(f"[{i}] {chunk}\n")






最后是答案生成阶段:选用的大模型是DeepSeek-V3.1。

这里使用的是硅基流动的API调用，需要先注册一个API KEY，地址：https://cloud.siliconflow.cn/me/account/ak

vim rag_llm.py

import requests
import json
from sentence_transformers import CrossEncoder
import chromadb
from sentence_transformers import SentenceTransformer
from typing import List


doc_file = r"/home/sam_admin/rag/doc.txt"

def split_into_chunks(doc_file: str) -> List[str]:
    with open(doc_file, mode='r', encoding='utf8') as file:
        content = file.read()

    return [chunk for chunk in content.split("\n\n")]

chunks = split_into_chunks(doc_file)

embedding_model = SentenceTransformer("shibing624/text2vec-base-chinese")

def embed_chunk(chunk: str) -> List[float]:
    embedding = embedding_model.encode(chunk, normalize_embeddings=True)
    return embedding.tolist()


embeddings = [embed_chunk(chunk) for chunk in chunks]

# EphemeralClient创建内存型向量数据库,数据不会写入磁盘，脚本运行结束之后数据就会清除
chromadb_client = chromadb.EphemeralClient()
chromadb_collection = chromadb_client.get_or_create_collection(name="default")

def save_embeddings(chunks: List[str], embeddings: List[List[float]]) -> None:
    for i, (chunk, embedding) in enumerate(zip(chunks, embeddings)):
        chromadb_collection.add(
            documents=[chunk],
            embeddings=[embedding],
            ids=[str(i)]
        )

save_embeddings(chunks, embeddings)

def retrieve(query: str, top_k: int) -> List[str]:
    query_embedding = embed_chunk(query)
    results = chromadb_collection.query(
        query_embeddings=[query_embedding],
        n_results=top_k
    )
    return results['documents'][0]

query = "哆啦A梦使用的3个秘密道具分别是什么？"
retrieved_chunks = retrieve(query, 5)

def rerank(query: str, retrieved_chunks: List[str], top_k: int) -> List[str]:
    cross_encoder = CrossEncoder('cross-encoder/mmarco-mMiniLMv2-L12-H384-v1')
    pairs = [(query, chunk) for chunk in retrieved_chunks]
    scores = cross_encoder.predict(pairs)

    scored_chunks = list(zip(retrieved_chunks, scores))
    scored_chunks.sort(key=lambda x: x[1], reverse=True)

    return [chunk for chunk, _ in scored_chunks][:top_k]

reranked_chunks = rerank(query, retrieved_chunks, 3)


url = "https://api.siliconflow.cn/v1/chat/completions"
prompt = f"""你是一位知识助手，请根据用户的问题和下列片段生成准确的回答。
   用户问题: {query}
   相关片段:{"\n\n".join(reranked_chunks)}
   请基于上述内容作答，不要编造信息。"""
print(f"{prompt}\n\n---\n")
payload = {
    "model": "deepseek-ai/DeepSeek-V3.1",
    "messages": [
        {
            "role": "user",
            "content": prompt
        }
    ],
}
headers = {
    "Authorization": "Bearer API-KEY",
    "Content-Type": "application/json"
}
response = requests.request("POST", url, json=payload, headers=headers)
print(json.loads(response.text)["choices"][0]["message"]["content"])



最后可以看出模型返回的结果是没有问题的，是正确答案。

关键字：