从62%到89%：一次RAG检索增强生成系统的三阶段优化记录

Python 3.11.6
langchain 0.1.12
langchain-community 0.0.28
pymilvus 2.3.4
sentence-transformers 2.6.1
FlagEmbedding 1.2.10
torch 2.2.1 (CUDA 12.1)
openai 1.14.2

import re
import numpy as np
from sentence_transformers import SentenceTransformer
from langchain.text_splitter import RecursiveCharacterTextSplitter

class SemanticChunker:
    def __init__(self, model_name="BAAI/bge-large-zh-v1.5",
                 threshold=0.72, max_chunk=384, overlap=80):
        self.model = SentenceTransformer(model_name)
        self.threshold = threshold
        self.max_chunk = max_chunk
        self.overlap = overlap
        self.splitter = RecursiveCharacterTextSplitter(
            chunk_size=max_chunk,
            chunk_overlap=overlap,
            separators=["\n\n", "\n", "。", "！", "？", "；", "，", " ", ""],
            keep_separator=True,
        )

    def _split_sentences(self, text):
        # 保留代码块完整性
        parts = re.split(r'(\n\n+)', text)
        sents = []
        for p in parts:
            if p.strip():
                sents.extend(re.split(r'(? self.max_chunk:
                chunks.append(cur)
                # overlap：把上一块尾部拼到新块
                tail = cur[-self.overlap:] if len(cur) > self.overlap else cur
                cur = tail + nxt
            else:
                cur += nxt
        if cur:
            chunks.append(cur)
        # 二次递归切分兜底
        final = []
        for c in chunks:
            final.extend(self.splitter.split_text(c) if len(c) > self.max_chunk else [c])
        return final

from FlagEmbedding import FlagReranker
from pymilvus import Collection
from sentence_transformers import SentenceTransformer

class RagRetriever:
    def __init__(self, collection: Collection, top_k=20, rerank_top=5):
        self.collection = collection
        self.embedder = SentenceTransformer("BAAI/bge-large-zh-v1.5")
        self.reranker = FlagReranker("BAAI/bge-reranker-large",
                                     use_fp16=True)
        self.top_k = top_k
        self.rerank_top = rerank_top
        # BGE推荐的query指令前缀
        self.query_prefix = "为这个句子生成表示以用于检索相关文章："

    def _embed_query(self, query: str):
        return self.embedder.encode(
            self.query_prefix + query,
            normalize_embeddings=True
        ).tolist()

    def retrieve(self, query: str):
        # 1) 向量粗排
        self.collection.load()
        res = self.collection.search(
            data=[self._embed_query(query)],
            anns_field="embedding",
            param={"metric_type": "IP", "params": {"ef": 128}},
            limit=self.top_k,
            output_fields=["text", "doc_id", "chunk_id"],
        )
        candidates = [
            {"text": h.entity.get("text"),
             "doc_id": h.entity.get("doc_id"),
             "score": float(h.distance)}
            for h in res[0]
        ]
        if not candidates:
            return []

        # 2) Rerank精排
        pairs = [[query, c["text"]] for c in candidates]
        scores = self.reranker.compute_score(pairs, normalize=True)
        for c, s in zip(candidates, scores):
            c["rerank_score"] = float(s)
        candidates.sort(key=lambda x: x["rerank_score"], reverse=True)
        return candidates[:self.rerank_top]