BM25向量数据库实战：从原理到Python实现

1次阅读

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

在搜索引擎和推荐系统的核心技术中，BM25 算法一直扮演着重要角色。作为一种基于概率模型的检索算法，BM25 通过考虑词频 (TF) 和逆文档频率 (IDF) 来计算文档与查询的相关性得分。与传统的关键词匹配相比，BM25 能更准确地反映词语在文档中的重要性。

但随着数据量的爆炸式增长和用户对语义搜索需求的提升，单纯的 BM25 算法开始面临挑战：

无法有效处理同义词和语义相关性
难以适应多模态数据检索
扩展性受限

将 BM25 与向量数据库结合，可以充分发挥两者的优势：BM25 提供精确的词汇级匹配，而向量数据库则支持高效的相似度搜索和语义理解。

在实现文本检索系统时，我们通常面临两种主要技术路线的选择：

传统倒排索引
优点：精确匹配效率高，实现简单
缺点：缺乏语义理解，扩展性差
典型应用：Elasticsearch, Solr
向量化检索
优点：支持语义搜索，易于扩展
缺点：计算成本高，需要预处理
典型应用：FAISS, Milvus

实际测试数据表明(基于 MS MARCO 数据集)：

指标	BM25+ 倒排索引	BM25+ 向量数据库
准确率 @10	0.32	0.41
召回率	0.58	0.67
查询延迟(ms)	12	25

import math
from collections import defaultdict

class BM25:
    def __init__(self, k1=1.5, b=0.75):
        self.k1 = k1
        self.b = b
        self.documents = []
        self.avgdl = 0
        self.df = defaultdict(int)
        self.idf = {}

    def add_document(self, document):
        """添加文档到索引"""
        self.documents.append(document)
        self.avgdl = sum(len(d) for d in self.documents) / len(self.documents)

        # 更新词频统计
        for word in set(document):
            self.df[word] += 1

    def calculate_idf(self):
        """计算逆文档频率"""
        N = len(self.documents)
        for word, freq in self.df.items():
            self.idf[word] = math.log((N - freq + 0.5) / (freq + 0.5) + 1)

    def get_score(self, query, document):
        """计算查询与文档的相关性得分"""
        score = 0.0
        doc_len = len(document)

        for word in query:
            if word not in document:
                continue

            tf = document.count(word)
            # BM25 公式
            numerator = self.idf[word] * tf * (self.k1 + 1)
            denominator = tf + self.k1 * (1 - self.b + self.b * doc_len / self.avgdl)
            score += numerator / denominator

        return score

import numpy as np
from pymilvus import connections, Collection

class VectorSearch:
    def __init__(self, host='localhost'):
        connections.connect(host=host)
        self.collection = Collection("bm25_vectors")  # 假设已创建集合

    def search(self, query_vector, top_k=10):
        """在向量数据库中执行近似最近邻搜索"""
        search_params = {
            "metric_type": "L2",
            "params": {"nprobe": 10}
        }

        results = self.collection.search(data=[query_vector],
            anns_field="embedding",
            param=search_params,
            limit=top_k,
            output_fields=["doc_id", "bm25_score"]
        )

        return results[0]

def hybrid_search(query, bm25_model, vector_db):
    """混合检索流程"""
    # 步骤 1: BM25 初步筛选
    bm25_scores = []
    for doc in bm25_model.documents:
        score = bm25_model.get_score(query, doc)
        bm25_scores.append((doc['id'], score))

    # 取 Top 1000 初步结果
    bm25_scores.sort(key=lambda x: x[1], reverse=True)
    candidate_ids = [doc_id for doc_id, _ in bm25_scores[:1000]]

    # 步骤 2: 向量精排
    query_vector = get_embedding(query)  # 假设已有嵌入函数
    vector_results = vector_db.search(query_vector, top_k=100)

    # 步骤 3: 结果融合
    final_results = []
    for hit in vector_results:
        doc_id = hit.entity.get('doc_id')
        bm25_score = hit.entity.get('bm25_score', 0)
        vector_score = hit.score

        # 组合评分(可调整权重)
        combined_score = 0.7 * vector_score + 0.3 * bm25_score
        final_results.append({
            'doc_id': doc_id,
            'score': combined_score
        })

    return sorted(final_results, key=lambda x: x['score'], reverse=True)

在 100 万文档规模下的实测数据：