RAG问答延迟降低42%：chunk重构与embedding选型及rerank重排全记录

Python: 3.10.12
langchain: 0.1.0 → 0.2.1（优化后升级）
chromadb: 0.4.22
sentence-transformers: 2.2.2
FlagEmbedding: 1.2.6（rerank用）
torch: 2.1.0+cu118
CUDA: 11.8
单卡: A10 24G（仅推理，batch_size=32）

from langchain.text_splitter import RecursiveCharacterTextSplitter

text_splitter = RecursiveCharacterTextSplitter(
    chunk_size=800,
    chunk_overlap=150,
    separators=[
        "\n## ",   # 二级标题
        "\n### ",  # 三级标题
        "\n\n",    # 空行
        "\n",      # 换行
        "。",      # 句号
        "；",      # 分号
        " ",       # 空格
    ],
    keep_separator=True,  # 关键：保留分隔符，避免标题丢失
)

# 对每篇文档执行
chunks = text_splitter.split_text(document)

# 给每个chunk加上父文档标题，增强上下文
chunk_with_title = f"# {doc_title}\n{chunk}"

from FlagEmbedding import BGEM3FlagModel
import chromadb
from chromadb.utils import embedding_functions

# 初始化bge-m3，使用dense+sparse混合检索
model = BGEM3FlagModel(
    'BAAI/bge-m3',
    use_fp16=True,
    device='cuda:0'
)

def encode_texts(texts):
    outputs = model.encode(
        texts,
        batch_size=32,
        max_length=1024,
        return_dense=True,
        return_sparse=True,
        return_colbert_vecs=False  # 省显存
    )
    # 混合向量：dense + sparse加权（sparse权重0.3）
    dense = outputs['dense_vecs']
    sparse = outputs['lexical_weights']
    return dense, sparse

# chromadb自带的embedding_function需要统一接口，
# 这里我们直接离线生成向量后写入
collection = client.get_or_create_collection(
    name="docs_v2",
    metadata={"hnsw:space": "cosine"}
)

# 写入时同时存dense和sparse（sparse存为metadata）
for i, chunk in enumerate(chunks):
    dense, sparse = encode_texts([chunk])
    # sparse weights需要序列化存储，这里简化为存top-50词
    sparse_simplified = {k: v for k, v in list(sparse[0].items())[:50]}
    collection.add(
        ids=[f"chunk_{i}"],
        embeddings=dense.tolist(),
        metadatas=[{"text": chunk, "sparse": json.dumps(sparse_simplified)}]
    )

def hybrid_search(query, top_k=20):
    q_dense, q_sparse = encode_texts([query])
    # chromadb只支持dense检索，sparse部分我们手动算BM25
    dense_results = collection.query(
        query_embeddings=q_dense.tolist(),
        n_results=top_k
    )
    # sparse部分：用jieba分词后算BM25，与dense分数加权
    bm25_scores = compute_bm25(query, collection)
    # 最终分数 = 0.7 * dense_cosine + 0.3 * normalized_bm25
    final_scores = 0.7 * dense_sim + 0.3 * bm25_norm
    return sorted(zip(ids, final_scores), key=lambda x: x[1], reverse=True)[:top_k]

from FlagEmbedding import FlagReranker

reranker = FlagReranker(
    'BAAI/bge-reranker-base',
    use_fp16=True,
    device='cuda:0'
)

def rerank(query, candidates, top_n=5):
    # candidates格式：[(id, text, score), ...]
    pairs = [[query, cand[1]] for cand in candidates]
    scores = reranker.compute_score(
        pairs,
        normalize=True,  # 返回0-1之间的sigmoid分数
        batch_size=16
    )
    # 合并分数并排序
    scored = [(cand[0], cand[1], cand[2], scores[i]) for i, cand in enumerate(candidates)]
    # 注意：这里我们直接用rerank分数排序，不再与向量分数融合
    # 因为rerank是交互式，更可靠
    scored.sort(key=lambda x: x[3], reverse=True)
    return scored[:top_n]

# 使用原始logits + 温度缩放
logits = reranker.compute_score(pairs, normalize=False)
scaled = [1 / (1 + np.exp(-(x / 0.3))) for x in logits]  # 温度0.3