RAG召回率从61%到89%：chunk重切、embedding换型与rerank排序的工程实践

Python 3.10.12
langchain 0.1.16
langchain-community 0.0.36
faiss-cpu 1.7.4
transformers 4.38.2
sentence-transformers 2.5.1
torch 2.2.1+cu121

from langchain.text_splitter import MarkdownHeaderTextSplitter, RecursiveCharacterTextSplitter

# 第一遍：按标题切出顶级语义块
headers_to_split_on = [
    ("#", "H1"),
    ("##", "H2"),
    ("###", "H3"),
]
md_splitter = MarkdownHeaderTextSplitter(headers_to_split_on=headers_to_split_on)

# 第二遍：对每个标题块内部，如果过长再递归切分
chunk_size = 512  # 字符级，比之前256大一倍，因为语义块本身有上下文
chunk_overlap = 128

def hybrid_split(text: str):
    sections = md_splitter.split_text(text)
    final_chunks = []
    for section in sections:
        # 关键：把标题拼回到内容里，让embedding能感知上下文
        header_text = " > ".join(section.metadata.values()) if section.metadata else ""
        content = f"{header_text}\n{section.page_content}"

        if len(content) <= chunk_size:
            final_chunks.append(content)
        else:
            # 用递归切分器兜底，但保持overlap让边界信息不丢失
            sub_splitter = RecursiveCharacterTextSplitter(
                chunk_size=chunk_size,
                chunk_overlap=chunk_overlap,
                separators=["\n\n", "\n", "。", "；", " ", ""]  # 注意中文标点优先
            )
            sub_chunks = sub_splitter.split_text(content)
            final_chunks.extend(sub_chunks)
    return final_chunks

from sentence_transformers import SentenceTransformer

class BGEEmbedding:
    def __init__(self, model_path: str = "/models/bge-large-zh-v1.5"):
        # 关键参数：normalize_embeddings=True，query指令必须明确
        self.model = SentenceTransformer(model_path, device="cuda:0")
        self.model.max_seq_length = 512  # bge支持最长512，对chunk长度友好
        self.query_instruction = "为这个句子生成表示以用于检索相关文章："
        self.doc_instruction = ""  # 文档不需要前缀，或者可以留空

    def embed_query(self, text: str):
        return self.model.encode(
            self.query_instruction + text,
            normalize_embeddings=True,
            convert_to_numpy=True
        )

    def embed_documents(self, texts: list[str]):
        # 文档不加指令，直接编码
        return self.model.encode(
            texts,
            batch_size=32,
            normalize_embeddings=True,
            convert_to_numpy=True,
            show_progress_bar=False
        )

from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch

class Reranker:
    def __init__(self, model_path: str = "/models/bge-reranker-base"):
        self.tokenizer = AutoTokenizer.from_pretrained(model_path)
        self.model = AutoModelForSequenceClassification.from_pretrained(model_path)
        self.model.eval()
        self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
        self.model.to(self.device)

    def rerank(self, query: str, docs: list[str], top_k: int = 5):
        pairs = [[query, doc] for doc in docs]
        inputs = self.tokenizer(
            pairs,
            padding=True,
            truncation=True,
            max_length=512,
            return_tensors="pt"
        ).to(self.device)

        with torch.no_grad():
            scores = self.model(**inputs).logits.squeeze(-1)  # shape: [batch_size]
            scores = torch.sigmoid(scores).cpu().numpy()

        # 按分数降序，取top_k
        doc_scores = list(zip(docs, scores))
        doc_scores.sort(key=lambda x: x[1], reverse=True)
        return doc_scores[:top_k]

def evaluate_pipeline(pipeline, test_queries, gold_doc_ids, k=5):
    hit = 0
    mrr_sum = 0.0
    total = len(test_queries)

    for query, gold_id in zip(test_queries, gold_doc_ids):
        retrieved = pipeline.retrieve(query, top_k=k)
        retrieved_ids = [doc_id for doc_id, _ in retrieved]

        if gold_id in retrieved_ids:
            hit += 1
            rank = retrieved_ids.index(gold_id) + 1
            mrr_sum += 1.0 / rank

    hit_rate = hit / total
    mrr = mrr_sum / total
    return {"hit@k": round(hit_rate, 4), "mrr": round(mrr, 4)}