🗂

PythonでBM25のスコアを算出してみる

2024/07/31に公開

やること

フルテキスト検索で使われるBM25のスコアをPythonで求める

背景

AI Searchには様々なドキュメント検索の手法があり、代表的なものとしてフルテキスト検索があります。フルテキスト検索ではtf-idfを発展させたBM25というアルゴリズムが使われています。

参考記事

マイクロソフトの公式に説明があります。
https://learn.microsoft.com/ja-jp/azure/search/index-ranking-similarity

理論

BM25の数式は一見難解ですが以下の記事がわかりやすかったです。
https://qiita.com/KokiSakano/items/2a0f4c45caaa09cf1ab9#:~:text=BM25は特に検索アルゴリズム,があるというところである。

実装

今回はシンプルに実装してスコアを算出してみましょう。

from rank_bm25 import BM25Okapi
import numpy as np

# コーパスとクエリの定義
corpus = [
    "hello world hello",
    "hello good morning",
    "hello world",
    "python BM25 implementation"
]
query = "hello world"

# テキストをトークン化
def tokenize(text):
    return text.split()

tokenized_corpus = [tokenize(doc) for doc in corpus]
tokenized_query = tokenize(query)

# BM25の実装
def compute_bm25(corpus, query, k1, b):
    bm25 = BM25Okapi(corpus, k1=k1, b=b)
    scores = bm25.get_scores(query)
    return scores

# 初期パラメータ
k1 = 1.2
b = 0.75

# BM25スコアの計算
bm25_scores = compute_bm25(tokenized_corpus, tokenized_query, k1, b)

print("BM25 Scores with k1={} and b={}:".format(k1, b))
print(bm25_scores)

実行結果
BM25 Scores with k1=1.2 and b=0.75:
[0.1622842 0.11670238 0.13624324 0. ]

コメントなど

冒頭の記事を参考にBM25のパラメータはk=1.2、b=0.75が既定値ということで計算してみました。結果はクエリと一致する"hello world"ではなく、"hello world hello"がなぜか最大のスコアになりましたが、ちょっと理由はわかりませんでした。
 今回はBM25を試しましたが、ベクトル検索やセマンティック検索にもそれぞれスコアを算出する理論があるので、その辺もしっかり押さえていきたいと感じた次第です。

ヘッドウォータース

Discussion