KUAKE-LR-Rerank

KUAKE-LR-Rerank

轻量级医学文本 Rerank 模型，基于 sklearn LogisticRegression 实现 Pointwise Rerank，在 KUAKE-QTR（查询-Title 相关性）和 KUAKE-IR（医学段落检索）两个 CBLUE 基准任务上验证。

A lightweight medical text Rerank model built with sklearn LogisticRegression (Pointwise Rerank), evaluated on the CBLUE benchmark's KUAKE-QTR (query-title relevance) and KUAKE-IR (medical passage retrieval) tasks.

模型简介 / Overview

• 架构 / Architecture： sklearn Pipeline + FeatureUnion（12 维手工特征 + 3 路 TF-IDF）
• 训练算法 / Algorithm： Multinomial LogisticRegression（4 类：0/1/2/3，P(label ≥ 2) 作为排序分数 / ranking score）
• 模型大小 / Size： ~1.7 GB（主要由 TF-IDF 词典矩阵构成 / TF-IDF vocabulary matrices）
• 环境依赖 / Dependency： 纯 CPU 可跑，无需 GPU / CPU-only, no GPU required

特征构成 / Feature Composition

性能指标 / Performance

KUAKE-QTR（查询-title 相关性 / query-title relevance）

• Accuracy: 56.1%
• NDCG@10: 0.8689
• MRR: 0.6152

KUAKE-IR（医学段落检索 / medical passage retrieval）

采用 TF-IDF 初筛（top-200）→ LR Rerank 两阶段流水线，全量统计评估 / Two-stage pipeline (TF-IDF top-200 → LR Rerank), full-collection evaluation:

（具体数值请运行 infer_ir.py 查看输出 / Run infer_ir.py for exact values）

文件说明 / Files

使用示例 / Usage

import joblib
import numpy as np

# 加载模型 / load model
pipe = joblib.load("lr_rerank.pkl")

# 输入 / inputs
queries = ["上呼吸道感染吃什么药"]
titles  = ["上呼吸道感染的治疗药物"]

# 预测相关性概率 / predict relevance prob (label >= 2)
probs = pipe.predict_proba(list(zip(queries, titles)))
score = probs[:, 2] + probs[:, 3]  # P(label >= 2)
print(f"相关性分数 / Relevance score: {score[0]:.4f}")

项目仓库 / Project Repository

完整代码、数据集、报告位于 / Full code, data, and reports at: https://github.com/githuBlijingai/KUAKE-LR-Rerank

License

Apache 2.0