Reranker模型选型：Cohere vs BGE vs Jina对比

第一阶段：召回（Recall）
  Embedding向量检索，快速从百万文档中召回Top-K（K=20-50）
  优势：速度快（<50ms）
  劣势：精度有限（Embedding模型容量小，语义理解浅）

第二阶段：精排（Precision）
  Cross-Encoder Reranker对Top-K重新打分排序
  优势：精度高（Cross-Encoder能深度理解query-doc关系）
  劣势：速度慢（需对每个候选单独计算）

Bi-Encoder (Embedding模型)：
  Query ──► Encoder ──► [向量] ◄── Encoder ◄── Document
                                    │
                              余弦相似度
                              快但浅

Cross-Encoder (Reranker)：
  [Query, Document] ──► Encoder ──► 相关性分数
                              慢但深

import cohere

co = cohere.Client(api_key="your-key")

def cohere_rerank(query, documents, top_n=5):
    results = co.rerank(
        model="rerank-v3.5",
        query=query,
        documents=documents,
        top_n=top_n,
    )
    
    reranked = []
    for result in results.results:
        reranked.append({
            "document": documents[result.index],
            "relevance_score": result.relevance_score,
            "index": result.index,
        })
    
    return reranked

from FlagEmbedding import FlagReranker

# 加载模型
reranker = FlagReranker(
    "BAAI/bge-reranker-v2-m3",
    use_fp16=True,
    device="cuda",
)

def bge_rerank(query, documents, top_n=5):
    # 构建query-doc对
    pairs = [[query, doc] for doc in documents]
    
    # 计算相关性分数
    scores = reranker.compute_score(pairs, normalize=True)
    
    # 排序
    ranked = sorted(
        zip(documents, scores),
        key=lambda x: x[1],
        reverse=True
    )
    
    return [
        {"document": doc, "relevance_score": score}
        for doc, score in ranked[:top_n]
    ]

# 方式1: API调用
import requests

def jina_rerank_api(query, documents, top_n=5):
    response = requests.post(
        "https://api.jina.ai/v1/rerank",
        headers={
            "Authorization": "Bearer your-key",
            "Content-Type": "application/json",
        },
        json={
            "model": "jina-reranker-v2-base-multilingual",
            "query": query,
            "documents": documents,
            "top_n": top_n,
        },
    )
    
    results = response.json()["results"]
    return [
        {"document": documents[r["index"]], "relevance_score": r["relevance_score"]}
        for r in results
    ]

# 方式2: 本地部署
from transformers import AutoModelForSequenceClassification, AutoTokenizer
import torch

tokenizer = AutoTokenizer.from_pretrained("jinaai/jina-reranker-v2-base-multilingual")
model = AutoModelForSequenceClassification.from_pretrained(
    "jinaai/jina-reranker-v2-base-multilingual",
    torch_dtype=torch.float16,
).cuda()

def jina_rerank_local(query, documents, top_n=5):
    pairs = [[query, doc] for doc in documents]
    
    inputs = tokenizer(
        pairs,
        padding=True,
        truncation=True,
        max_length=8192,
        return_tensors="pt",
    ).to("cuda")
    
    with torch.no_grad():
        scores = model(**inputs).logits.squeeze(-1)
    
    ranked = sorted(zip(documents, scores.tolist()), key=lambda x: x[1], reverse=True)
    return [{"document": d, "relevance_score": s} for d, s in ranked[:top_n]]

def cascade_rerank(query, candidates, top_k=5):
    """级联Reranker：先快后精"""
    
    # 第一阶段：轻量Reranker快速筛选（BGE-base）
    light_reranker = FlagReranker("BAAI/bge-reranker-base", device="cuda")
    pairs = [[query, doc] for doc in candidates]
    scores = light_reranker.compute_score(pairs)
    
    # 取Top-10
    top_10 = [candidates[i] for i in np.argsort(scores)[-10:]]
    
    # 第二阶段：重量Reranker精排（BGE-gemma）
    heavy_reranker = FlagReranker("BAAI/bge-reranker-v2-gemma", device="cuda")
    pairs_2 = [[query, doc] for doc in top_10]
    scores_2 = heavy_reranker.compute_score(pairs_2, normalize=True)
    
    # 取Top-5
    final = sorted(zip(top_10, scores_2), key=lambda x: x[1], reverse=True)[:top_k]
    
    return final

def hybrid_rerank(query, candidates, top_n=5):
    """多个Reranker分数融合"""
    
    # BGE分数
    bge_scores = bge_reranker.compute_score(
        [[query, doc] for doc in candidates], normalize=True
    )
    
    # Jina分数
    jina_scores = jina_reranker.predict(
        [[query, doc] for doc in candidates]
    )
    jina_scores = min_max_normalize(jina_scores)
    
    # 分数融合
    fused_scores = 0.6 * np.array(bge_scores) + 0.4 * np.array(jina_scores)
    
    # 排序
    ranked = sorted(zip(candidates, fused_scores), key=lambda x: x[1], reverse=True)
    return ranked[:top_n]

from transformers import AutoModelForSequenceClassification, Trainer, TrainingArguments

def fine_tune_reranker(train_data, model_name="BAAI/bge-reranker-v2-m3"):
    """微调Reranker以适配特定领域"""
    
    model = AutoModelForSequenceClassification.from_pretrained(
        model_name, num_labels=1
    )
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    
    # train_data格式: [{"query": "...", "doc": "...", "label": 0/1}]
    
    def preprocess(examples):
        return tokenizer(
            examples["query"], examples["doc"],
            truncation=True, max_length=512, padding="max_length"
        )
    
    dataset = dataset.map(preprocess, batched=True)
    
    training_args = TrainingArguments(
        output_dir="./domain-reranker",
        num_train_epochs=3,
        per_device_train_batch_size=16,
        learning_rate=2e-5,
        warmup_ratio=0.1,
        bf16=True,
    )
    
    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=dataset,
    )
    
    trainer.train()
    return model

# 推荐配置
config = {
    "embedding_model": "BAAI/bge-m3",        # Embedding模型
    "vector_search_top_k": 20,                 # 初筛召回数
    "reranker_model": "BAAI/bge-reranker-v2-m3", # Reranker
    "reranker_top_k": 5,                       # Reranker输出数
    "reranker_max_length": 512,                # 最大长度
    "reranker_batch_size": 20,                 # 批量大小
}

# 完整RAG检索pipeline
def rag_retrieve(query, vector_store, reranker, config):
    # 第一步：向量检索召回Top-20
    candidates = vector_store.similarity_search(query, k=config["vector_search_top_k"])
    
    # 第二步：Reranker精排到Top-5
    pairs = [[query, doc.page_content] for doc in candidates]
    scores = reranker.compute_score(pairs, normalize=True)
    
    # 第三步：过滤低分文档
    ranked = sorted(zip(candidates, scores), key=lambda x: x[1], reverse=True)
    final_docs = [
        doc for doc, score in ranked[:config["reranker_top_k"]]
        if score > 0.1  # 分数阈值，过滤不相关文档
    ]
    
    return final_docs