RAG 重排序实战：Cohere Rerank vs BGE-Reranker vs Jina

向量检索 (Bi-Encoder)          重排序 (Cross-Encoder)
Q → [Embedding] → ← [Embedding] ← Doc    Q + Doc → [Cross-Encoder] → Score
速度快，精度中等                  速度慢，精度高
召回阶段 (Top-50)                精排阶段 (Top-5)

import cohere

client = cohere.Client(api_key="your-api-key")

def cohere_rerank(query: str, documents: list, top_n: int = 5):
    response = client.rerank(
        model="rerank-multilingual-v3.0",
        query=query,
        documents=documents,
        top_n=top_n,
        max_tokens_per_doc=4096
    )
    
    return [
        {
            "index": r.index,
            "document": documents[r.index],
            "relevance_score": r.relevance_score
        }
        for r in response.results
    ]

from FlagEmbedding import FlagReranker

# 加载模型
reranker = FlagReranker(
    'BAAI/bge-reranker-v2-m3',
    use_fp16=True  # 半精度加速
)

def bge_rerank(query: str, documents: list, top_n: int = 5):
    # 构造 query-document 对
    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]
    ]

from jina import Client

# 通过 API 调用
def jina_rerank_api(query: str, documents: list, top_n: int = 5):
    client = Client(host="https://api.jina.ai/v1/rerank")
    
    response = client.post(
        "/rerank",
        json={
            "model": "jina-reranker-v2-base-multilingual",
            "query": query,
            "documents": documents,
            "top_n": top_n
        },
        headers={"Authorization": "Bearer your-api-key"}
    )
    
    return response.json()["results"]

# 本地部署
from transformers import AutoModelForSequenceClassification, AutoTokenizer

def jina_rerank_local(query: str, documents: list, top_n: int = 5):
    model_name = "jinaai/jina-reranker-v2-base-multilingual"
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = AutoModelForSequenceClassification.from_pretrained(model_name)
    
    pairs = [[query, doc] for doc in documents]
    inputs = tokenizer(pairs, padding=True, truncation=True, return_tensors="pt", max_length=8192)
    
    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]]

class RAGPipelineWithRerank:
    def __init__(self):
        self.embedder = EmbeddingModel("bge-m3")
        self.vector_store = MilvusIndex(...)
        self.reranker = FlagReranker("bge-reranker-v2-m3")
        self.llm = LLM("gpt-5.5")
    
    def query(self, question: str):
        # 1. 向量检索 Top-50（宽召回）
        query_emb = self.embedder.encode(question)
        candidates = self.vector_store.search(query_emb, top_k=50)
        
        # 2. 关键词检索补充（可选）
        keyword_results = self.keyword_search(question, top_k=20)
        candidates = self._merge(candidates, keyword_results)
        
        # 3. 重排序 Top-5（精排）
        doc_texts = [c.text for c in candidates]
        reranked = self.reranker.compute_score(
            [[question, doc] for doc in doc_texts],
            normalize=True
        )
        top_indices = np.argsort(reranked)[::-1][:5]
        top_docs = [candidates[i] for i in top_indices]
        
        # 4. 生成
        return self.llm.generate(question, context=top_docs)

def batch_rerank(query: str, documents: list, batch_size: int = 32):
    """批量重排序，避免 OOM"""
    all_scores = []
    
    for i in range(0, len(documents), batch_size):
        batch = documents[i:i+batch_size]
        pairs = [[query, doc] for doc in batch]
        scores = reranker.compute_score(pairs, normalize=True)
        all_scores.extend(scores if isinstance(scores, list) else [scores])
    
    return all_scores

def rerank_with_context_expansion(query: str, chunks: list, reranker):
    """重排序后，为 Top-K 扩展上下文窗口"""
    # 1. 重排序得到 Top-5
    scores = reranker.compute_score(
        [[query, c.text] for c in chunks], normalize=True
    )
    top_indices = np.argsort(scores)[::-1][:5]
    
    # 2. 为每个 top chunk 扩展上下文（前后各 1 个 chunk）
    expanded = []
    for idx in top_indices:
        main_chunk = chunks[idx]
        prev_chunk = chunks[idx - 1] if idx > 0 else None
        next_chunk = chunks[idx + 1] if idx < len(chunks) - 1 else None
        
        expanded_text = ""
        if prev_chunk:
            expanded_text += prev_chunk.text[-200:] + "\n---\n"
        expanded_text += main_chunk.text
        if next_chunk:
            expanded_text += "\n---\n" + next_chunk.text[:200]
        
        expanded.append({
            "text": expanded_text,
            "main_chunk_id": main_chunk.id,
            "score": scores[idx]
        })
    
    return expanded