主流开源 Rerank 模型解析与选型指南（2026 版）

1

NDCG 是 “归一化折损累积增益”（Normalized Discounted Cumulative Gain）的缩写，是信息检索领域常用的评估指标。它的核心作用是衡量检索结果的相关性和排序质量：

1
2
3
4
5
6

# 核心依赖（适配所有模型）
pip install torch transformers sentence-transformers

# 专项依赖（按需安装）
pip install colbert-ai  # ColBERT v2专用
pip install accelerate  # 大模型推理加速（Qwen3-Reranker-4B）

  1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154

# 安装依赖：pip install transformers>=4.51.0 torch modelscope
import torch
from typing import List, Optional, Tuple
from modelscope import AutoTokenizer, AutoModelForCausalLM


class QwenReranker:
    """
    通义千问重排模型封装类

    基于ModelScope的Qwen3-Reranker系列模型，实现查询(Query)与文档(Document)的相关性评分及排序，
    模型通过判断文档是否满足查询需求，输出0-1之间的相关性得分（越接近1表示相关性越高）。

    基础用法示例：
        >>> ranker = QwenReranker("Qwen/Qwen3-Reranker-0.6B")
        >>> query = "人工智能在医疗领域的应用"
        >>> docs = ["AI辅助诊断系统可分析医学影像...", "人工智能在农业中的应用..."]
        >>> sorted_results = ranker.rerank_sorted(query, docs)
        >>> print(sorted_results)  # [(doc索引, 得分), ...]（按得分降序）
    """

    DEFAULT_INSTRUCTION = "Given a web search query, retrieve relevant passages that answer the query"

    def __init__(
            self,
            model_name_or_path: str = "Qwen/Qwen3-Reranker-0.6B",
            device: Optional[str] = None,
            max_length: int = 8192,
    ):
        """
        初始化重排模型

        Args:
            model_name_or_path: 模型名称或本地路径。
                1.若为模型名称：如果本地缓存（默认用户路径 ~/.cache/modelscope/hub ）中没有模型文件，则 from_pretrained() 会自动
                从modelscope.cn下载模型权重、配置和分词器文件。 如果本地缓存已存在该模型，代码会直接加载缓存中的文件，无需重新下载。
                2.若为本地模型的存储路径：则需要手动下载模型文件到路径地址（config.json、*.safetensors 等）.
            device: 运行设备（"cuda"/"cpu"），默认自动检测
            max_length: 默认最大文本长度
        """
        self.tokenizer = AutoTokenizer.from_pretrained(model_name_or_path, padding_side="left")
        self.max_length = max_length

        # 设备自动检测
        self.device = device if device else ("cuda" if torch.cuda.is_available() else "cpu")
        # 加载模型

        self.model = AutoModelForCausalLM.from_pretrained(model_name_or_path).eval()
        # We recommend enabling flash_attention_2 for better acceleration and memory saving.
        # model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen3-Reranker-0.6B", torch_dtype=torch.float16,
        # attn_implementation="flash_attention_2").cuda().eval()

        if device and device != "auto":  # 指定单卡
            self.model = self.model.to(device)

        # 特殊 token id
        self.token_true_id = self.tokenizer.convert_tokens_to_ids("yes")
        self.token_false_id = self.tokenizer.convert_tokens_to_ids("no")

        # prompt 模板
        self.prefix = (
            "<|im_start|>system\nJudge whether the Document meets the requirements based on the Query and the "
            "Instruct provided. Note that the answer can only be \"yes\" or "
            "\"no\".<|im_end|>\n<|im_start|>user\n")
        self.suffix = "<|im_end|>\n<|im_start|>assistant\n<think>\n\n</think>\n\n"
        self.prefix_tokens = self.tokenizer.encode(self.prefix, add_special_tokens=False)
        self.suffix_tokens = self.tokenizer.encode(self.suffix, add_special_tokens=False)

    @staticmethod
    def _format(instruction: Optional[str], query: str, doc: str) -> str:
        inst = instruction or QwenReranker.DEFAULT_INSTRUCTION
        return f"<Instruct>: {inst}\n<Query>: {query}\n<Document>: {doc}"

    def _tokenize(self, texts: List[str]):
        inputs = self.tokenizer(
            texts,
            padding=False,
            truncation="longest_first",
            return_attention_mask=False,
            max_length=self.max_length - len(self.prefix_tokens) - len(self.suffix_tokens),
        )
        # 拼接 prefix + text + suffix
        for idx, ids in enumerate(inputs["input_ids"]):
            inputs["input_ids"][idx] = self.prefix_tokens + ids + self.suffix_tokens
        # pad 到 batch
        inputs = self.tokenizer.pad(inputs, padding=True, return_tensors="pt", max_length=self.max_length)
        for key in inputs:
            inputs[key] = inputs[key].to(self.model.device)
        return inputs

    @torch.no_grad()
    def _compute_logits(self, inputs) -> List[float]:
        logits = self.model(**inputs).logits[:, -1, :]
        true_score = logits[:, self.token_true_id]
        false_score = logits[:, self.token_false_id]
        batch_scores = torch.stack([false_score, true_score], dim=1)
        batch_scores = torch.nn.functional.log_softmax(batch_scores, dim=1)
        scores = batch_scores[:, 1].exp().tolist()
        return scores

    def rerank_sorted(
            self,
            query: str,
            docs: List[str],
            instruction: Optional[str] = None,
            batch_size: int = 8,
    ) -> List[Tuple[int, float]]:
        """
        对单 query 多 doc 的场景，返回 [(doc_idx, score), ...] 按得分降序
        """
        queries = [query] * len(docs)
        pairs = [self._format(instruction, q, d) for q, d in zip(queries, docs)]

        scores = []
        for i in range(0, len(pairs), batch_size):
            batch_pairs = pairs[i: i + batch_size]
            inputs = self._tokenize(batch_pairs)
            scores.extend(self._compute_logits(inputs))

        return sorted(enumerate(scores), key=lambda x: x[1], reverse=True)


# ---------------- 使用示例 ----------------
if __name__ == "__main__":
    # 1. 初始化模型（CPU 也能跑，GPU 更快）
    # model_name_or_path: 模型名称或本地路径。
    # -若为模型名称：如果本地缓存（默认路径~ /.cache/modelscope/）中没有模型文件，from_pretrained()
    #  会自动从modelscope.cn下载模型权重、配置和分词器文件。 如果本地缓存已存在该模型，代码会直接加载缓存中的文件，无需重新下载。
    # -若为本地模型的存储路径：则需要手动下载模型文件到路径地址（config.json、model.safetensors(或pytorch_model.bin)、tokenizer.json
    #  tokenizer_config.json、special_tokens_map.json、vocab.txt）.
    ranker = QwenReranker(
        # 由于本地硬件资源有限，替换为Qwen3-Reranker-0.6B模型.
        # model_name_or_path="Qwen/Qwen3-Reranker-4B",
        model_name_or_path="Qwen/Qwen3-Reranker-0.6B",
    )

    # 2. 测试查询+文档
    query = "人工智能在医疗领域的应用"
    docs = [
        "人工智能驱动的农业无人机可实现作物病虫害监测，精准率达85%以上",  # 无关（完全跨领域）
        "AI辅助诊断系统可快速分析CT、MRI等医学影像，使肺癌、乳腺癌等疾病的早期检出率提升30%以上",  # 高相关
        "深度学习算法在金融风控中的应用案例，有效降低信贷违约率15%",  # 无关（金融领域）
        "卷积神经网络在皮肤癌图像分类任务上达到91%的准确率，为基层医疗机构提供了低成本诊断方案",  # 中高相关（技术细分，间接相关）
        "智能模拟在生物医药研发中的应用，主要用于药物分子结构模拟，不属于临床医疗应用范畴",  # 干扰项（语义相近但领域不符）
        "人工智能技术在医疗领域的核心应用包括AI辅助诊断、个性化治疗方案生成、医疗影像分析等",  # 高相关
        "机器学习算法在医保风控系统中的应用，可识别虚假就医报销行为，属于AI在医疗管理的边缘场景",  # 低相关（跨领域AI，弱关联）
        "医疗大数据平台通过机器学习算法整合患者电子病历，为医院管理决策提供数据支撑",  # 中相关（泛医疗AI，关联性减弱）
    ]

    # 3. 自动排序
    print("\n【按得分降序】")
    sorted_res = ranker.rerank_sorted(query, docs)
    for rank, (idx, score) in enumerate(sorted_res, 1):
        print(f"{rank}. {score:.4f}\t{docs[idx]}")

  1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162

import logging
from typing import List, Tuple, Optional, Dict, Any

import torch
from sentence_transformers import CrossEncoder

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
)
logger = logging.getLogger(__name__)


class MxbaiReranker:
    """
    MixedBread AI Rerank 模型封装类
    支持多语言（德/英/西/法/中等）的文本重排
    """

    def __init__(
            self,
            model_name_or_path: str = "mixedbread-ai/mxbai-rerank-large-v2",
            max_length: int = 8192,
            device: Optional[str] = None,
            model_kwargs: Optional[Dict[str, Any]] = None
    ):
        """
        初始化重排模型

        Args:
            model_name_or_path: 模型名称或本地路径
            max_length: 文本最大长度
            device: 运行设备 (e.g., "cpu", "cuda", "cuda:0")
            model_kwargs: 传递给CrossEncoder的额外参数
        """
        self.model_name_or_path = model_name_or_path
        self.max_length = max_length
        # 设备自动检测
        self.device = device if device else ("cuda" if torch.cuda.is_available() else "cpu")
        self.model_kwargs = model_kwargs or {}

        # 加载模型
        self._load_model()

    def _load_model(self) -> None:
        """加载CrossEncoder模型"""
        try:
            logger.info(f"正在加载模型: {self.model_name_or_path}")
            self.model = CrossEncoder(
                model_name_or_path=self.model_name_or_path,
                max_length=self.max_length,
                device=self.device,
                **self.model_kwargs
            )
            logger.info("模型加载成功")
        except Exception as e:
            logger.error(f"模型加载失败: {e}")
            raise

    def rerank_sorted(
            self,
            query: str,
            passages: List[str],
            batch_size: int = 16,
            return_scores: bool = True
    ) -> List[Tuple[str, float]] | List[str]:
        """
        对段落进行重排

        Args:
            query: 多语言查询文本
            passages: 待重排的段落列表
            batch_size: 批量推理大小,可根据CPU/GPU调整，CPU建议设为8/16
            return_scores: 是否返回得分

        Returns:
            排序后的段落列表（可选包含得分）
        """
        if not query:
            raise ValueError("查询文本不能为空")
        if not passages:
            raise ValueError("段落列表不能为空")

        try:
            # 构造输入对
            inputs = [[query, passage] for passage in passages]

            # 批量推理
            logger.debug(f"开始推理，批次大小: {batch_size}")
            # 推理打分（batch_size可根据CPU/GPU调整，CPU建议设为8/16）
            scores = self.model.predict(
                inputs,
                batch_size=batch_size,
                show_progress_bar=False
            )

            # 排序
            ranked_pairs = sorted(
                zip(passages, scores),
                key=lambda x: x[1],
                reverse=True
            )

            logger.info(f"重排完成，共处理 {len(passages)} 个段落")

            if return_scores:
                return ranked_pairs
            else:
                return [pair[0] for pair in ranked_pairs]

        except Exception as e:
            logger.error(f"重排失败: {e}")
            raise

    def release(self):
        """显式释放模型资源（建议手动调用）"""
        try:
            if hasattr(self, 'model') and self.model is not None:
                del self.model
                if torch.cuda.is_available():
                    torch.cuda.empty_cache()
                logger.info("模型资源已释放")
        except Exception as e:
            logger.warning(f"释放模型资源时出错: {e}")


# 使用示例
if __name__ == "__main__":

    # 初始化重排器
    reranker = MxbaiReranker(
        # 模型名称或本地路径
        # 这里使用本地模型路径，从huggingface.co或modelscope.cn中下载模型(mxbai-rerank-large-v2)文件，存储本地路径下。
        # model_name_or_path="mixedbread-ai/mxbai-rerank-large-v2", # 模型名称,自动下载到~/.cache/huggingface/
        model_name_or_path="../../data/models_reranker_data/mixedbread-ai/mxbai-rerank-large-v2",  # 本地模型路径
        max_length=8192,
        # device="cuda"  # 如果有GPU可以指定
    )

    # 测试数据
    query = "人工智能在医疗领域的应用"
    docs = [
        "人工智能驱动的农业无人机可实现作物病虫害监测，精准率达85%以上-E",  # 无关（完全跨领域）
        "AI辅助诊断系统可快速分析CT、MRI等医学影像，使肺癌、乳腺癌等疾病的早期检出率提升30%以上-A",  # 高相关
        "深度学习算法在金融风控中的应用案例，有效降低信贷违约率15%-E",  # 无关（金融领域）
        "卷积神经网络在皮肤癌图像分类任务上达到91%的准确率，为基层医疗机构提供了低成本诊断方案-B",  # 中高相关（技术细分，间接相关）
        "智能模拟在生物医药研发中的应用，主要用于药物分子结构模拟，不属于临床医疗应用范畴-D",  # 干扰项（语义相近但领域不符）
        "人工智能技术在医疗领域的核心应用包括AI辅助诊断、个性化治疗方案生成、医疗影像分析等-A",  # 高相关
        "机器学习算法在医保风控系统中的应用，可识别虚假就医报销行为，属于AI在医疗管理的边缘场景-D",  # 低相关（跨领域AI，弱关联）
        "医疗大数据平台通过机器学习算法整合患者电子病历，为医院管理决策提供数据支撑-C",  # 中相关（泛医疗AI，关联性减弱）
    ]

    # 执行重排
    results = reranker.rerank_sorted(query, docs, batch_size=8)

    # 打印结果
    print(f"\nmxbai-rerank-large-v2 重排结果：")
    print(f"query:{query}")
    print("-" * 80)
    for i, (passage, score) in enumerate(results, 1):
        print(f"第{i}名：得分 {score:.4f} | 内容：{passage}")

1

从结果可见 mxbai-rerank-large-v2模型对中文的重排并不好。

  1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140

import os
from typing import List, Tuple, Optional

import torch
from FlagEmbedding import FlagReranker


class RerankerBgeV2:
    """
    BGE-Reranker-V2 重排模型处理器

    基于BAAI推出的bge-reranker-v2-m3模型实现文本相关性重排，用于提升检索系统中
    召回阶段候选文档的排序精度，核心通过计算query-doc对的语义相似度分数实现重排。
    适用于中文/多语言场景的检索后重排任务。
    """

    def __init__(
            self,
            model_path: str = "../../data/models_reranker_data/BAAI/bge-reranker-v2-m3",
            use_fp16: bool = True,
            device: str = "cpu"
    ):
        """
        初始化重排模型
            model_path: 模型本地存储路径，需包含完整的模型文件（config.json、model.safetensors等）
            use_fp16: 是否启用FP16精度推理，GPU环境下推荐True，CPU环境会自动强制关闭
            device: 模型运行设备，支持"cpu"、"cuda"、"cuda:0"等格式，为空时自动检测可用设备
        """
        # 转换为绝对路径，避免相对路径导致的文件查找失败
        self.model_path = os.path.abspath(model_path)
        # 设备自动检测
        self.device = device if device else ("cuda" if torch.cuda.is_available() else "cpu")
        # CPU环境下FP16无性能收益且可能导致兼容性问题，强制关闭
        self.use_fp16 = use_fp16 if device != "cpu" else False
        # 重排模型实例，延迟初始化
        self.reranker: Optional[FlagReranker] = None

        # 初始化模型
        self._init_model()

    def _init_model(self) -> None:
        """加载并重初始化重排模型"""
        try:
            self.reranker = FlagReranker(
                model_name_or_path=self.model_path,
                use_fp16=self.use_fp16,
                device=self.device
            )
            print(f"✅ 重排模型加载成功，设备：{self.device}，FP16：{self.use_fp16}")
        except Exception as e:
            raise RuntimeError(f"模型初始化失败：{e}")

    def rerank_sorted(
            self,
            query: str,
            candidate_docs: List[str],
            top_k: Optional[int] = None
    ) -> List[Tuple[str, float]]:
        """
        对候选文档列表按与查询语句的相关性进行重排

        Args:
            query: 用户查询语句（文本）
            candidate_docs: 待重排的候选文档列表
            top_k: 返回前k个结果（None返回全部）
        Returns:
           重排结果列表，元素为(文档文本, 相关性分数)，按分数降序排列
        """
        if not candidate_docs:
            raise ValueError("候选文档列表不能为空")

        if self.reranker is None:
            raise RuntimeError("模型未初始化，请检查模型路径和设备配置")

        # 构造query-doc文本对,格式：[[query1, doc1], [query1, doc2], ...]
        text_pairs = [[query, doc] for doc in candidate_docs]

        # 批量计算相关性分数
        try:
            scores = self.reranker.compute_score(text_pairs)
        except Exception as e:
            raise RuntimeError(f"分数计算失败：{e}")

        # 按分数降序排序
        ranked_results = sorted(
            zip(candidate_docs, scores),
            key=lambda x: x[1],
            reverse=True
        )

        # 截取top_k结果
        if top_k is not None and top_k > 0:
            ranked_results = ranked_results[:top_k]

        return ranked_results


# 示例：BGE-Reranker-V2模型的完整使用流程
if __name__ == "__main__":

    # 1. 初始化重排处理器
    # 需下载模型文件到存储路径../../data/models_reranker_data/BAAI，模型下载地址：
    # - HuggingFace: https://huggingface.co/BAAI/bge-reranker-v2-m3
    # - ModelScope: https://www.modelscope.cn/models/BAAI/bge-reranker-v2-m3
    # 需下载的核心文件：config.json、model.safetensors(或pytorch_model.bin)、 tokenizer.json、
    #  tokenizer_config.json、special_tokens_map.json、vocab.json。
    reranker = RerankerBgeV2(
        model_path="../../data/models_reranker_data/BAAI/bge-reranker-v2-m3",
        # device="cuda:0",  # 如需使用GPU可显式指定
        # use_fp16=True     # GPU环境下建议开启FP16
    )

    # 2. 准备测试数据
    query = "人工智能在医疗领域的应用"
    docs = [
        "人工智能驱动的农业无人机可实现作物病虫害监测，精准率达85%以上",  # 无关（完全跨领域）
        "AI辅助诊断系统可快速分析CT、MRI等医学影像，使肺癌、乳腺癌等疾病的早期检出率提升30%以上",  # 高相关
        "深度学习算法在金融风控中的应用案例，有效降低信贷违约率15%",  # 无关（金融领域）
        "卷积神经网络在皮肤癌图像分类任务上达到91%的准确率，为基层医疗机构提供了低成本诊断方案",  # 中高相关（技术细分，间接相关）
        "智能模拟在生物医药研发中的应用，主要用于药物分子结构模拟，不属于临床医疗应用范畴",  # 干扰项（语义相近但领域不符）
        "人工智能技术在医疗领域的核心应用包括AI辅助诊断、个性化治疗方案生成、医疗影像分析等",  # 高相关
        "机器学习算法在医保风控系统中的应用，可识别虚假就医报销行为，属于AI在医疗管理的边缘场景",  # 低相关（跨领域AI，弱关联）
        "医疗大数据平台通过机器学习算法整合患者电子病历，为医院管理决策提供数据支撑",  # 中相关（泛医疗AI，关联性减弱）
    ]

    # 3. 执行重排
    results = reranker.rerank_sorted(
        query=query,
        candidate_docs=docs,
        top_k=None  # 可改为None返回全部
    )

    # 4. 打印结果
    print("\n" + "=" * 50)
    print(f"bge-reranker-v2-m3 重排结果（分数越高越相关）：")
    print(f"查询语句：{query}")
    print("=" * 50)
    for idx, (doc, score) in enumerate(results, 1):
        print(f"Top{idx} | 分数：{score:.4f} | 文档：{doc}")
    print("=" * 50 + "\n")

1

从结果可见 bge-reranker-v2-m3模型对中文的重排相对较好。

  1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145

from typing import List, Tuple, Optional, Union

import torch
from sentence_transformers import CrossEncoder


class CrossEncoderReranker:

    def __init__(
            self,
            model_name_or_path: str = "../data/models_data/cross-encoder/ms-marco-MiniLM-L6-v2",
            device: Optional[str] = None,
            batch_size: int = 16
    ):
        """
        初始化重排器。
        基于CrossEncoder的文档重排器，用于检索系统中对召回阶段的候选文档进行语义相关性重排，提升检索精度。
        适配采用CrossEncoder架构系列模型, 如ms-marco-MiniLM-L-12-v2、bge-reranker-large、jina-rerank-v3、mxbai-rerank-large-v2等

        Args:
            model_name_or_path: str，模型名称（HuggingFace Hub规范名）或本地存储路径：
                - 1.模型名称：本地缓存（默认~/.cache/huggingface/）无该模型时，自动从Hub下载权重/配置/分词器；缓存已存在则直接加载，无需重复下载。
                - 2.本地路径：需手动下载完整模型文件（包含config.json、model.safetensors/pytorch_model.bin等）到本地路径地址。
            device: 模型运行设备，None则自动检测（优先使用CUDA，无则使用CPU）
            batch_size: 推理批次大小，建议CPU设8/16，GPU可根据显存适当增大（默认16）
        """
        # 设备自动适配
        self.device = device if device else ("cuda" if torch.cuda.is_available() else "cpu")
        self.batch_size = batch_size
        self.model_name_or_path = model_name_or_path
        self.model: Optional[CrossEncoder] = None

        # 初始化模型
        self._load_model()

    def _load_model(self) -> None:
        """内部方法：加载CrossEncoder模型"""
        try:
            self.model = CrossEncoder(self.model_name_or_path, device=self.device)
            print(f"✅重排模型加载完成 | 设备：{self.device} | 批次大小：{self.batch_size}")
        except Exception as e:
            raise RuntimeError(f"模型加载失败：{e}\n请检查：1. 模型路径是否正确 2. 网络是否正常（首次下载需联网）")

    def rerank_sorted(
            self,
            query: str,
            candidate_docs: List[str],
            return_scores: bool = True,
            top_k: Optional[int] = None
    ) -> Union[List[str], List[Tuple[str, float]]]:
        """
        对候选文档进行相关性重排

        Args:
            query: 用户查询语句
            candidate_docs: 召回阶段的候选文档列表（建议top-20以内，保证重排效率）
            return_scores: 是否返回相关性得分（得分越高，相关性越强）
            top_k: 返回前k个最相关文档，None返回全部

        Returns:
            重排后的文档列表（或文档+得分的元组列表），按相关性降序排列

        Raises:
            ValueError: 候选文档列表为空时抛出
            RuntimeError: 模型未成功加载时抛出
        """
        # 输入合法性校验
        if not candidate_docs:
            raise ValueError("候选文档列表不能为空")

        if self.model is None:
            raise RuntimeError("模型未初始化，请检查模型加载是否成功")

        # 构造query-doc配对（CrossEncoder必需输入格式）
        doc_pairs = [[query, doc] for doc in candidate_docs]

        # 批量推理计算相关性得分
        try:
            scores = self.model.predict(doc_pairs, batch_size=self.batch_size)
        except Exception as e:
            raise RuntimeError(f"重排打分失败：{e}")

        # 按得分降序排序
        sorted_indices = scores.argsort()[::-1]
        ranked_results = [
            (candidate_docs[i], float(scores[i])) for i in sorted_indices
        ]

        # 截取top_k结果
        if top_k is not None and top_k > 0:
            ranked_results = ranked_results[:top_k]

        # 适配返回格式
        if return_scores:
            return ranked_results
        else:
            return [item[0] for item in ranked_results]

    def release(self):
        """显式释放模型资源（建议手动调用）"""
        if self.model is not None:
            self.model = None
            if torch.cuda.is_available():
                torch.cuda.empty_cache()
            print("🔌 重排模型资源已释放")


# 示例使用
if __name__ == "__main__":
    # 1. 创建重排器实例
    reranker = CrossEncoderReranker(
        # model_name_or_path="cross-encoder/ms-marco-MiniLM-L6-v2",  # 使用HuggingFace远程加载模型
        # 手动从huggingface.co或modelscope.cn中下载模型文件，存储本地路径下
        model_name_or_path="../../data/models_reranker_data/cross-encoder/ms-marco-MiniLM-L6-v2",
        # device="cpu",  # 强制使用CPU
        batch_size=16
    )

    # 2. 准备测试数据
    query = "人工智能在医疗领域的应用"
    docs = [
        "人工智能驱动的农业无人机可实现作物病虫害监测，精准率达85%以上-E",  # 无关（完全跨领域）
        "AI辅助诊断系统可快速分析CT、MRI等医学影像，使肺癌、乳腺癌等疾病的早期检出率提升30%以上-A",  # 高相关
        "深度学习算法在金融风控中的应用案例，有效降低信贷违约率15%-E",  # 无关（金融领域）
        "卷积神经网络在皮肤癌图像分类任务上达到91%的准确率，为基层医疗机构提供了低成本诊断方案-B",  # 中高相关（技术细分，间接相关）
        "智能模拟在生物医药研发中的应用，主要用于药物分子结构模拟，不属于临床医疗应用范畴-D",  # 干扰项（语义相近但领域不符）
        "人工智能技术在医疗领域的核心应用包括AI辅助诊断、个性化治疗方案生成、医疗影像分析等-A",  # 高相关
        "机器学习算法在医保风控系统中的应用，可识别虚假就医报销行为，属于AI在医疗管理的边缘场景-D",  # 低相关（跨领域AI，弱关联）
        "医疗大数据平台通过机器学习算法整合患者电子病历，为医院管理决策提供数据支撑-C",  # 中相关（泛医疗AI，关联性减弱）
    ]

    # 3. 执行重排
    # 返回带得分的结果
    ranked_docs_with_scores = reranker.rerank_sorted(
        query=query,
        candidate_docs=docs,
        return_scores=True,
        top_k=None  # 可指定返回前N个，如top_k=2
    )

    # 4. 输出结果
    print("\n=== ms-marco-MiniLM-L6-v2 重排结果===")
    print(f"query:{query}")
    for idx, (doc, score) in enumerate(ranked_docs_with_scores, 1):
        print(f"TOP-{idx} | 得分：{score:.4f}  文档：{doc}")

1

从执行结果可看出ms-marco-MiniLM-L6-v2 模型对中文的重排效果很差

  1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157

from typing import List, Tuple, Optional, Union, Dict, Any
import torch
from transformers import AutoModel


class JinaRerankV3:
    """
    基于 transformers.AutoModel 封装的 jina-rerank-v3 重排器
    """
    DEFAULT_MODEL_NAME = "jinaai/jina-rerank-v3"

    def __init__(
            self,
            model_name_or_path: str = DEFAULT_MODEL_NAME,
            device: Optional[str] = None,
            batch_size: int = 32,  # jina-rerank-v3 批量推理最优批次
            trust_remote_code: bool = True,
            dtype: Union[str, torch.dtype] = "auto"
    ):
        """
        初始化 jina-rerank-v3 重排器（基于 AutoModel 原生接口）

        Args:
            model_name_or_path: 模型本地路径或HuggingFace名称
            device: 运行设备（auto/cuda/cpu）
            batch_size: 推理批次大小
            trust_remote_code: 是否允许加载模型仓库中自定义的代码（如 modeling_*.py、等），并执行这些代码来初始化模型/配置。
            dtype: 模型数据类型（auto/float16/float32）
        """
        # 设备适配
        self.device = device if device else ("cuda" if torch.cuda.is_available() else "cpu")
        self.batch_size = batch_size
        self.model_name_or_path = model_name_or_path
        self.trust_remote_code = trust_remote_code
        self.dtype = dtype if isinstance(dtype, torch.dtype) else dtype

        # 模型初始化
        self.model: Optional[AutoModel] = None
        self._load_model()

    def _load_model(self) -> None:
        """重载模型加载方法：使用 AutoModel 加载 jina-rerank-v3"""
        try:
            # 加载原生 AutoModel（jina官方推荐方式）
            self.model = AutoModel.from_pretrained(
                self.model_name_or_path,
                trust_remote_code=self.trust_remote_code,
                dtype=self.dtype,
                device_map=self.device  # 自动分配设备
            )
            # 设置评估模式
            self.model.eval()
            print(
                f"✅ Jina-Rerank-V3 模型加载完成 | 设备：{self.device} | 批次大小：{self.batch_size} | "
                f"数据类型：{self.dtype} | 评估模式：{self.model.training}"
            )
        except Exception as e:
            raise RuntimeError(
                f"Jina-Rerank-V3 模型加载失败：{e}\n"
                "排查建议：\n"
                "1. 模型文件是否完整 \n"
                "2. trust_remote_code 是否设为True \n"
                "3. transformers版本是否≥4.36.0"
            )

    def rerank_sorted(
            self,
            query: str,
            candidate_docs: List[str],
            return_scores: bool = True,
            top_k: Optional[int] = None
    ) -> Union[List[str], List[Tuple[str, float]]]:
        """
        核心重排方法
        调用 jina-rerank-v3 原生 rerank 方法，返回排序后的结果

        Args:
            query: 查询语句（支持多语言）
            candidate_docs: 候选文档列表
            return_scores: 是否返回相关性得分
            top_k: 返回前k个结果（None返回全部）

        Returns:
            排序后的文档列表 或 (文档, 得分) 元组列表
        """
        # 输入合法性校验
        if not candidate_docs:
            raise ValueError("候选文档列表不能为空")
        if self.model is None:
            raise RuntimeError("模型未初始化，请检查模型加载是否成功")

        # 核心推理：调用 jina 原生 rerank 方法
        try:
            with torch.no_grad():  # 禁用梯度计算（节省显存）
                results = self.model.rerank(
                    query=query,
                    documents=candidate_docs,
                    top_n=top_k if top_k else len(candidate_docs)
                )
        except Exception as e:
            raise RuntimeError(f"重排推理失败：{e}")

        # 解析结果（按得分降序排列，jina原生结果已排序）
        ranked_results = [
            (result["document"], float(result["relevance_score"]))
            for result in results
        ]

        # 适配返回格式
        if return_scores:
            return ranked_results
        else:
            return [item[0] for item in ranked_results]

    def release(self) -> None:
        """资源释放：清理模型和GPU缓存"""
        if self.model is not None:
            del self.model
            self.model = None
            if torch.cuda.is_available():
                torch.cuda.empty_cache()
            print("🔌 Jina-Rerank-V3 模型资源已释放")


# -------------------------- 测试示例 --------------------------
if __name__ == "__main__":
    # 1. 初始化重排器（本地模型路径）
    reranker = JinaRerankV3(
        model_name_or_path="../../data/models_reranker_data/jinaai/jina-reranker-v3",
        batch_size=16
    )

    # 2. 测试数据（多语言混合）
    query = "人工智能在医疗领域的应用"
    docs = [
        "人工智能驱动的农业无人机可实现作物病虫害监测，精准率达85%以上-E",  # 无关（完全跨领域）
        "AI辅助诊断系统可快速分析CT、MRI等医学影像，使肺癌、乳腺癌等疾病的早期检出率提升30%以上-A",  # 高相关
        "深度学习算法在金融风控中的应用案例，有效降低信贷违约率15%-E",  # 无关（金融领域）
        "卷积神经网络在皮肤癌图像分类任务上达到91%的准确率，为基层医疗机构提供了低成本诊断方案-B",  # 中高相关（技术细分，间接相关）
        "智能模拟在生物医药研发中的应用，主要用于药物分子结构模拟，不属于临床医疗应用范畴-D",  # 干扰项（语义相近但领域不符）
        "人工智能技术在医疗领域的核心应用包括AI辅助诊断、个性化治疗方案生成、医疗影像分析等-A",  # 高相关
        "机器学习算法在医保风控系统中的应用，可识别虚假就医报销行为，属于AI在医疗管理的边缘场景-D",  # 低相关（跨领域AI，弱关联）
        "医疗大数据平台通过机器学习算法整合患者电子病历，为医院管理决策提供数据支撑-C",  # 中相关（泛医疗AI，关联性减弱）
    ]

    # 3. 重排结果
    print("===== 重排结果=====")
    results_with_scores = reranker.rerank_sorted(
        query=query,
        candidate_docs=docs,
        return_scores=True
    )
    for idx, (doc, score) in enumerate(results_with_scores, 1):
        print(f"{idx}. 得分：{score:.4f} 文档：{doc[:100]}")

    # 4. 释放资源
    reranker.release()

  1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290

# 安装 ColBERT  pip install colbert-ai
# 额外依赖（适配模型加载） pip install torch transformers faiss-cpu  # GPU版可装 faiss-gpu
from typing import List, Tuple, Optional, Union
import os
import torch
from colbert import Indexer, Searcher
from colbert.infra import Run, RunConfig, ColBERTConfig


class ColBERTReranker:
    """
    ColBERT v2 重排器（适配百万级候选集的高效检索重排）
    """
    # 默认配置（平衡精度/速度/存储）
    DEFAULT_CHECKPOINT = "colbert-ir/colbertv2.0"
    DEFAULT_TOP_K = 100

    def __init__(
            self,
            checkpoint: str = DEFAULT_CHECKPOINT,
            index_root: str = "./data/colbert_index",
            nbits: int = 2,  # 2bit量化，百万级数据仅占GB级存储
            doc_maxlen: int = 512,
            nranks: int = 1,
            experiment_name: str = "default",
            device=None
    ):
        """
        初始化 ColBERT 重排器

        Args:
            checkpoint: ColBERT 模型权重（设置为HuggingFace模型名称/本地路径）
                - 1.若为模型名称：如果本地缓存（默认路径~ /.cache/huggingface/）中没有该模型，会自动从 Hugging Face Hub 下载模型权重、配置和分词器文件。
                若本地已存在该模型，代码会直接加载缓存中的文件，无需重新下载。由于国内无法访问huggingface，需要设置代理或镜像源方式。
                - 2.若为本地路径：则需要手动下载模型文件到路径地址.
            index_root: 索引存储根路径
            nbits: 量化位数（1/2/4/8，位数越低存储越小、精度略降）
            doc_maxlen: 文档最大长度
            nranks: 并行进程数（单机设1，分布式可增大）
            experiment_name: 实验名称（用于ColBERT日志）
        """
        # 基础配置
        self.checkpoint = checkpoint
        self.index_root = os.path.abspath(index_root)  # 索引根路径
        self.experiment_name = experiment_name
        self.nbits = nbits
        self.doc_maxlen = doc_maxlen
        self.nranks = nranks
        self.device = device if device else ("cuda" if torch.cuda.is_available() else "cpu")
        # 核心组件
        self.indexer: Optional[Indexer] = None
        self.searcher: Optional[Searcher] = None
        self.collection: Optional[List[str]] = None  # 存储原始段落集合（索引构建后缓存）
        # 日志提示
        print(f"🔧 ColBERT 初始化完成 | 模型：{checkpoint} | 索引路径：{self.index_root}")

    def _get_colbert_config(self) -> ColBERTConfig:
        """构建 ColBERT 配置对象（封装通用配置逻辑）"""
        config = ColBERTConfig(
            nbits=self.nbits,
            root=self.index_root,
            experiment=self.experiment_name,  # 关联实验名称
            doc_maxlen=self.doc_maxlen,
            bsize=128,  # 批量编码大小,额外优化参数（百万级场景）
        )
        return config

    def build_index(
            self,
            passages: List[str],
            index_name: str = "english_kb",
            overwrite: bool = False
    ) -> Indexer | None:
        """
        构建 ColBERT 量化索引（支持百万级段落）

        Args:
            passages: 段落列表（支持百万级数据）
            index_name: 索引名称（存储在index_path下的子目录）
            overwrite: 是否覆盖已有索引

        Returns:
            初始化后的 Indexer 实例
        """
        # 参数校验
        if not passages:
            raise ValueError("段落列表不能为空，无法构建索引")

        # 索引路径校验，Indexer默认使用的索引路径由root（实验根路径）/experiment（实验名）/indexes/index_name（索引名）拼接
        full_index_path = os.path.join(self.index_root, self.experiment_name, "indexes", index_name)
        if os.path.exists(full_index_path) and not overwrite:
            print(f"索引已存在：{full_index_path} | 如需覆盖请设置 overwrite=True")
            return

        # 缓存原始段落（用于检索时映射ID到文本）
        self.collection = passages

        try:
            print(f"🚀 开始构建 ColBERT 索引 | 段落数量：{len(passages):,} | 索引名称：{index_name}")
            run_config = RunConfig(
                nranks=self.nranks,
                root=self.index_root,  # 实验根路径
                experiment=self.experiment_name,  # 实验名
            )
            with Run().context(run_config):
                _config = self._get_colbert_config()
                self.indexer = Indexer(checkpoint=self.checkpoint, config=_config)
                # 构建索引（自动分批处理百万级数据）
                index_path = self.indexer.index(
                    name=index_name,
                    collection=passages,
                    overwrite=overwrite
                )

            print(f"✅ ColBERT 索引构建完成 | 存储路径：{index_path}")
            return self.indexer
        except Exception as e:
            raise RuntimeError(f"ColBERT 索引构建失败：{str(e)}") from e

    def load_index(self, index_name: str = "english_kb", collection: Optional[List[str]] = None) -> Searcher:
        """
        加载已构建的索引（避免重复构建，提升复用效率）

        Args:
            index_name: 索引名称
            collection: 原始段落列表（若未缓存需传入，用于映射ID到文本）

        Returns:
            初始化后的 Searcher 实例

        Raises:
            FileNotFoundError: 索引路径不存在
        """
        full_index_path = os.path.join(self.index_root, self.experiment_name, "indexes", index_name)

        if not os.path.exists(full_index_path):
            raise FileNotFoundError(f"索引不存在：{full_index_path} | 请先调用 build_index 构建索引")

        try:
            print(f"📥 加载 ColBERT 索引 | 路径：{full_index_path}")
            # RunConfig 关联路径和实验名
            run_config = RunConfig(
                nranks=self.nranks,
                root=self.index_root,
                experiment=self.experiment_name,
            )
            with Run().context(run_config):
                config = self._get_colbert_config()
                self.searcher = Searcher(
                    index=index_name,
                    checkpoint=self.checkpoint,
                    config=config
                )
            # 关联原始集合
            if collection is not None:
                self.searcher.collection = collection
            print(f"✅ ColBERT 索引加载完成 | 索引规模：{len(self.searcher.collection):,} 段落")
            return self.searcher
        except Exception as e:
            raise RuntimeError(f"ColBERT 索引加载失败：{str(e)}") from e

    def rerank(
            self,
            query: str,
            top_k: int = DEFAULT_TOP_K,
            index_name: str = "english_kb",
            load_index_if_needed: bool = True,
            collection: Optional[List[str]] = None
    ) -> List[Tuple[str, float]]:
        """
        基于 ColBERT 索引的高效重排

        Args:
            query: 查询语句（ColBERT v2 推荐英文，中文需适配中文版模型）
            top_k: 返回前k条最相关段落
            index_name: 索引名称
            load_index_if_needed: 若未加载索引，是否自动加载
            collection: 原始段落列表（若未缓存需传入，用于映射ID到文本）

        Returns:
            排序结果列表：[(段落文本, 相关性得分), ...]（按得分降序）

        Raises:
            RuntimeError: 索引未加载且自动加载失败

        """
        # 参数校验
        if top_k <= 0:
            raise ValueError(f"top_k 必须大于0，当前值：{top_k}")
        if not query.strip():
            raise ValueError("查询语句不能为空")

        # 自动加载索引
        if self.searcher is None and load_index_if_needed:
            self.load_index(index_name=index_name,collection=collection)

        # 校验索引是否就绪
        if self.searcher is None:
            raise RuntimeError("索引未加载 | 请先调用 build_index 或 load_index")

        try:
            print(f"开始 ColBERT 重排 | 查询：{query[:50]}... | Top-K：{top_k}")
            with Run().context(RunConfig(nranks=self.nranks)):
                # 核心检索（预编码+后期交互，百万级低延迟）
                results = self.searcher.search(query, k=top_k)
                # 解析结果：(段落ID, 得分) → (段落文本, 得分)
                pid_list, _, score_list = results  # 拆分3个列表（忽略第二个辅助列表）
                ranked_results = []
                for pid, score in zip(pid_list, score_list):  # 按位置一一对应
                    try:
                        passage_text = self.searcher.collection[pid]
                        ranked_results.append((passage_text, float(score)))
                    except KeyError:
                        print(f"警告：pid {pid} 在collection中不存在，跳过")

            print(f"ColBERT 重排完成 | 返回结果数：{len(ranked_results)}")
            return ranked_results
        except Exception as e:
            raise RuntimeError(f"ColBERT 重排失败：{str(e)}") from e

    def release(self):
        """释放资源"""
        if self.indexer is not None:
            del self.indexer
        if self.searcher is not None:
            del self.searcher
        self.collection = None
        print("✅ ColBERT 资源已释放")


# -------------------------- 测试示例 --------------------------
if __name__ == "__main__":
    # 1. 初始化 ColBERT 重排器
    colbert_reranker = ColBERTReranker(
        # 模型权重路径（二选一）：
        # 1. 填写HuggingFace模型名称，默认从HuggingFace自动加载,如：checkpoint="colbert-ir/colbertv2.0"
        # 2. 填写模型本地存储路径：需提前下载模型文件到指定路径
        #    模型文件获取地址：https://modelscope.cn/models/colbert-ir/colbertv2.0/
        #    需下载文件：config.json、model.safetensors、special_tokens_map.json、tokenizer.json、tokenizer_config.json、vocab.txt
        checkpoint="../../data/models_reranker_data/colbert-ir/colbertv2",  # 本地路径
        # checkpoint="./colbertv2_models",  # 本地路径
        index_root="./colbert_index",
        # 可选扩展配置（按需启用）：
        # nbits=2,                # 量化位数（1/2/4/8，位数越低存储越小、精度略降）
        # doc_maxlen=512,         # 文档最大token长度（适配BERT类模型）
        # experiment_name="colbert_v2_rerank",  # 实验名（用于索引路径分级）
        # device="cuda",          # 运行设备（自动优先GPU，无则用CPU）
        # nranks=1,               # 并行进程数（单机建议设1，分布式可增大）
    )

    # 2. 模拟百万级段落（实际可替换为真实数据）
    test_passages = [f"Document {i}: RAG optimization for large-scale knowledge base {i}" for i in range(1000)]
    # 插入目标段落（高相关）
    test_passages[100] = "Document 100: ColBERT v2 achieves 10x throughput for million-scale reranking"
    test_passages[200] = "Document 200: ColBERT v2 is optimized for million-scale reranking with 2-bit quantization"
    test_passages[400] = "Document 400: Memory-efficient index for 1M+ passages using ColBERT v2 2-bit quantization"
    # 插入无关段落（测试重排精度）
    test_passages[150] = "Document 150: Weather forecast for New York on day 2 - sunny with 25°C"
    test_passages[250] = "Document 250: How to bake a cake - step by step guide for beginners"
    test_passages[350] = "Document 350: Car maintenance tips for gasoline engines 3"

    # 3. 构建索引（首次运行耗时，后续可复用）
    try:
        colbert_reranker.build_index(
            passages=test_passages,
            index_name="english_kb",
            overwrite=False  # 覆盖已有索引（测试用）
        )
    except RuntimeError as e:
        print(f"索引构建失败：{e}")
        exit(0)

    # 4. 执行重排
    query = "Which model is efficient for million-scale English reranking?"
    try:
        results = colbert_reranker.rerank(
            query=query,
            top_k=5,
            index_name="english_kb",
            collection=test_passages
        )
    except RuntimeError as e:
        raise RuntimeError(f"重排失败：{str(e)}") from e

    print("\n===== ColBERT v2 百万级候选重排结果（前5条）=====")
    for i, (passage, score) in enumerate(results, 1):
        print(f"\n{i}. 得分：{score:.4f} 文档：{passage[:80]}")

    # 5. 释放资源
    colbert_reranker.release()