基于 sklearn LogisticRegression 实现的轻量化 Pointwise Rerank（重排序模型），在 CBLUE 基准的 KUAKE-QTR（查询-标题相关性）与 KUAKE-IR（医学段落检索）数据集上完成训练、评估与推理，代码与模型权重全部开源。
对搜索引擎初步召回的结果进行二次排序，把最相关的排到最前面
从海量文档库中快速挑出可能相关的候选集。例如从 95 万篇医学段落中，通过 TF-IDF 余弦相似度取出 Top-200 候选。
对候选集逐条精细打分，把最相关的排到最前面。本项目对每个 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 让手动特征与三路 TF-IDF 特征并行计算，避免串行开销，最终拼接为约 15000 维稀疏向量。
LogisticRegression 采用 multinomial softmax，输出 0/1/2/3 四级相关性概率分布。
取 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 相关 / 不相关 |
四步完成训练、评估与推理
Python ≥ 3.8，安装 scikit-learn、pandas、numpy、jieba、python-Levenshtein 等依赖。
pip install scikit-learn pandas numpy jieba python-Levenshtein matplotlib seaborn joblib scipy
GroupKFold 三折交叉验证（按 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 × 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。