オープンソース · Apache-2.0 ライセンス

KUAKE-LR-Rerank

sklearn LogisticRegression による軽量な Pointwise Rerank(リランキングモデル)。CBLUE ベンチマークの KUAKE-QTR(クエリ・タイトル関連度)と KUAKE-IR(医学段落検索)データセットで学習・評価・推論を完了し、コードとモデル重みを完全公開しています。

Rerank とは?

検索エンジンが一次召回した結果を再ソートし、最も関連するものを先頭に配置します

第 1 段階:検索・召回

大規模な文書群から関連しそうな候補を素早く抽出します。例えば 95 万件の医学段落から TF-IDF コサイン類似度で Top-200 候補を取得します。

第 2 段階:精密リランク

候補を 1 件ずつ精密にスコアリングし、最も関連するものを先頭に配置します。本プロジェクトは各 query-document ペアを独立に採点し、スコアの高い順に並べます。

わかりやすい例え:図書館で本を探すときと同じです。第 1 段階の「検索」は司書が本棚から関連しそうな 200 冊を素早く選ぶこと、第 2 段階の「Rerank」はあなたがタイトルと要約を見て最も関連する 10 冊を机に置くことです。

Rerank 手法の比較

本プロジェクトは Pointwise 手法を採用し、効果とエンジニアリングコストを両立

手法考え方計算量代表例
Pointwise各 query-doc ペアを独立に採点低Logistic Regression、BERT シングルタワー
Pairwise文書ペアの相対順序を比較中RankNet、LambdaRank
Listwiseランキングリスト全体の指標を直接最適化高ListNet、LambdaMART

Pipeline アーキテクチャ

sklearn FeatureUnion で 4 組の特徴量を並列計算し、連結して LogisticRegression に投入

              ┌──────────────┐
              │ query + title│
              └──────┬───────┘
         ┌───────────┼───────────────────┐
         ▼           ▼           ▼       ▼
   ┌─────────┐ ┌─────────┐ ┌─────────┐ ┌──────────┐
   │Manual   │ │Query    │ │Title    │ │Concat    │
   │(12-d)   │ │TF-IDF   │ │TF-IDF   │ │TF-IDF    │
   └────┬────┘ └────┬────┘ └────┬────┘ └────┬─────┘
        └───────────┼───────────┼───────────┘
                    ▼
           ┌────────────────┐
           │  FeatureUnion  │
           │ (concat ~15k-d) │
           └───────┬────────┘
                   ▼
          ┌────────────────┐
          │ LogisticRegr.  │
          │ Softmax 4-cls  │
          └───────┬────────┘
                   ▼
          ┌────────────────┐
          │ Rerank Score   │
          │ P(label≥2)     │
          └────────────────┘

特徴量の並列計算

FeatureUnion により手動特徴量と 3 系統の TF-IDF 特徴量を並列計算し、直列処理のオーバーヘッドを回避。最終的に約 15,000 次元の疎ベクトルに連結します。

Softmax 4 分類

LogisticRegression は multinomial softmax を採用し、0/1/2/3 の 4 段階関連度の確率分布を出力します。

ランキングスコア

P(label>=2) を最終的な Rerank スコアとして採用、すなわち「やや関連以上」の累積確率です。

12 次元の手動特徴量

単語重複、文字列類似度、長さ特徴量、古典的 IR スコアを融合

#特徴量説明
1overlap_ratioquery と title の共通 token 数 / query token 数
2jaccard共通 token 数 / 和集合 token 数
3q_hit_titlequery が title にヒットした token の割合
4t_hit_querytitle が query にヒットした token の割合
5edit_dist_norm正規化編集距離(Levenshtein)
6lcs_len最長共通部分文字列の長さ
7lcs_ratio最長共通部分文字列 / min(len1, len2)
8digit_match一致する数字の個数(用量・時間などの一致性)
9len_ratiotitle_len / query_len
10len_diff絶対長の差
11query_lenquery の文字数
12bm25_score軽量 BM25 スコア(古典的 IR 関連度)
BM25 スコアは特徴量重要度で第 2 位(重み 1.192)となり、古典的 IR 特徴量がランキングタスクで依然不可欠であることを裏付けています。

IR 二段階検索フロー

95 万段落から効率的に検索・精密ランキング

段落庫 (958,846件) --> TF-IDF ベクトル化 --> コサイン類似度 --> Top-200 --> LR Rerank --> Top-10

第 1 段階(一次選別)

96 万段落を TfidfVectorizer で疎行列にベクトル化し、各 query について疎行列乗算でコサイン類似度を計算して Top-200 候補を取得します。

第 2 段階(Rerank)

Top-200 候補について特徴ベクトル(12 次元手動 + TF-IDF)を構築し、LR モデルで P(label>=2) をランキングスコアとして予測し、Top-10 を出力します。

エンジニアリング最適化:一次選別の TF-IDF 行列とベクトル化器は models/ir_cache/ にキャッシュされます。初回実行は約 2-3 分、以降はキャッシュ読み込みで数秒です。

評価結果

CBLUE ベンチマーク検証セットでの実測性能

KUAKE-QTR(クエリ・タイトル関連度)

0.8689
NDCG@10
ランキング品質が高い
0.5613
Accuracy
4 分類、ランダムベースライン約 25%
0.6152
MRR
最初の関連文書は平均 2 位

KUAKE-IR(医学段落検索、1000 クエリ)

0.7158
NDCG@10
96 万段落から検索・ランキング
0.0176
MRR
各 query につき関連 doc は 1 件のみ注釈
0.0280
Recall@10
注釈が不完全なため、実際の召回はさらに多い可能性

特徴量重要度 Top-5

順位特徴量重み
1f1858(concat TF-IDF 特徴量)1.219
2f12(BM25 スコア)1.192
3f18600.969
4f11400.953
5f33550.926

データセット

CBLUE 中国語生物医学言語理解評価ベンチマークに基づく

データセット用途規模ラベル
KUAKE-QTR Query-Title 関連度分類 24,174 train / 2,913 dev / 5,465 test 4 分類(0 無関連 ~ 3 完全関連)
KUAKE-IR 医学段落検索 958,846 段落 / 1,000 クエリ注釈 query-doc 関連 / 無関連

クイックスタート

4 ステップで学習・評価・推論を完了

依存関係のインストール

Python >= 3.8。scikit-learn、pandas、numpy、jieba、python-Levenshtein などの依存関係をインストールします。

pip install scikit-learn pandas numpy jieba python-Levenshtein matplotlib seaborn joblib scipy

学習 + 評価

GroupKFold 3 分割交差検証(query でグループ化しリークを防止)+ GridSearchCV ハイパーパラメータ探索、計 48 回の学習。

python -c "import sys; sys.path.insert(0, '.'); from src.train_eval import main; main()"

KUAKE-QTR テスト推論

5,465 件の予測結果を出力。predict_label と predict_score を含みます。

python -c "import sys; sys.path.insert(0, '.'); from src.infer_qtr import main; main()"

KUAKE-IR 検索 + Rerank

1,000 query x 10 doc の計 10,000 行の検索ランキング結果を出力。全量検索には約 2-4 GB のメモリが必要です。

python -c "import sys; sys.path.insert(0, '.'); from src.infer_ir import main; main()"

モデル使用例

学習済み Pipeline を読み込むだけで直接推論できます

import joblib
import pandas as pd

# 加载模型
pipeline = joblib.load("models/lr_rerank.pkl")

# 单条预测
df = pd.DataFrame([{
    "query": "糖尿病饮食注意事项",
    "title": "糖尿病患者饮食指南",
    "label": -1  # dummy
}])

probs = pipeline.predict_proba(df)
score = probs[0, 2] + probs[0, 3]   # P(label≥2) 作为排序分
pred = pipeline.predict(df)[0]

print(f"预测标签: {pred}, 排序分数: {score:.4f}")

リソース

コード、モデル重み、技術ドキュメントをすべて公開

プロジェクト成果物

モデル、予測結果、可視化レポート一式

ファイル形式説明
models/lr_rerank.pklpickle完全な sklearn Pipeline、推論可能(1.7 GB、LFS 管理)
output/KUAKE-QTR_test_pred.jsonJSON5,465 件のテスト予測
output/KUAKE-IR_dev_pred.tsvTSV検索ランキング結果
output/eval_report.txtText完全な評価レポート
output/confusion_matrix.pngPNG混同行列ヒートマップ
output/feature_importance.pngPNG特徴量重要度 Top-30
output/lr_rerank_report.htmlHTML全工程テクニカルレポート(SVG 図解付き)

さらなる最適化

Pointwise から意味的特徴量への進化の道筋

Pairwise 学習

LambdaRank でランキング指標を直接最適化し、Pointwise 分類を置き換えます。

意味的特徴量

事前学習済み言語モデル(BERT など)を導入し、TF-IDF に代わる深い意味的特徴量を抽出します。

ドメイン特徴量

医学エンティティマッチング、同義語拡張などの医学ドメイン特徴量を追加します。

より大きな候補集合

一次選別を Top-200 から Top-500 または Top-1000 に拡張します。

オープンソースで共に発展

本プロジェクトは Apache-2.0 で公開されています。Issue と PR を歓迎し、中国語医学検索とリランキング技術の発展に共に取り組みましょう。