sklearn LogisticRegression による軽量な Pointwise Rerank(リランキングモデル)。CBLUE ベンチマークの KUAKE-QTR(クエリ・タイトル関連度)と KUAKE-IR(医学段落検索)データセットで学習・評価・推論を完了し、コードとモデル重みを完全公開しています。
検索エンジンが一次召回した結果を再ソートし、最も関連するものを先頭に配置します
大規模な文書群から関連しそうな候補を素早く抽出します。例えば 95 万件の医学段落から TF-IDF コサイン類似度で Top-200 候補を取得します。
候補を 1 件ずつ精密にスコアリングし、最も関連するものを先頭に配置します。本プロジェクトは各 query-document ペアを独立に採点し、スコアの高い順に並べます。
本プロジェクトは Pointwise 手法を採用し、効果とエンジニアリングコストを両立
| 手法 | 考え方 | 計算量 | 代表例 |
|---|---|---|---|
| Pointwise | 各 query-doc ペアを独立に採点 | 低 | Logistic Regression、BERT シングルタワー |
| Pairwise | 文書ペアの相対順序を比較 | 中 | RankNet、LambdaRank |
| Listwise | ランキングリスト全体の指標を直接最適化 | 高 | ListNet、LambdaMART |
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 次元の疎ベクトルに連結します。
LogisticRegression は multinomial softmax を採用し、0/1/2/3 の 4 段階関連度の確率分布を出力します。
P(label>=2) を最終的な Rerank スコアとして採用、すなわち「やや関連以上」の累積確率です。
単語重複、文字列類似度、長さ特徴量、古典的 IR スコアを融合
| # | 特徴量 | 説明 |
|---|---|---|
| 1 | overlap_ratio | query と title の共通 token 数 / query token 数 |
| 2 | jaccard | 共通 token 数 / 和集合 token 数 |
| 3 | q_hit_title | query が title にヒットした token の割合 |
| 4 | t_hit_query | title が query にヒットした token の割合 |
| 5 | edit_dist_norm | 正規化編集距離(Levenshtein) |
| 6 | lcs_len | 最長共通部分文字列の長さ |
| 7 | lcs_ratio | 最長共通部分文字列 / min(len1, len2) |
| 8 | digit_match | 一致する数字の個数(用量・時間などの一致性) |
| 9 | len_ratio | title_len / query_len |
| 10 | len_diff | 絶対長の差 |
| 11 | query_len | query の文字数 |
| 12 | bm25_score | 軽量 BM25 スコア(古典的 IR 関連度) |
95 万段落から効率的に検索・精密ランキング
段落庫 (958,846件) --> TF-IDF ベクトル化 --> コサイン類似度 --> Top-200 --> LR Rerank --> Top-10
96 万段落を TfidfVectorizer で疎行列にベクトル化し、各 query について疎行列乗算でコサイン類似度を計算して Top-200 候補を取得します。
Top-200 候補について特徴ベクトル(12 次元手動 + TF-IDF)を構築し、LR モデルで P(label>=2) をランキングスコアとして予測し、Top-10 を出力します。
CBLUE ベンチマーク検証セットでの実測性能
| 順位 | 特徴量 | 重み |
|---|---|---|
| 1 | f1858(concat TF-IDF 特徴量) | 1.219 |
| 2 | f12(BM25 スコア) | 1.192 |
| 3 | f1860 | 0.969 |
| 4 | f1140 | 0.953 |
| 5 | f3355 | 0.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()"
5,465 件の予測結果を出力。predict_label と predict_score を含みます。
python -c "import sys; sys.path.insert(0, '.'); from src.infer_qtr import main; main()"
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}")
コード、モデル重み、技術ドキュメントをすべて公開
完全なソースコード、学習スクリプト、評価レポート。Issue と PR を歓迎します。
見る学習済みモデル重みは Hugging Face で公開されており、直接ダウンロードして再利用できます。
見るRerank の概念図解、Pipeline アーキテクチャ図、12 次元特徴量の詳細を含む全工程レポート。
モデル、予測結果、可視化レポート一式
| ファイル | 形式 | 説明 |
|---|---|---|
| models/lr_rerank.pkl | pickle | 完全な sklearn Pipeline、推論可能(1.7 GB、LFS 管理) |
| output/KUAKE-QTR_test_pred.json | JSON | 5,465 件のテスト予測 |
| output/KUAKE-IR_dev_pred.tsv | TSV | 検索ランキング結果 |
| output/eval_report.txt | Text | 完全な評価レポート |
| output/confusion_matrix.png | PNG | 混同行列ヒートマップ |
| output/feature_importance.png | PNG | 特徴量重要度 Top-30 |
| output/lr_rerank_report.html | HTML | 全工程テクニカルレポート(SVG 図解付き) |
Pointwise から意味的特徴量への進化の道筋
LambdaRank でランキング指標を直接最適化し、Pointwise 分類を置き換えます。
事前学習済み言語モデル(BERT など)を導入し、TF-IDF に代わる深い意味的特徴量を抽出します。
医学エンティティマッチング、同義語拡張などの医学ドメイン特徴量を追加します。
一次選別を Top-200 から Top-500 または Top-1000 に拡張します。
本プロジェクトは Apache-2.0 で公開されています。Issue と PR を歓迎し、中国語医学検索とリランキング技術の発展に共に取り組みましょう。