rerank.py

"""
EpiRAG - rerank.py
------------------
Cross-encoder reranking using sentence-transformers.
Loads a lightweight cross-encoder (ms-marco-MiniLM-L-6-v2) and reranks
the top-K chunks returned by vector search / hybrid retrieval by true
relevance scores, significantly improving answer quality.
"""

from __future__ import annotations

_cross_encoder = None

CROSS_ENCODER_MODEL = "cross-encoder/ms-marco-MiniLM-L-6-v2"


def _get_cross_encoder():
    global _cross_encoder
    if _cross_encoder is None:
        try:
            from sentence_transformers import CrossEncoder
            print(f"  [Rerank] Loading cross-encoder: {CROSS_ENCODER_MODEL}", flush=True)
            _cross_encoder = CrossEncoder(CROSS_ENCODER_MODEL)
            print("  [Rerank] Cross-encoder ready.", flush=True)
        except Exception as e:
            print(f"  [Rerank] Could not load cross-encoder ({e}). Skipping rerank.", flush=True)
            _cross_encoder = False   # sentinel - don't retry
    return _cross_encoder if _cross_encoder else None


def rerank(query: str, chunks: list[dict], top_n: int | None = None) -> list[dict]:
    """
    Rerank `chunks` using a cross-encoder given `query`.
    Falls back gracefully (returns original order) if model unavailable.

    Args:
        query:   The user's search query.
        chunks:  List of chunk dicts (must have a 'text' key).
        top_n:   If provided, only return the top N chunks after reranking.

    Returns:
        Reranked list of chunk dicts (descending relevance).
    """
    if not chunks:
        return chunks

    ce = _get_cross_encoder()
    if ce is None:
        return chunks

    try:
        pairs  = [(query, c["text"]) for c in chunks]
        scores = ce.predict(pairs).tolist()

        # Attach cross-encoder score and sort descending
        for chunk, score in zip(chunks, scores):
            chunk["ce_score"] = round(float(score), 4)

        ranked = sorted(chunks, key=lambda c: c.get("ce_score", 0.0), reverse=True)
        return ranked[:top_n] if top_n else ranked

    except Exception as e:
        print(f"  [Rerank] Failed ({e}), returning original order.", flush=True)
        return chunks