Reranking Techniques

# Cross-encoder reranking with sentence-transformers
from sentence_transformers import CrossEncoder
import numpy as np

# Load a cross-encoder model
reranker = CrossEncoder("cross-encoder/ms-marco-MiniLM-L-12-v2")

def rerank_documents(
    query: str,
    documents: list[str],
    top_k: int = 5,
) -> list[tuple[str, float]]:
    """Rerank documents using a cross-encoder."""
    # Create query-document pairs
    pairs = [(query, doc) for doc in documents]

    # Score all pairs
    scores = reranker.predict(pairs)

    # Sort by score (descending)
    scored_docs = list(zip(documents, scores))
    scored_docs.sort(key=lambda x: x[1], reverse=True)

    return scored_docs[:top_k]

# Example: Initial retrieval returns 20 candidates, reranking picks top 5
candidates = [
    "To reset your password, go to Settings > Security",
    "Password requirements: 8 chars, 1 uppercase, 1 number",
    "Account recovery options include email and phone",
    "Two-factor authentication adds an extra security layer",
    "Contact support for locked account assistance",
    "Our pricing plans start at $9.99/month",
    "API documentation is available at docs.example.com",
]

reranked = rerank_documents(
    query="How do I change my password?",
    documents=candidates,
    top_k=3,
)

for doc, score in reranked:
    print(f"[{score:.4f}] {doc}")

// Cohere Rerank API
import { CohereClient } from "cohere-ai";

const cohere = new CohereClient({
  token: process.env.COHERE_API_KEY!,
});

async function cohereRerank(
  query: string,
  documents: string[],
  topK: number = 5
): Promise<{ text: string; score: number; index: number }[]> {
  const response = await cohere.rerank({
    model: "rerank-english-v3.0",
    query,
    documents,
    topN: topK,
    returnDocuments: true,
  });

  return response.results.map(r => ({
    text: documents[r.index],
    score: r.relevanceScore,
    index: r.index,
  }));
}

// Integration with RAG pipeline
async function ragWithReranking(question: string): Promise<string> {
  // Step 1: Retrieve broad set of candidates
  const candidates = await vectorStore.similaritySearch(question, 20);

  // Step 2: Rerank with Cohere
  const reranked = await cohereRerank(
    question,
    candidates.map(c => c.pageContent),
    5
  );

  // Step 3: Use top reranked results as context
  const context = reranked.map(r => r.text).join("\n\n");

  // Step 4: Generate answer
  const response = await anthropic.messages.create({
    model: "claude-sonnet-4-20250514",
    max_tokens: 1024,
    system: `Answer based on context:\n${context}`,
    messages: [{ role: "user", content: question }],
  });

  return response.content[0].type === "text" ? response.content[0].text : "";
}

# Cohere Rerank in Python
import cohere

co = cohere.Client()

def rerank_with_cohere(
    query: str,
    documents: list[str],
    top_n: int = 5,
) -> list[dict]:
    results = co.rerank(
        model="rerank-english-v3.0",
        query=query,
        documents=documents,
        top_n=top_n,
        return_documents=True,
    )

    return [
        {
            "text": r.document.text,
            "score": r.relevance_score,
            "original_index": r.index,
        }
        for r in results.results
    ]

# Usage
docs = ["doc1...", "doc2...", "doc3..."]
reranked = rerank_with_cohere("my query", docs)
for r in reranked:
    print(f"[{r['score']:.4f}] {r['text'][:80]}...")

// LLM-based reranking with Claude
async function llmRerank(
  query: string,
  documents: { id: string; content: string }[],
  topK: number = 5
): Promise<string[]> {
  const docList = documents
    .map((d, i) => `[Doc ${i + 1}] ${d.content.slice(0, 200)}`)
    .join("\n\n");

  const response = await client.messages.create({
    model: "claude-sonnet-4-20250514",
    max_tokens: 256,
    messages: [
      {
        role: "user",
        content: `Given the query: "${query}"

Rank these documents by relevance. Return ONLY the document numbers in order of relevance, most relevant first.
Format: 3, 1, 5, 2, 4

Documents:
${docList}`,
      },
    ],
  });

  const text = (response.content[0] as any).text;
  const indices = text.match(/\d+/g)?.map(Number) || [];

  return indices
    .filter(i => i >= 1 && i <= documents.length)
    .slice(0, topK)
    .map(i => documents[i - 1].content);
}