BM25 + Embedding 混合检索 实现

我们做一个标准的：

🔎 BM25（关键词召回） + Embedding（语义召回） + 融合排序

---

一、整体架构

                用户 Query
                     ↓
        ┌─────────────────────────┐
        │                         │
   BM25 关键词检索           Embedding 语义检索
        │                         │
        └──────────┬──────────────┘
                   ↓
              候选结果合并
                   ↓
              融合打分 / RRF
                   ↓
                 TopK

---

二、实现步骤（完整示例）

下面给你一个最小可用版本。

假设：

• 文档是 List[str]
• 使用 rank_bm25
• 使用 sentence-transformers 做 embedding
• 使用 FAISS 做向量搜索（推荐）

---

1️⃣ 安装依赖

pip install rank-bm25
pip install sentence-transformers
pip install faiss-cpu

---

2️⃣ 构建混合检索类

import jieba
import numpy as np
import faiss
from rank_bm25 import BM25Okapi
from sentence_transformers import SentenceTransformer
from typing import List


class HybridRetriever:
    def __init__(self, model_name="all-MiniLM-L6-v2"):
        # Embedding 模型
        self.model = SentenceTransformer(model_name)

        self.bm25 = None
        self.documents = None
        self.embeddings = None
        self.index = None  # FAISS 索引

    # ======================
    # 1. 构建索引
    # ======================
    def build(self, documents: List[str]):
        self.documents = documents

        # ---------- BM25 ----------
        tokenized_docs = [list(jieba.cut(doc)) for doc in documents]
        self.bm25 = BM25Okapi(tokenized_docs)

        # ---------- Embedding ----------
        self.embeddings = self.model.encode(documents)

        # 归一化（用于 cosine）
        self.embeddings = np.array(self.embeddings).astype("float32")
        faiss.normalize_L2(self.embeddings)

        dim = self.embeddings.shape[1]
        self.index = faiss.IndexFlatIP(dim)  # 内积 = cosine
        self.index.add(self.embeddings)

    # ======================
    # 2. 混合检索
    # ======================
    def search(self, query: str, top_k=5, bm25_weight=0.5):

        # ---------- BM25 ----------
        tokenized_query = list(jieba.cut(query))
        bm25_scores = self.bm25.get_scores(tokenized_query)

        # 取 BM25 TopN
        bm25_top = np.argsort(bm25_scores)[::-1][:top_k*5]

        # ---------- Embedding ----------
        query_vec = self.model.encode([query]).astype("float32")
        faiss.normalize_L2(query_vec)

        D, I = self.index.search(query_vec, top_k*5)

        embedding_top = I[0]

        # ---------- 候选集合 ----------
        candidate_ids = set(bm25_top).union(set(embedding_top))

        # ---------- 融合打分（加权） ----------
        results = []

        for idx in candidate_ids:
            bm25_score = bm25_scores[idx]

            # embedding 相似度
            emb_score = 0
            if idx in embedding_top:
                emb_score = 1  # 简化处理

            final_score = bm25_weight * bm25_score + \
                          (1 - bm25_weight) * emb_score

            results.append((idx, final_score))

        # 排序
        results.sort(key=lambda x: x[1], reverse=True)

        return [self.documents[i] for i, _ in results[:top_k]]

---

三、这里发生了什么？

🔹 BM25 负责：

• 精确关键词匹配
• 函数名
• 错误码
• 专业术语

---

🔹 Embedding 负责：

• 同义表达
• 自然语言问法
• 语义理解
• 模糊问题

---

🔹 融合方式

这里用了：

加权融合

更高级一点可以用：

• RRF（推荐）
• Cross-Encoder rerank（更强）

---

四、更工业级版本（推荐）

真正生产系统一般是：

1️⃣ BM25 召回 100 条
2️⃣ Embedding 召回 100 条
3️⃣ 合并
4️⃣ 用 Cross-Encoder 重新排序

Cross-Encoder 是：

把 query 和 document 拼在一起

让模型判断相关性

效果最好。

---

五、什么时候必须用混合检索？

如果你的系统：

• 是代码问答
• 是 Bug 定位
• 是知识库问答
• 是 RAG 系统
• 用户输入是自然语言

👉 强烈建议混合。

---

六、为什么混合效果好？

因为：

BM25 = 高精度召回

Embedding = 高语义召回

两者覆盖不同类型问题。

单独用任何一个都会有盲区。

---

七、如果你想接入你现有代码

你现在有：

KeywordIndex (TF-IDF)

升级路径：

1. 把 TF-IDF 换成 BM25
2. 新增 EmbeddingIndex
3. 写一个 HybridRetriever
4. 对外只暴露一个 search()

---

八、我给你一个实战建议

如果你是做：

• ChatBI
• 代码分析系统
• Bug 定位系统

推荐结构：

Hybrid Retrieval
      ↓
LLM Rerank
      ↓
最终答案生成

这是当前主流 RAG 架构。

---