搭建简易版的Rag系统

Rag系统搭建

简易RAG（检索增强生成）搭建概述

一、核心原理

RAG = 检索（Retrieval）+ 增强（Augmentation）+ 生成（Generation）

不用微调大模型，先从本地文档找相关内容 → 把内容塞进Prompt → 让大模型基于资料回答，解决AI幻觉、知识库过时问题。

二、整体流程

文档加载与分块
读取PDF/TXT/MD等文档，切成小块（Chunk），避免上下文过长、检索不准。

python 复制代码

def split_into_chunks(doc_file):
    with open(doc_file,'r', encoding='utf-8') as file:
        content = file.read()
        return [chunk for chunk in content.split('\n\n')]

chunks = split_into_chunks("你的文档路径",)  #将文档划分为几个片段

# 导入完成后可以进行检查
# for i,chunk in enumerate(chunks):
#     print(f"[{i},{chunk}")

文本向量化 & 存入向量库
用Embedding模型把文本块转成向量，存入轻量向量库（如FAISS、Chroma）。

python 复制代码

# 如果运行不显示以下错误，可以跳过这几行代码
#(MaxRetryError("HTTPSConnectionPool(host='huggingface.co', port=443): Max retries exceeded with url: #/shibing624/text2vec-base-chinese/resolve/main/modules.json (Caused by #ConnectTimeoutError(<HTTPSConnection(host='huggingface.co', port=443) at 0x1b570fb5f60>, 'Connection to huggingface.co #timed out. (connect timeout=10)'))"), '(Request ID: 9a7e30b1-4e49-4bae-9f40-8036cc31cfa9)')' thrown while requesting HEAD #https://huggingface.co/shibing624/text2vec-base-chinese/resolve/main/modules.json
#Retrying in 1s [Retry 1/5].
# 用以下代码的含义：
# 1. 用CPU跑，不用显卡
# 2. 走国内镜像下载模型，更快更稳
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"

# 从以下这几行开始写
from sentence_transformers import SentenceTransformer
embedding_model = SentenceTransformer('shibing624/text2vec-base-chinese')

def embed_chunk(chunk):
    embedding = embedding_model.encode(chunk)
    return embedding.tolist()

'''  进行测试
test_embedding = embed_chunk("测试内容")
print(len(test_embedding))
print(test_embedding)
'''

用户提问→相似度检索
用户问题也转向量，在向量库召回最相似的N个文档片段。

python 复制代码

import chromadb

chromadb_client = chromadb.EphemeralClient()  #创建临时保存
chromadb_collection = chromadb_client.get_or_create_collection(name="default")

def save_embedding(chunk,embedding):  # 保存块与之对应的embedding
    ids = [str(i) for i in range(len(chunk))]
    chromadb_collection.add(
        documents = chunk,
        embeddings = embeddings,
        ids = ids
    )
save_embedding(chunks,embeddings)

def retrieve(query,top_k):   # 选出相似度较高的前top_k个片段
    query_embedding = embed_chunk(query)
    result = chromadb_collection.query(
        query_embeddings=[query_embedding],
        n_results=top_k
    )
    return result['documents'][0]

query = "这篇文章的主旨大意是什么？"   # 用户提出的问题
receive_chunks = retrieve(query,5)  #找到整个关键片段

# 查看相似的片段
for i ,chunk in enumerate(receive_chunks):
    print(f"{i},{chunk}\n")

from sentence_transformers import CrossEncoder
def rerank(quary,receive_chunks,top_k):  #给找到的关键字重新排序，在选出关键片段
    cross_encode = CrossEncoder('cross-encoder/mmarco-mMiniLMv2-L12-H384-v1')
    pairs = [(quary,chunk) for chunk in receive_chunks]
    scores = cross_encode.predict(pairs)
    chunk_with_score_list = [(chunk,score)
                        for chunk,score in zip(receive_chunks,scores)]
    chunk_with_score_list.sort(key=lambda pairs:pairs[1],reverse=True)
    return [chunk for chunk,_ in chunk_with_score_list][:top_k]

rerank_chunks = rerank(query,receive_chunks,3)
# 查看相思的片段
for i,chunk in enumerate(rerank_chunks):
    print(f"[{i}]{chunk}\n")

构造Prompt→大模型生成答案
把问题+检索到的片段拼成提示词，传给大模型，生成回答。

python 复制代码

from ollama import chat  #调用大模型，给出作答
def generate(query,chunks):
    chunks_content = "\n\n".join(chunks)
    prompt = f""" 你是一个知识助手，请根据用户的问题和下列产生的片段生成准确的回答。
    用户问题:{query}
    相关片段:{chunks_content}
    请基于上述内容作答，不要编造信息"""
    print(f"{prompt}\n\n---\n")

    response = chat(
        model='qwen2.5:0.5b',
        messages=[{"role": "user", "content": prompt}]
    )
    return response["message"]["content"]

answer = generate(query,chunks)
print(answer)  #做出回答