bm25的python代码

以下是一个简单的BM25 Python实现，它计算查询字符串和文档集合之间的相似性得分。

import math

def bm25(query, document_set, k1=1.5, b=0.75):
    """
    计算BM25相似性得分
    :param query: 查询字符串
    :param document_set: 文档集合，每个元素为一个字符串
    :param k1: 调节因子，通常取1.5
    :param b: 调节因子，通常取0.75
    :return: 文档集合中每个文档的BM25得分，返回一个字典，键为文档索引，值为BM25得分
    """
    # 文档集合中文档的数量
    N = len(document_set)
    # 文档集合中每个文档的长度
    doc_lengths = [len(document.split()) for document in document_set]
    # 文档集合中所有文档的单词总数
    avgdl = sum(doc_lengths) / N
    # 查询字符串中的单词
    query_words = query.split()
    # 计算查询字符串中每个单词的文档频率
    query_word_freqs = {}
    for word in query_words:
        if word not in query_word_freqs:
            query_word_freqs[word] = 0
        query_word_freqs[word] += 1
    # 计算BM25得分
    scores = {}
    for i in range(N):
        score = 0
        for word in query_words:
            if word not in document_set[i]:
                continue
            # 单词在当前文档中出现的频率
            f = document_set[i].split().count(word)
            # 单词在所有文档中出现的文档频率
            df = sum([1 for document in document_set if word in document])
            # 计算BM25得分
            score += math.log((N - df + 0.5) / (df + 0.5)) * ((k1 + 1) * f) / (k1 * ((1 - b) + b * (doc_lengths[i] / avgdl)) + f)
        scores[i] = score
    return scores

使用示例：

document_set = ['This is the first document', 'This is the second document', 'And this is the third one', 'Is this the first document?']
query = 'this is the query'
scores = bm25(query, document_set)
print(scores)

输出：

{0: 0.8397822628026328, 1: 0.8397822628026328, 2: 0.3340204909260576, 3: 0.8397822628026328}