Query → 向量检索 Top-50 → 重排序 → Top-5 → LLM 生成

# 两阶段检索流程
def two_stage_retrieve(
    query: str,
    vector_index,          # 双塔向量索引
    reranker,              # 交叉编码器重排序
    first_stage_k: int = 50,   # 初筛数量
    final_k: int = 5,          # 最终数量
) -> list:
    # Stage 1: 向量检索（快，召回多）
    candidates = vector_index.search(query, top_k=first_stage_k)

    # Stage 2: 重排序（精，筛少）
    pairs = [(query, doc["content"]) for doc in candidates]
    scores = reranker.predict(pairs)

    # 排序并取 Top-K
    ranked = sorted(zip(candidates, scores), key=lambda x: x[1], reverse=True)
    return [doc for doc, score in ranked[:final_k]]

import cohere

class CohereReranker:
    """Cohere Rerank 封装"""
    def __init__(self, api_key: str, model: str = "rerank-multilingual-v3.0"):
        self.client = cohere.Client(api_key)
        self.model = model

    def rerank(
        self,
        query: str,
        documents: list[str],
        top_n: int = 5,
        max_chunks_per_doc: int = None,
    ) -> list[dict]:
        """
        对文档列表重排序
        返回: [{"index": int, "relevance_score": float, "document": str}, ...]
        """
        response = self.client.rerank(
            model=self.model,
            query=query,
            documents=documents,
            top_n=top_n,
            max_chunks_per_doc=max_chunks_per_doc,
        )

        results = []
        for result in response.results:
            results.append({
                "index": result.index,
                "relevance_score": result.relevance_score,
                "document": documents[result.index],
            })
        return results


# 在 RAG 中使用
reranker = CohereReranker(api_key="your-api-key")

# 假设向量检索已返回 50 篇候选
candidates = vector_index.search(query, top_k=50)
doc_texts = [doc["content"] for doc in candidates]

# 重排序取 Top-5
reranked = reranker.rerank(query, doc_texts, top_n=5)
for r in reranked:
    print(f"Score: {r['relevance_score']:.4f} | {r['document'][:100]}...")

import torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer

class BGEReranker:
    """bge-reranker 本地部署"""
    def __init__(
        self,
        model_name: str = "BAAI/bge-reranker-v2-m3",
        device: str = "cuda",
        max_length: int = 512,
        batch_size: int = 32,
    ):
        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
        self.model = AutoModelForSequenceClassification.from_pretrained(model_name)
        self.model.to(device)
        self.model.eval()
        self.device = device
        self.max_length = max_length
        self.batch_size = batch_size

    @torch.no_grad()
    def compute_scores(
        self,
        pairs: list[tuple[str, str]],  # [(query, doc), ...]
    ) -> list[float]:
        """计算每对 (query, doc) 的相关性分数"""
        all_scores = []

        for i in range(0, len(pairs), self.batch_size):
            batch = pairs[i:i + self.batch_size]
            texts = [(q, d) for q, d in batch]

            inputs = self.tokenizer(
                texts,
                padding=True,
                truncation=True,
                max_length=self.max_length,
                return_tensors="pt",
            ).to(self.device)

            scores = self.model(**inputs).logits.view(-1).float()
            all_scores.extend(scores.cpu().tolist())

        return all_scores

    def rerank(
        self,
        query: str,
        documents: list[str],
        top_n: int = 5,
    ) -> list[dict]:
        """重排序主入口"""
        pairs = [(query, doc) for doc in documents]
        scores = self.compute_scores(pairs)

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

        results = []
        for idx, score in ranked[:top_n]:
            results.append({
                "index": idx,
                "relevance_score": score,
                "document": documents[idx],
            })
        return results


# 使用
reranker = BGEReranker("BAAI/bge-reranker-v2-m3", device="cuda")
results = reranker.rerank(query, candidates, top_n=5)

# 基于 Gemma-2B 的重排序器，效果接近 Cohere
class BGEGemmaReranker:
    def __init__(self, model_name="BAAI/bge-reranker-v2-gemma"):
        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
        self.model = AutoModelForSequenceClassification.from_pretrained(
            model_name,
            torch_dtype=torch.bfloat16,
            device_map="auto",
        )
        self.model.eval()

    @torch.no_grad()
    def rerank(self, query, documents, top_n=5, max_length=8192):
        pairs = [(query, doc) for doc in documents]
        inputs = self.tokenizer(
            pairs,
            padding=True,
            truncation=True,
            max_length=max_length,
            return_tensors="pt",
        ).to(self.model.device)

        scores = self.model(**inputs).logits.view(-1).float()
        ranked = sorted(enumerate(scores.tolist()), key=lambda x: x[1], reverse=True)
        return [{"index": i, "score": s, "document": documents[i]} for i, s in ranked[:top_n]]

import json
import random
from typing import List, Tuple

def build_training_data(
    queries: List[str],
    corpus: List[str],
    positive_pairs: List[Tuple[int, int]],  # (query_idx, doc_idx) 正样本
    negative_ratio: int = 4,
) -> List[dict]:
    """
    构造重排序训练数据
    """
    data = []
    for q_idx, d_idx in positive_pairs:
        # 正样本
        data.append({
            "query": queries[q_idx],
            "document": corpus[d_idx],
            "label": 1,
        })
        # 负采样
        neg_indices = random.sample(
            [i for i in range(len(corpus)) if i != d_idx],
            negative_ratio,
        )
        for neg_idx in neg_indices:
            data.append({
                "query": queries[q_idx],
                "document": corpus[neg_idx],
                "label": 0,
            })
    return data

from transformers import (
    AutoModelForSequenceClassification,
    AutoTokenizer,
    Trainer,
    TrainingArguments,
)
from datasets import Dataset
import torch

def train_custom_reranker(
    train_data: List[dict],
    base_model: str = "BAAI/bge-reranker-base",
    output_dir: str = "./custom-reranker",
    epochs: int = 3,
    lr: float = 2e-5,
    batch_size: int = 16,
):
    """微调自定义重排序模型"""
    tokenizer = AutoTokenizer.from_pretrained(base_model)
    model = AutoModelForSequenceClassification.from_pretrained(
        base_model,
        num_labels=2,
    )

    # 编码数据
    def encode(examples):
        features = tokenizer(
            list(zip(examples["query"], examples["document"])),
            padding="max_length",
            truncation=True,
            max_length=512,
            return_tensors="pt",
        )
        features["labels"] = examples["label"]
        return features

    dataset = Dataset.from_list(train_data)
    dataset = dataset.map(encode, batched=True, remove_columns=dataset.column_names)

    training_args = TrainingArguments(
        output_dir=output_dir,
        num_train_epochs=epochs,
        per_device_train_batch_size=batch_size,
        learning_rate=lr,
        warmup_ratio=0.1,
        weight_decay=0.01,
        save_strategy="epoch",
        evaluation_strategy="no",
        fp16=True,
    )

    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=dataset,
        tokenizer=tokenizer,
    )

    trainer.train()
    model.save_pretrained(output_dir)
    tokenizer.save_pretrained(output_dir)
    return model

def hierarchical_rerank(
    query: str,
    candidates: list[dict],
    fast_reranker,          # bge-reranker-base（快速粗排）
    strong_reranker,        # bge-reranker-v2-gemma（精排）
    fast_top_n: int = 20,   # 粗排保留数量
    final_top_n: int = 5,   # 精排保留数量
) -> list[dict]:
    """
    两级重排序：先快排筛减，再强排精选
    """
    # Stage 1: 快速重排序（粗排）
    docs = [c["content"] for c in candidates]
    fast_results = fast_reranker.rerank(query, docs, top_n=fast_top_n)

    # Stage 2: 强力重排序（精排）
    refined_docs = [r["document"] for r in fast_results]
    final_results = strong_reranker.rerank(query, refined_docs, top_n=final_top_n)

    # 映射回原始索引
    return final_results

# 动态批处理：将多个 query 的候选合并处理
class BatchRerankServer:
    def __init__(self, reranker, max_batch_size=64, max_wait_ms=50):
        self.reranker = reranker
        self.max_batch_size = max_batch_size
        self.max_wait_ms = max_wait_ms
        self.queue = []

    async def rerank(self, query, documents, top_n=5):
        # 入队等待批处理
        self.queue.append((query, documents, top_n))
        # ... 批处理逻辑

from functools import lru_cache
import hashlib

@lru_cache(maxsize=10000)
def cached_rerank(query: str, doc_hash: str, docs_tuple: tuple) -> tuple:
    """LRU 缓存重排序结果"""
    return tuple(reranker.rerank(query, list(docs_tuple)))

# 使用时计算文档 hash 作为缓存 key
def rerank_with_cache(query, documents):
    doc_hash = hashlib.md5(
        "".join(documents).encode()
    ).hexdigest()
    return cached_rerank(query, doc_hash, tuple(documents))

# 将 bge-reranker 导出为 ONNX 格式，推理速度提升 1.5-3x
def export_to_onnx(model, tokenizer, output_path):
    model.eval()
    dummy_input = tokenizer(
        ["hello"], ["world"],
        padding=True, truncation=True,
        return_tensors="pt"
    )
    torch.onnx.export(
        model,
        (dummy_input["input_ids"], dummy_input["attention_mask"]),
        output_path,
        opset_version=14,
        input_names=["input_ids", "attention_mask"],
        output_names=["logits"],
        dynamic_axes={
            "input_ids": {0: "batch", 1: "seq"},
            "attention_mask": {0: "batch", 1: "seq"},
        },
    )

数据隐私要求高？
├── 是 → 不能用 Cohere API
│   ├── 有 GPU → bge-reranker-v2-m3
│   └── 仅 CPU → bge-reranker-base
└── 否 → 可以用 Cohere
    ├── 预算充足 + 追求最高质量 → Cohere Rerank v3.5
    └── 预算有限 → bge-reranker（本地部署）

领域差异大？
├── 是 → bge-reranker-base 微调（自训练）
└── 否 → 直接用预训练模型

延迟要求？
├── <100ms → bge-reranker-base (GPU) 或层级重排序
├── 100-500ms → bge-reranker-large / Cohere API
└── >500ms 可接受 → bge-reranker-v2-gemma