从学术检索到生产级 RAG：BM25 算法的核心转化与极简实现

从学术检索到生产级 RAG：BM25 算法的核心转化与极简实现

在企业级检索增强生成（RAG）系统开发中，如何从本地知识库快速定位与用户问题最相关的文档，直接影响最终生成效果。不少团队会优先考虑向量数据库或深度学习检索方案（如 Dense Retrieval），但在实际部署中，传统 BM25 算法凭借其零显存占用、CPU 快速响应和关键词精准匹配能力，仍是可靠的选择。

一、学术方案与企业数据的适配难题

学术界常用 MS MARCO 等公开数据集评估新型检索器（如双塔模型）的 NDCG 指标。然而，当这类模型直接应用于企业 RAG 环境时，常遇到实际问题：

核心工程目标：用最简代码实现经典 BM25 算法，作为高性价比的本地召回方案。

二、BM25 检索流程解析

BM25 在 TF-IDF 基础上引入文档长度惩罚和词频截断机制，使相关性评分更贴近人工评估。其处理流程如下：

graph TD
A[注册文档库] –> B[分词并统计基础指标]
B –> C[计算平均文档长度 avgdl]
C –> D[计算逆文档频率 IDF]
E[用户查询] –> F[查询分词]
F –> G[计算词频 TF]
G –> H[代入 BM25 公式得分]
H –> I[按得分排序]
I –> J[返回 Top-N 结果]

通过长度惩罚因子，该算法能有效过滤仅因篇幅长而高频出现关键词的文档，提升召回质量。

三、纯 Python 实现的 BM25 引擎

以下代码实现不依赖外部数值计算库，具备以下特性：

• 支持文档注册与动态索引更新
• 自动计算文档平均长度
• 内置 IDF 负值防护机制

# bm25_retriever.py
import math
from typing import List, Dict, Tuple

class BM25Retriever:
def __init__(self, k1: float = 1.5, b: float = 0.75):
self.k1 = k1
self.b = b
self.doc_count = 0
self.avg_doc_len = 0.0
self.doc_lengths = []
self.doc_term_freqs = []
self.doc_freqs = {}
self.idfs = {}
self.raw_documents = []

def register_documents(self, documents: List[str]):
self.raw_documents = documents
self.doc_count = len(documents)

total_len = 0
for doc in documents:
words = self._tokenize(doc)
doc_len = len(words)
total_len += doc_len
self.doc_lengths.append(doc_len)

tf_dict = {}
for word in words:
tf_dict[word] = tf_dict.get(word, 0) + 1
self.doc_term_freqs.append(tf_dict)

for word in tf_dict.keys():
self.doc_freqs[word] = self.doc_freqs.get(word, 0) + 1

self.avg_doc_len = total_len / self.doc_count if self.doc_count else 0.0
self._calculate_idfs()

def _tokenize(self, text: str) -> List[str]:
normalized = text.lower()
for char in ".,!?[](){}:;\\"'<>~`|-_+=":
normalized = normalized.replace(char, " ")
return [w for w in normalized.split() if w.strip()]

def _calculate_idfs(self):
for word, df in self.doc_freqs.items():
raw_idf = math.log((self.doc_count – df + 0.5) / (df + 0.5) + 1.0)
self.idfs[word] = max(raw_idf, 1e-4)

def retrieve(self, query: str, top_n: int = 3) -> List[Tuple[str, float]]:
query_words = self._tokenize(query)
scores = []

for idx in range(self.doc_count):
score = 0.0
doc_len = self.doc_lengths[idx]
tf_dict = self.doc_term_freqs[idx]

for word in query_words:
if word not in tf_dict:
continue

tf = tf_dict[word]
idf = self.idfs.get(word, 0.0)

numerator = tf * (self.k1 + 1)
denominator = tf + self.k1 * (1 – self.b + self.b * (doc_len / self.avg_doc_len))
score += idf * (numerator / denominator)

scores.append((idx, score))

scores.sort(key=lambda x: x[1], reverse=True)

return [(self.raw_documents[doc_idx], score)
for doc_idx, score in scores[:top_n]]

# 测试示例
if __name__ == "__main__":
mock_docs = [
"In AI, RAG is a method to fetch documents using vector search.",
"BM25 is a term matching algorithm based on tf-idf statistics.",
"RAG applications often use BM25 to search exact keywords from databases.",
"Vector search has semantic understanding, but fails at precise product code matching."
]

retriever = BM25Retriever()
retriever.register_documents(mock_docs)

query_text = "BM25 keyword search RAG"
hits = retriever.retrieve(query_text, top_n=2)

print(f"Query: '{query_text}'")
for doc, score in hits:
print(f"[Score: {score:.4f}] Doc: {doc}")

质量评估

主要修改：