Reranking — Cross-Encoders, ColBERT & Contextual Compression

from sentence_transformers import CrossEncoder
import time

# Load once at startup — this is a 90 MB model
reranker = CrossEncoder("cross-encoder/ms-marco-MiniLM-L-6-v2", max_length=512)


def rerank_with_cross_encoder(
    query: str,
    candidates: list[dict],
    top_k: int = 5,
) -> list[dict]:
    """
    Rerank candidate documents using a cross-encoder.

    Args:
        query: the user query string
        candidates: list of dicts with at least a "text" key
        top_k: number of top documents to return after reranking

    Returns:
        top_k documents sorted by cross-encoder relevance score, descending
    """
    if not candidates:
        return []

    # Build (query, passage) pairs for the cross-encoder
    pairs = [(query, c["text"]) for c in candidates]

    t0 = time.perf_counter()
    # predict() returns a float score per pair (higher = more relevant)
    scores = reranker.predict(pairs, batch_size=32, show_progress_bar=False)
    latency_ms = (time.perf_counter() - t0) * 1000

    # Attach scores and sort
    for candidate, score in zip(candidates, scores):
        candidate["rerank_score"] = float(score)

    sorted_candidates = sorted(candidates, key=lambda x: x["rerank_score"], reverse=True)
    print(f"Reranking {len(candidates)} candidates took {latency_ms:.1f} ms")
    return sorted_candidates[:top_k]

from sentence_transformers import SentenceTransformer

embed_model = SentenceTransformer("BAAI/bge-large-en-v1.5")


def retrieve_and_rerank(
    query: str,
    retriever_fn,      # fn(embedding, k) -> list[dict with id, text, score]
    k_retrieve: int = 50,
    k_final: int = 5,
) -> dict:
    """
    Full pipeline: embed query → retrieve top-50 → rerank → return top-5.
    Returns timings so you can monitor each stage's latency.
    """
    timings = {}

    # Stage 1: embed the query
    t0 = time.perf_counter()
    query_vec = embed_model.encode([query], normalize_embeddings=True)[0].tolist()
    timings["embed_ms"] = (time.perf_counter() - t0) * 1000

    # Stage 2: ANN retrieval — fast, approximate
    t0 = time.perf_counter()
    candidates = retriever_fn(query_vec, k=k_retrieve)
    timings["retrieve_ms"] = (time.perf_counter() - t0) * 1000

    # Stage 3: cross-encoder reranking — slow, accurate
    t0 = time.perf_counter()
    final_results = rerank_with_cross_encoder(query, candidates, top_k=k_final)
    timings["rerank_ms"] = (time.perf_counter() - t0) * 1000

    timings["total_ms"] = sum(timings.values())

    return {"results": final_results, "timings": timings}

import cohere

co = cohere.Client(api_key="COHERE_API_KEY")


def cohere_rerank(
    query: str,
    candidates: list[dict],
    top_k: int = 5,
    model: str = "rerank-english-v3.0",
) -> list[dict]:
    """Rerank using the Cohere Rerank API."""
    documents = [c["text"] for c in candidates]

    response = co.rerank(
        model=model,
        query=query,
        documents=documents,
        top_n=top_k,
        return_documents=False,  # we keep the originals
    )

    # response.results: list of RerankResult(index, relevance_score)
    reranked = []
    for result in response.results:
        doc = dict(candidates[result.index])
        doc["rerank_score"] = result.relevance_score
        reranked.append(doc)

    return reranked

# ColBERT via the RAGatouille library (wraps the official ColBERT implementation)
from ragatouille import RAGPretrainedModel

colbert = RAGPretrainedModel.from_pretrained("colbert-ir/colbertv2.0")

# Index documents (this creates a ColBERT-specific index)
colbert.index(
    collection=["document text 1", "document text 2", ...],
    index_name="my_docs",
    max_document_length=256,
    split_documents=True,
)

# Search
results = colbert.search(
    query="how does ColBERT compute relevance?",
    k=10,
)

RRF(d) = Σᵢ 1 / (k + rankᵢ(d))

def reciprocal_rank_fusion(
    ranked_lists: list[list[str]],
    k: int = 60,
) -> list[tuple[str, float]]:
    """
    Merge multiple ranked lists using RRF.

    Args:
        ranked_lists: each inner list is a ranked sequence of document IDs
                      from one retrieval source, best-first
        k: constant (default 60) — dampens the advantage of top ranks

    Returns:
        list of (doc_id, rrf_score) sorted descending
    """
    scores: dict[str, float] = {}
    for ranked_list in ranked_lists:
        for rank, doc_id in enumerate(ranked_list, start=1):
            scores[doc_id] = scores.get(doc_id, 0.0) + 1.0 / (k + rank)
    return sorted(scores.items(), key=lambda x: x[1], reverse=True)

COMPRESS_PROMPT = """You are a document extraction assistant.

Given a USER QUERY and a DOCUMENT CHUNK, extract only the sentences from the chunk
that are directly relevant to answering the query.

Return ONLY the extracted sentences verbatim. If nothing is relevant, return an empty string.
Do not paraphrase or summarise.

USER QUERY: {query}

DOCUMENT CHUNK:
{chunk}

Relevant sentences:"""

from groq import Groq

groq_client = Groq()


def compress_chunk(query: str, chunk_text: str) -> str:
    """Extract only query-relevant sentences from a chunk."""
    response = groq_client.chat.completions.create(
        model="llama-3.1-8b-instant",
        messages=[{
            "role": "user",
            "content": COMPRESS_PROMPT.format(query=query, chunk=chunk_text),
        }],
        temperature=0.0,
        max_tokens=300,
    )
    compressed = response.choices[0].message.content.strip()
    # Fall back to original if compression returned empty
    return compressed if len(compressed) > 20 else chunk_text


def compress_results(query: str, results: list[dict]) -> list[dict]:
    """Apply contextual compression to a list of reranked results."""
    compressed = []
    for r in results:
        original_len = len(r["text"].split())
        compressed_text = compress_chunk(query, r["text"])
        compressed_len = len(compressed_text.split())
        compressed.append({
            **r,
            "text": compressed_text,
            "compression_ratio": compressed_len / original_len if original_len > 0 else 1.0,
        })
    return compressed

def full_reranking_pipeline(
    query: str,
    dense_results: list[dict],
    sparse_results: list[dict],    # BM25 results
    top_k: int = 5,
    compress: bool = False,
) -> dict:
    """
    Complete pipeline: RRF merge → cross-encoder rerank → optional compression.
    """
    # Merge dense and BM25 results with RRF
    dense_ids = [r["id"] for r in dense_results]
    sparse_ids = [r["id"] for r in sparse_results]
    merged_ids_scores = reciprocal_rank_fusion([dense_ids, sparse_ids])

    # Build result map for lookup
    result_map = {r["id"]: r for r in dense_results + sparse_results}
    candidates = [result_map[doc_id] for doc_id, _ in merged_ids_scores if doc_id in result_map]

    # Cross-encoder reranking on top-50 merged candidates
    reranked = rerank_with_cross_encoder(query, candidates[:50], top_k=top_k)

    if compress:
        reranked = compress_results(query, reranked)

    return {"results": reranked, "candidate_count": len(candidates)}