BM25算法在语义相似度检索中的实战应用与优化指南

1次阅读

共计 2499 个字符，预计需要花费 7 分钟才能阅读完成。

在信息检索领域，如何高效准确地计算文本语义相似度是开发者面临的常见挑战。传统的关键词匹配方法（如布尔检索）往往无法处理同义词、多义词等复杂语义场景，导致检索结果相关性差。而基于统计的语言模型 BM25（Best Matching 25）通过引入词频（TF）、逆文档频率（IDF）和文档长度归一化等机制，显著提升了检索效果。

相比 TF-IDF，BM25 的优势在于：

对词频的非线性处理：避免高频词的过度权重
文档长度归一化：解决长文档与短文档的公平性比较
可调节参数：适应不同数据分布特性

TF-IDF 是信息检索的基础算法，其核心思想是：

词频（TF）：单词在文档中出现的频率
逆文档频率（IDF）：单词在整个语料库中的稀有程度

而 BM25 在 TF-IDF 基础上进行了三项重要改进：

饱和 TF：使用参数 k 控制词频的影响上限
文档长度归一化：通过参数 b 调节长文档的惩罚力度
查询项权重：考虑查询词在查询中的分布

BM25 评分公式的核心部分：

score(D,Q) = Σ IDF(qi) * (f(qi,D) * (k1 + 1)) / (f(qi,D) + k1 * (1 - b + b * |D| / avgdl))

其中：
– f(qi,D)：词 qi 在文档 D 中的词频
– |D|：文档 D 的长度（词数）
– avgdl：语料库中文档的平均长度
– k1, b：可调节参数（通常 k1∈[1.2,2.0], b=0.75）

import math
from collections import defaultdict
import numpy as np

class BM25:
    def __init__(self, k1=1.5, b=0.75):
        self.k1 = k1
        self.b = b
        self.doc_lengths = []
        self.avgdl = 0
        self.doc_freqs = []
        self.idf = {}
        self.doc_len = 0
        self.doc_count = 0

    def fit(self, documents):
        """构建 BM25 模型"""
        self.doc_count = len(documents)
        self.doc_lengths = [len(doc) for doc in documents]
        self.avgdl = sum(self.doc_lengths) / self.doc_count

        # 计算词项文档频率
        freq = defaultdict(int)
        for doc in documents:
            seen = set()
            for word in doc:
                if word not in seen:
                    freq[word] += 1
                    seen.add(word)

        # 计算 IDF
        for word, count in freq.items():
            self.idf[word] = math.log((self.doc_count - count + 0.5) / (count + 0.5) + 1)

    def score(self, query, document):
        """计算单个文档的 BM25 得分"""
        score = 0.0
        doc_len = len(document)
        frequencies = defaultdict(int)
        for word in document:
            frequencies[word] += 1

        for word in query:
            if word not in frequencies:
                continue
            idf = self.idf.get(word, 0)
            tf = frequencies[word]
            numerator = tf * (self.k1 + 1)
            denominator = tf + self.k1 * (1 - self.b + self.b * (doc_len / self.avgdl))
            score += idf * (numerator / denominator)

        return score

    def search(self, query, documents):
        """对所有文档进行检索排序"""
        scores = [(i, self.score(query, doc)) for i, doc in enumerate(documents)]
        return sorted(scores, key=lambda x: x[1], reverse=True)

# 示例用法
if __name__ == "__main__":
    # 示例文档集（已分词）docs = [["苹果", "手机", "发布", "新品"],
        ["华为", "发布", "新款", "智能手机"],
        ["苹果", "公司", "财报", "超预期"],
        ["智能手机", "市场", "竞争", "激烈"]
    ]

    bm25 = BM25()
    bm25.fit(docs)

    # 查询（已分词）query = ["苹果", "手机"]
    results = bm25.search(query, docs)

    print("检索结果排序：")
    for idx, score in results:
        print(f"文档 {idx+1}: {docs[idx]} 得分: {score:.4f}")