用 BM25 与向量混合检索提升 RAG 召回

# 依赖：pip install rank_bm25==0.2.2 sentence-transformers==2.7.0 numpy==1.26.4
# 纯向量检索对关键词精确匹配（如专有名词、型号、编号）不敏感，容易漏掉含关键字的片段。
# BM25 基于词频统计，擅长精确关键词匹配；二者互补。混合检索先分别召回再加权融合，召回完整度通常优于单路。
import numpy as np
from rank_bm25 import BM25Okapi
from sentence_transformers import SentenceTransformer

docs = [
    "产品A型号 X-200 支持私有化部署，适用于金融场景。",
    "产品B型号 Y-500 仅提供 SaaS，适合中小团队。",
    "RAG 系统先用向量检索召回片段，再让大模型生成答案。",
    "金融场景对数据合规要求高，私有化部署更受青睐。",
]
tokenized = [d.split() for d in docs]
bm25 = BM25Okapi(tokenized)
embedder = SentenceTransformer("BAAI/bge-small-zh")
doc_vecs = embedder.encode(docs, normalize_embeddings=True)

def hybrid_search(query: str, k: int = 2, alpha: float = 0.5):
    # BM25 路径：按词匹配打分
    bm25_scores = np.array(bm25.get_scores(query.split()))
    # 向量路径：语义相似度
    qv = embedder.encode([query], normalize_embeddings=True)[0]
    vec_scores = doc_vecs @ qv
    # 各自归一化后加权融合
    bm25_n = (bm25_scores - bm25_scores.min()) / (bm25_scores.max() - bm25_scores.min() + 1e-9)
    vec_n = (vec_scores - vec_scores.min()) / (vec_scores.max() - vec_scores.min() + 1e-9)
    fused = alpha * vec_n + (1 - alpha) * bm25_n
    return [docs[i] for i in np.argsort(-fused)[:k]]

if __name__ == "__main__":
    print(hybrid_search("金融场景 私有化部署 哪个型号"))

['产品A型号 X-200 支持私有化部署，适用于金融场景。', '金融场景对数据合规要求高，私有化部署更受青睐。']