src/retrieval/reranker.py

"""
src/retrieval/reranker.py
Cross-encoder reranking (BAAI/bge-reranker-v2-m3).

Pipeline:
  Stage 1 → Chroma dense retrieval → top-20 candidates
  Stage 2 → CrossEncoder scores each (query, chunk) pair → returns top-5

Latency: ~0.5–2s CPU, ~0.1–0.3s GPU per query.
Quality:  +5–15% MRR@10 on scientific QA vs dense-only.
"""
import time
from dataclasses import dataclass


@dataclass
class RankedChunk:
    text: str
    metadata: dict
    distance: float
    rerank_score: float


class CrossEncoderReranker:
    DEFAULT_MODEL = "BAAI/bge-reranker-v2-m3"

    def __init__(self, model_name: str = None):
        from sentence_transformers import CrossEncoder
        self.model_name = model_name or self.DEFAULT_MODEL
        self.model = CrossEncoder(self.model_name, max_length=512)
        print(f"✅ Reranker loaded: {self.model_name}")

    def rerank(self, query: str, retrieval_result: dict, top_k: int = 5) -> list:
        docs = retrieval_result["documents"]
        metas = retrieval_result["metadatas"]
        dists = retrieval_result["distances"]
        if not docs:
            return []
        scores = self.model.predict([(query, d) for d in docs], show_progress_bar=False)
        ranked = sorted(
            [RankedChunk(text=d, metadata=m, distance=dist, rerank_score=float(s))
             for d, m, dist, s in zip(docs, metas, dists, scores)],
            key=lambda x: x.rerank_score, reverse=True,
        )
        return ranked[:top_k]

    def rerank_with_timing(self, query, retrieval_result, top_k=5):
        t0 = time.time()
        return self.rerank(query, retrieval_result, top_k), time.time() - t0


def ranked_to_result(ranked):
    return {
        "documents": [r.text for r in ranked],
        "metadatas": [r.metadata for r in ranked],
        "distances": [r.distance for r in ranked],
        "rerank_scores": [r.rerank_score for r in ranked],
    }