Rerankers

"""Second-stage reranking with transparent query-document interaction features."""

import re
from typing import Dict, List, Sequence

from .models import SearchResult
from .text import content_terms, tokenize


class CrossFeatureReranker:
    """A deterministic stand-in for a neural cross-encoder.

    Unlike a bi-encoder, this scorer computes features over the query and
    candidate together: coverage, phrase presence, term proximity, title
    overlap, and matching years.  It is useful for teaching the pipeline
    boundary; production deployments should train or evaluate a proper
    cross-encoder/late-interaction model on their domain.
    """

    name = "cross-feature"

    def __init__(
        self,
        retrieval_weight: float = 0.30,
        coverage_weight: float = 0.38,
        phrase_weight: float = 0.12,
        proximity_weight: float = 0.10,
        title_weight: float = 0.05,
        year_weight: float = 0.05,
    ) -> None:
        self.retrieval_weight = retrieval_weight
        self.coverage_weight = coverage_weight
        self.phrase_weight = phrase_weight
        self.proximity_weight = proximity_weight
        self.title_weight = title_weight
        self.year_weight = year_weight

    @staticmethod
    def _proximity(query_terms: Sequence[str], document_terms: Sequence[str]) -> float:
        positions = [
            index for index, term in enumerate(document_terms) if term in set(query_terms)
        ]
        if len(set(document_terms) & set(query_terms)) < 2 or not positions:
            return 0.0
        span = max(positions) - min(positions) + 1
        return min(1.0, len(set(query_terms)) / span)

    def rerank(
        self, query: str, results: Sequence[SearchResult], k: int
    ) -> List[SearchResult]:
        if k <= 0 or not results:
            return []
        query_terms = content_terms(query)
        query_set = set(query_terms)
        query_phrase = " ".join(query_terms)
        query_years = set(re.findall(r"\b(?:19|20)\d{2}\b", query))
        raw_scores = [result.score for result in results]
        low, high = min(raw_scores), max(raw_scores)

        rescored = []
        for result in results:
            entities = result.chunk.metadata.get("entities", ())
            if isinstance(entities, str):
                entities = (entities,)
            identity = result.chunk.title + " " + " ".join(str(entity) for entity in entities)
            document_terms = content_terms(identity + " " + result.chunk.text)
            document_set = set(document_terms)
            title_set = set(content_terms(result.chunk.title))
            coverage = len(query_set & document_set) / max(1, len(query_set))
            phrase = float(bool(query_phrase) and query_phrase in " ".join(document_terms))
            proximity = self._proximity(query_terms, document_terms)
            title = len(query_set & title_set) / max(1, len(query_set))
            years = set(re.findall(r"\b(?:19|20)\d{2}\b", result.chunk.text))
            year = float(not query_years or bool(query_years & years))
            retrieval = (result.score - low) / (high - low) if high > low else 1.0
            score = (
                self.retrieval_weight * retrieval
                + self.coverage_weight * coverage
                + self.phrase_weight * phrase
                + self.proximity_weight * proximity
                + self.title_weight * title
                + self.year_weight * year
            )
            components: Dict[str, float] = dict(result.component_scores)
            components.update(
                {
                    "rerank_retrieval": retrieval,
                    "rerank_coverage": coverage,
                    "rerank_phrase": phrase,
                    "rerank_proximity": proximity,
                    "rerank_title": title,
                    "rerank_year": year,
                }
            )
            rescored.append((score, result, components))
        rescored.sort(key=lambda item: (-item[0], item[1].chunk.id))
        return [
            SearchResult(
                chunk=result.chunk,
                score=score,
                rank=rank,
                retriever=result.retriever + "+" + self.name,
                component_scores=components,
            )
            for rank, (score, result, components) in enumerate(rescored[:k], start=1)
        ]