RAG检索增强生成系统优化：chunk策略、embedding切换与rerank效果对比

Python 3.10.13
langchain 0.1.20
langchain-community 0.0.38
chromadb 0.4.24
sentence-transformers 2.7.0
FlagEmbedding 1.2.10
torch 2.2.1 + cu121
openai 1.30.1

from langchain.text_splitter import MarkdownHeaderTextSplitter, RecursiveCharacterTextSplitter
import tiktoken

def semantic_chunk(md_text: str, max_tokens: int = 256, overlap_ratio: float = 0.1):
    # 1. 按标题层级切
    headers = [("#", "h1"), ("##", "h2"), ("###", "h3")]
    md_splitter = MarkdownHeaderTextSplitter(headers_to_split_on=headers)
    sections = md_splitter.split_text(md_text)

    # 2. 段落内按token切，保留10%重叠
    enc = tiktoken.get_encoding("cl100k_base")
    token_splitter = RecursiveCharacterTextSplitter.from_tiktoken_encoder(
        encoding_name="cl100k_base",
        chunk_size=max_tokens,
        chunk_overlap=int(max_tokens * overlap_ratio),
        separators=["\n\n", "\n", "。", "；", " ", ""]
    )

    chunks = []
    for sec in sections:
        header_path = " > ".join([f"{k}:{v}" for k, v in sec.metadata.items()])
        for sub in token_splitter.split_text(sec.page_content):
            # 把标题路径拼到chunk前面，增强上下文
            chunks.append({
                "text": f"[{header_path}]\n{sub}",
                "metadata": sec.metadata
            })
    return chunks

from sentence_transformers import SentenceTransformer
import numpy as np

class BGEEmbedder:
    def __init__(self, model_path="BAAI/bge-large-zh-v1.5", device="cuda"):
        self.model = SentenceTransformer(model_path, device=device)
        self.query_instruction = "为这个句子生成表示以用于检索相关文章："

    def encode_queries(self, queries, batch_size=32):
        queries = [self.query_instruction + q for q in queries]
        return self.model.encode(
            queries, batch_size=batch_size,
            normalize_embeddings=True,  # 必须归一化，配合内积
            show_progress_bar=False
        )

    def encode_passages(self, passages, batch_size=64):
        return self.model.encode(
            passages, batch_size=batch_size,
            normalize_embeddings=True,
            show_progress_bar=False
        )

# 重建索引
embedder = BGEEmbedder()
texts = [c["text"] for c in chunks]
embs = embedder.encode_passages(texts)
collection.add(
    ids=[str(i) for i in range(len(texts))],
    embeddings=embs.tolist(),
    documents=texts,
    metadatas=[c["metadata"] for c in chunks]
)

from FlagEmbedding import FlagReranker

class RerankRetriever:
    def __init__(self, collection, embedder, reranker_path="BAAI/bge-reranker-large"):
        self.collection = collection
        self.embedder = embedder
        self.reranker = FlagReranker(reranker_path, use_fp16=True)

    def retrieve(self, query, recall_k=20, top_k=5):
        # 1. 向量召回
        q_emb = self.embedder.encode_queries([query])[0]
        res = self.collection.query(
            query_embeddings=[q_emb.tolist()],
            n_results=recall_k
        )
        docs = res["documents"][0]
        ids = res["ids"][0]

        # 2. rerank
        pairs = [[query, d] for d in docs]
        scores = self.reranker.compute_score(pairs, normalize=True)

        # 3. 重排取Top-K
        ranked = sorted(zip(ids, docs, scores), key=lambda x: x[2], reverse=True)
        return ranked[:top_k]