04-04-01 — Re-ranking

O problema com embeddings para ranking final

Embeddings (bi-encoders) codificam query e documento separadamente. A comparação por cosine similarity é rápida, mas superficial — o modelo não "vê" query e documento ao mesmo tempo, então não pode capturar relações sutis entre eles.

Cross-encoders fazem diferente: recebem [query, documento] como entrada única e produzem um score de relevância. Como processam os dois juntos, capturam nuances muito melhor. Mas são lentos — não é prático rodar um cross-encoder sobre 1 milhão de chunks para cada query.

A solução elegante: use bi-encoder para recuperar rápido um conjunto maior de candidatos (top-100), depois cross-encoder para reordenar esse conjunto menor com mais precisão.

Opção 1: Cohere Rerank

API gerenciada, zero infra, excelente qualidade. Suporta português. Ideal para quem não está no Azure ou quer a melhor qualidade disponível via API.

import cohere

co = cohere.Client("SUA_API_KEY")

def rerank_with_cohere(query: str, candidates: list[dict], top_n: int = 5) -> list[dict]:
    """
    candidates: lista de dicts com "id" e "content"
    """
    rerank_results = co.rerank(
        query=query,
        documents=[c["content"] for c in candidates],
        top_n=top_n,
        model="rerank-multilingual-v3.0"  # suporta português
    )

    reranked = []
    for r in rerank_results.results:
        candidate = candidates[r.index].copy()
        candidate["rerank_score"] = r.relevance_score
        reranked.append(candidate)

    return reranked

# Pipeline completo
initial_candidates = vector_search(query, k=50)      # 50 candidatos rápidos
final_results = rerank_with_cohere(query, initial_candidates, top_n=5)  # 5 melhores

Opção 2: Azure AI Semantic Ranker

Integrado ao Azure AI Search. Funciona como cross-encoder sobre os candidatos do hybrid search. Vantagem: sem latência de rede adicional, já está dentro do Azure AI Search. Desvantagem: só funciona com índices Azure AI Search.

# Já visto em 04-02-03, mas aqui o foco é o reranking
results = search_client.search(
    search_text=query,
    vector_queries=[VectorizedQuery(vector=embed(query), k_nearest_neighbors=50, fields="content_vector")],
    query_type="semantic",
    semantic_configuration_name="my-semantic-config",
    top=5,
    # Retornar o caption (trecho mais relevante identificado pelo ranker)
    query_caption="extractive",
    query_answer="extractive"  # resposta direta se encontrada
)

for r in results:
    print(f"Score: {r['@search.reranker_score']:.3f}")  # score do semantic ranker
    print(f"Caption: {r['@search.captions'][0].text if r.get('@search.captions') else 'N/A'}")

Opção 3: Cross-encoder local (self-hosted)

Para quem quer controle total e latência mínima. Modelos open-source como cross-encoder/ms-marco-MiniLM-L-6-v2 são pequenos (~22MB) e podem rodar no mesmo servidor da aplicação.

from sentence_transformers import CrossEncoder

# Carregar modelo cross-encoder (download ~22MB na primeira vez)
cross_encoder = CrossEncoder("cross-encoder/ms-marco-MiniLM-L-6-v2")

def rerank_local(query: str, candidates: list[dict], top_n: int = 5) -> list[dict]:
    pairs = [(query, c["content"]) for c in candidates]
    scores = cross_encoder.predict(pairs)

    ranked = sorted(
        zip(candidates, scores),
        key=lambda x: x[1],
        reverse=True
    )
    return [{"chunk": c, "score": float(s)} for c, s in ranked[:top_n]]

Comparativo das opções