RAG 系统进阶：混合检索 + 重排序如何让大模型回答更精准？

pip install langchain chromadb sentence-transformers rank_bm25 openai tiktoken pypdf

from langchain.document_loaders import PyPDFDirectoryLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter

loader = PyPDFDirectoryLoader("./docs/")
raw_docs = loader.load()

text_splitter = RecursiveCharacterTextSplitter(
    chunk_size=512,
    chunk_overlap=50,
    separators=["\n\n", "\n", "。", "！", "？", "；", "，", " ", ""]
)
chunks = text_splitter.split_documents(raw_docs)
print(f"共生成 {len(chunks)} 个文本块")

from sentence_transformers import SentenceTransformer
import chromadb
from chromadb.utils import embedding_functions

# 初始化嵌入模型（本地加载）
embed_model = SentenceTransformer("BAAI/bge-large-zh-v1.5")
embedding_fn = embedding_functions.SentenceTransformerEmbeddingFunction(
    model_name="BAAI/bge-large-zh-v1.5"
)

# 创建 Chroma 向量库
client = chromadb.PersistentClient(path="./chroma_db")
collection = client.get_or_create_collection(
    name="langchain_docs",
    embedding_function=embedding_fn
)

# 批量插入（为演示，简化为循环）
for i, chunk in enumerate(chunks):
    collection.add(
        ids=[f"chunk_{i}"],
        documents=[chunk.page_content],
        metadatas=[chunk.metadata]
    )

from rank_bm25 import BM25Okapi
import jieba  # 中文分词

# 对所有文本块进行分词
tokenized_corpus = [list(jieba.cut(chunk.page_content)) for chunk in chunks]
bm25 = BM25Okapi(tokenized_corpus)

# 保存原始文本列表供后续使用
corpus_texts = [chunk.page_content for chunk in chunks]

def hybrid_search(query: str, top_k: int = 20, alpha: float = 0.5):
    # 1. 向量检索
    vec_results = collection.query(
        query_texts=[query],
        n_results=top_k * 2  # 多召回一些用于融合
    )
    vec_ids = vec_results['ids'][0]
    vec_scores = vec_results['distances'][0]  # 距离越小越相似
    
    # 2. BM25 检索
    tokenized_query = list(jieba.cut(query))
    bm25_scores = bm25.get_scores(tokenized_query)
    # 取 top_k*2 个索引
    bm25_top_indices = sorted(range(len(bm25_scores)), key=lambda i: bm25_scores[i], reverse=True)[:top_k*2]
    bm25_top_scores = [bm25_scores[i] for i in bm25_top_indices]
    
    # 3. 构建排序字典 (doc_id -> RRF score)
    rrf_scores = {}
    k = 60
    
    # 处理向量结果
    for rank, doc_id in enumerate(vec_ids):
        rrf_scores[doc_id] = rrf_scores.get(doc_id, 0) + 1 / (k + rank + 1)
    
    # 处理 BM25 结果（需要将索引映射为 doc_id，这里我们使用 "chunk_{i}" 格式）
    for rank, idx in enumerate(bm25_top_indices):
        doc_id = f"chunk_{idx}"
        rrf_scores[doc_id] = rrf_scores.get(doc_id, 0) + 1 / (k + rank + 1)
    
    # 排序取 top_k
    sorted_docs = sorted(rrf_scores.items(), key=lambda x: x[1], reverse=True)[:top_k]
    final_ids = [doc_id for doc_id, _ in sorted_docs]
    
    # 获取对应文本
    final_texts = []
    for doc_id in final_ids:
        idx = int(doc_id.split("_")[1])
        final_texts.append(corpus_texts[idx])
    
    return final_texts, final_ids

from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch

rerank_model_name = "BAAI/bge-reranker-large"
tokenizer = AutoTokenizer.from_pretrained(rerank_model_name)
model = AutoModelForSequenceClassification.from_pretrained(rerank_model_name)
model.eval()

def rerank(query: str, candidate_texts: list, top_k: int = 5):
    pairs = [[query, text] for text in candidate_texts]
    with torch.no_grad():
        inputs = tokenizer(pairs, padding=True, truncation=True, return_tensors="pt", max_length=512)
        scores = model(**inputs, return_dict=True).logits.view(-1, ).float()
        # 得分越高越相关
        sorted_indices = torch.argsort(scores, descending=True).tolist()
        sorted_texts = [candidate_texts[i] for i in sorted_indices]
        sorted_scores = [scores[i].item() for i in sorted_indices]
    return sorted_texts[:top_k], sorted_scores[:top_k]

import openai

def generate_answer(query: str, contexts: list):
    context_str = "\n\n".join([f"[文档{i+1}] {text}" for i, text in enumerate(contexts)])
    prompt = f"""你是一个专业的技术助手，请基于以下参考文档回答用户问题。如果文档中没有相关信息，请明确告知。

参考文档：
{context_str}

用户问题：{query}

请给出简洁、准确的回答（不超过200字）："""
    
    response = openai.ChatCompletion.create(
        model="gpt-4",
        messages=[{"role": "user", "content": prompt}],
        temperature=0.1,
    )
    return response.choices[0].message.content