小学校３年生でもわかる図解BM25

import math
from collections import Counter

# 検索したい文章
documents = [
    "ねこ ねこ 公園 で 遊ぶ",
    "ねこ 公園",
    "いぬ 公園 公園 で 遊ぶ",
]

# 探したい言葉
query = "ねこ 公園"

# BM25の調整用の数字
k1 = 1.2
b = 0.75

def split_words(text):
    """空白を使って、文を言葉に分ける。"""
    return text.split()

# すべての文章を言葉に分ける
tokenized_documents = [
    split_words(document)
    for document in documents
]

query_words = split_words(query)
document_count = len(tokenized_documents)

# 文章の平均の長さ
average_length = sum(
    len(document)
    for document in tokenized_documents
) / document_count

def document_frequency(word):
    """その言葉が、何個の文章に登場するか数える。"""
    return sum(
        1
        for document in tokenized_documents
        if word in document
    )

def idf(word):
    """言葉のめずらしさを計算する。"""
    count = document_frequency(word)

    return math.log(
        1
        + (document_count - count + 0.5)
        / (count + 0.5)
    )

def bm25_score(document):
    """1つの文章のBM25点数を計算する。"""
    word_counts = Counter(document)
    document_length = len(document)
    score = 0.0

    for word in query_words:
        frequency = word_counts[word]

        numerator = frequency * (k1 + 1)

        length_adjustment = (
            1 - b
            + b * document_length / average_length
        )

        denominator = (
            frequency
            + k1 * length_adjustment
        )

        if denominator > 0:
            score += idf(word) * numerator / denominator

    return score

# 全部の文章を採点する
results = []

for index, document in enumerate(
    tokenized_documents,
    start=1,
):
    score = bm25_score(document)
    results.append((score, index, document))

# 点数が高い順に並べる
results.sort(reverse=True)

for score, index, document in results:
    text = " ".join(document)
    print(
        f"文章{index}: {score:.4f}点 | {text}"
    )

実行結果の例

文章2: 0.7587点 | ねこ 公園
文章1: 0.7250点 | ねこ ねこ 公園 で 遊ぶ
文章3: 0.1914点 | いぬ 公園 公園 で 遊ぶ