ranxを使って検索システムのオフライン評価をする

ranxはPythonで実装されたランキング評価用のライブラリです。検索システムの評価に使える以下のような指標が用意されていたり、それらの結果をRRFのようなアルゴリズムで組み合わせたり、論文用にテーブルをエクスポートする機能を備えています。本記事では、ranxの簡単な使い方を紹介します。

インストール

pipを使っている場合は以下のコマンドを実行してインストールします。

pip install ranx

基本的な使い方

基本的には、以下に示すようにQrelsとRunのインスタンスを作成します。Qrelsは、評価用の正解データを格納するためのクラスです。一方、Runは評価対象のモデルによって推定された関連性スコアを格納するためのクラスです。ここではどちらもPythonの辞書から作成していますが、PandasのDataFrameやファイルから作成することもできます。

from ranx import Qrels, Run

# クエリ`q_1`の正解は`d_12`と`d_25`で、それぞれの関連性スコアは5と3
qrels_dict = {
    "q_1": { "d_12": 5, "d_25": 3 },
    "q_2": { "d_11": 6, "d_22": 1 }
}

# モデルによる推定結果をクエリの場合と同様に格納
run_dict = {
    "q_1": {
        "d_12": 0.9, "d_23": 0.8, "d_25": 0.7,
        "d_36": 0.6, "d_32": 0.5, "d_35": 0.4
    },
    "q_2": {
        "d_12": 0.9, "d_11": 0.8, "d_25": 0.7,
        "d_36": 0.6, "d_22": 0.5, "d_35": 0.4
    }
}

qrels = Qrels(qrels_dict)
run = Run(run_dict)

QrelsとRunを用意できたら、evaluate関数を使って評価をします。evaluate関数は、評価用の正解データとモデルの推定結果および評価指標を受け取り、指定した評価指標のスコアを返します。評価指標は文字列か文字列のリストで指定することができ、後者であれば複数の評価指標のスコアを一度に計算できます。

from ranx import evaluate

# 1つの指標に対してスコアを計算する
evaluate(qrels, run, "ndcg@5")
>>> 0.7861

# 複数の指標のスコアを一度に計算する
evaluate(qrels, run, ["map@5", "mrr"])
>>> {"map@5": 0.6416, "mrr": 0.75}

複数の異なるモデルや同じモデルのパラメーターを変えて実行した結果を比較するにはcompare関数を使います。compare関数は評価用の正解データと複数の推定結果、評価指標およびしきい値を受け取り、比較結果を表すReportのインスタンスを返します。Reportのメソッドを呼び出すことで、比較結果をLatex形式や辞書に変換できます。

from ranx import compare

# 異なる実行結果を比較し、両側t検定を実行します。
report = compare(
    qrels=qrels,
    runs=[run_1, run_2, run_3, run_4, run_5],
    metrics=["map@100", "mrr@100", "ndcg@10"],
    max_p=0.01
)

実行結果は以下のように表示されます。

#    Model    MAP@100    MRR@100    NDCG@10
---  -------  --------   --------   ---------
a    model_1  0.320ᵇ     0.320ᵇ     0.368ᵇᶜ
b    model_2  0.233      0.234      0.239
c    model_3  0.308ᵇ     0.309ᵇ     0.330ᵇ
d    model_4  0.366ᵃᵇᶜ   0.367ᵃᵇᶜ   0.408ᵃᵇᶜ
e    model_5  0.405ᵃᵇᶜᵈ  0.406ᵃᵇᶜᵈ  0.451ᵃᵇᶜᵈ

モデルの推定結果を組み合わせるにはfuse関数を使います。fuse関数は複数の推定結果や統合方法を受け取り、統合結果を表すRunのインスタンスを返します。統合方法によっては最適化対象となるハイパーパラメーターが存在します（RRFにおけるkなど）が、そのような場合はoptimize_fusion関数を使って最適なハイパーパラメーターを推定することもできます。

from ranx import fuse, optimize_fusion

best_params = optimize_fusion(
    qrels=train_qrels,
    runs=[train_run_1, train_run_2, train_run_3],
    norm="min-max",     # The norm. to apply before fusion
    method="wsum",      # The fusion algorithm to use (Weighted Sum)
    metric="ndcg@100",  # The metric to maximize
)

combined_test_run = fuse(
    runs=[test_run_1, test_run_2, test_run_3],  
    norm="min-max",       
    method="wsum",        
    params=best_params,
)

おわりに

本記事では検索システムの評価に使えるranxの基本的な使い方について紹介しました。最近はRAG（Retrieval Augmented Generation）を使ったシステムを作っていたのですが、そのコンポーネントである検索システムの評価をするためにライブラリを探していたところ、ranxを見つけました。RAGなどの作成を支援してくれるLlamaIndexのようなフレームワークには評価用の機能が含まれていることが多いですが、ranxを使うことで、フレームワークに依存しない評価を行うことができます。興味があれば試してみてください。

Ahogrammer

Deep Dive Into NLP, ML and Cloud

ranxを使って検索システムのオフライン評価をする

インストール

基本的な使い方

おわりに

参考資料