src/reader.py

#!/usr/bin/env python3
"""
The extractive reader: find the answer SPAN inside retrieved passages.

WHY THIS COMPONENT CARRIES THE LATENCY PLAN
-------------------------------------------
src/extractability.py measured that 58.6% of answerable queries have their answer
present verbatim (exact or token-subsequence) in the gold passage. Those queries
never need a language model -- copy the span and you are done in single-digit
milliseconds, on a MIG slice where generating 64 tokens costs 640-1000 ms.

The reader also produces the CONFIDENCE that src/router.py routes on, and that
src/router.calibrate_threshold turns into a guaranteed error bound. One score,
three jobs: answer, route, abstain.

TWO IMPLEMENTATIONS
-------------------
LexicalSpanReader   no model, no GPU, ~0.3 ms. IDF-weighted overlap + answer-type
                    priors. Works today, and is the honest baseline the neural
                    reader must beat.
NeuralSpanReader    bge-m3 encoder + start/end heads, trained by weak supervision
                    (src/span_labels.py locates `Answer` inside
                    `Translated_passages`). Loads a checkpoint if one exists and
                    falls back to lexical otherwise, so the pipeline never breaks
                    because training has not finished.

CALIBRATION MATTERS MORE THAN ACCURACY HERE
-------------------------------------------
Coverage at a fixed error bound is set by how well confidence separates right from
wrong, not by raw accuracy. A reader that is 70% accurate and well calibrated
beats one that is 75% accurate and badly calibrated, because the calibrated one
can answer far more queries at the same guaranteed error rate. So the score is
built from a MARGIN (best vs runner-up), which is what separates, rather than an
absolute similarity, which is not.
"""
from __future__ import annotations

import math
import re
import unicodedata
from dataclasses import dataclass
from typing import Sequence

from src.chunkers.base import split_sentences, words_of
from src.router import Passage, Span
from src.textnorm import normalise

__all__ = ["LexicalSpanReader", "NeuralSpanReader", "AnswerType", "detect_answer_type"]

_WS = re.compile(r"\s+")
# Digits across the scripts in MSMARCO-XI: ASCII, Devanagari, Bengali, Gurmukhi,
# Gujarati, Odia, Tamil, Telugu, Kannada, Malayalam, Arabic-Indic (Urdu).
_DIGITS = re.compile(r"[0-9०-९০-৯੦-੯૦-૯"
                     r"୦-୯௦-௯౦-౯೦-೯"
                     r"൦-൯٠-٩۰-۹]")


class AnswerType:
    NUMERIC = "NUMERIC"
    PERSON = "PERSON"
    LOCATION = "LOCATION"
    ENTITY = "ENTITY"
    DESCRIPTION = "DESCRIPTION"


# Cue words, English plus the Indic equivalents that survive translation.
_CUES = {
    AnswerType.NUMERIC: (
        "how many", "how much", "how long", "how old", "how far", "how fast",
        "how tall", "how deep", "how high", "how wide", "how heavy", "how big",
        "what year", "what age", "what speed", "what time", "what temperature",
        "what percentage", "what is the average", "when did", "when was", "when is",
        "cost", "price", "average", "percentage", "temperature", "salary", "rate",
        "कितन", "कितना", "कितने", "कब", "एत्तनै",
        "எத்தனை", "எவ்வளவு", "எப்போது", "എത്ര", "എപ്പോൾ",
        "ఎన్ని", "ఎంత", "ఎప్పుడు", "ಎಷ್ಟು", "ಯಾವಾಗ",
        "کتنا", "کتنے", "کب", "কত", "কখন", "કેટલા", "ਕਿੰਨਾ", "କେତେ"),
    AnswerType.PERSON: ("who is", "who was", "who did", "who are", "whose",
                        "कौन", "யார்", "ആര്", "ఎవరు", "ಯಾರು", "کون"),
    AnswerType.LOCATION: ("where is", "where are", "where did", "location of",
                          "कहां", "कहाँ", "எங்கே", "എവിടെ", "ఎక్కడ", "ಎಲ್ಲಿ", "کہاں"),
    AnswerType.ENTITY: ("what is the name", "which company", "which country",
                        "what type", "what kind"),
}


def detect_answer_type(query: str) -> str:
    """Cheap answer-type prior from the query text.

    NB: query_type exists in the dataset but NOT at inference time, so it has to
    be inferred. It is only a prior -- extractability showed the type spread is
    24 pp, too narrow to route on, so it nudges span scores rather than deciding
    anything.
    """
    q = query.lower()
    for t, cues in _CUES.items():
        if any(c in q for c in cues):
            return t
    return AnswerType.DESCRIPTION


_ES_STEMS = ("s", "x", "z", "ch", "sh")


_IND_PREFIX = 5


def _fold_indic(t: str) -> str:
    """Fold an Indic token to a prefix so inflected forms match.

    The comment below used to say exact match is enough because query and
    passage share a language. It is not: these morphologies are agglutinative,
    so a live Kannada query for ಕಾರ್ಪೊರೇಷನ್ met ಕಾರ್ಪೊರೇಷನ್‌ಗಳು in the passage
    and overlapped on nothing. Every candidate then scored 0, _confidence
    returned 0 for a best score of 0, and the reader shipped whichever sentence
    happened to sort first while reporting conf 0.00. It was not choosing.

    A prefix cannot end on a combining mark: Indic vowel signs and viramas are
    separate code points, and cutting before the base they attach to leaves a
    fragment that matches things it should not.
    """
    if len(t) <= _IND_PREFIX:
        return t
    import unicodedata

    p = t[:_IND_PREFIX]
    while p and unicodedata.category(p[-1]) in ("Mn", "Mc"):
        p = p[:-1]
    return p or t[:_IND_PREFIX]


def _stem(t: str) -> str:
    """Crude suffix strip for Latin script so 'eagle' matches 'eagles'.

    Order is load-bearing: stripping "es" before "s" turns 'eagles' into 'eagl'
    while 'eagle' stays 'eagle', so the two stop matching -- the exact bug this
    function exists to prevent. "es" is only a plural marker after s/x/z/ch/sh
    (boxes, dishes); everywhere else the plural is a bare "s".

    Deliberately not applied to Indic tokens: their morphology is agglutinative
    and chopping trailing characters destroys meaning. Query and passage are in
    the SAME language, so exact match already works there.
    """
    if not t.isascii():
        return _fold_indic(t)
    if len(t) < 4 or not t.isalpha():
        return t
    if t.endswith("ies") and len(t) > 4:
        return t[:-3] + "y"
    if t.endswith("es") and len(t) > 4 and t[:-2].endswith(_ES_STEMS):
        return t[:-2]
    if t.endswith("s") and not t.endswith("ss") and len(t) > 3:
        return t[:-1]
    if t.endswith("ing") and len(t) > 5:
        return t[:-3]
    if t.endswith("ed") and len(t) > 4:
        return t[:-2]
    return t


def _norm(s: str) -> str:
    # Shared with every metric in the repo. Was a `[^\w\s]` regex, which deleted
    # Indic vowel marks and split each word at the gap -- so this scorer was
    # matching consonant fragments, not words, in 13 of the 14 languages.
    return normalise(s, True)


def _norm_tokens(s: str) -> list[str]:
    return [_stem(t) for t in _norm(s).split()]


_STOP = {"the","a","an","of","is","are","was","were","to","in","for","on","and",
         "do","does","did","how","what","when","where","who","which","that","this",
         "it","its","be","been","by","with","as","at","from","or","not"}


def _idf(passages: Sequence[str]) -> dict[str, float]:
    """IDF over the retrieved set. Rare query terms should dominate the match;
    without this, stopwords decide the span."""
    n = max(1, len(passages))
    df: dict[str, int] = {}
    for p in passages:
        for t in set(_norm_tokens(p)):
            df[t] = df.get(t, 0) + 1
    return {t: math.log(1 + n / c) for t, c in df.items()}


# Median gold-answer length by type, measured on MSMARCO-XI (src/extractability.py).
_TARGET_WORDS = {
    AnswerType.NUMERIC: 9.0,
    AnswerType.LOCATION: 8.0,
    AnswerType.PERSON: 10.0,
    AnswerType.ENTITY: 12.0,
    AnswerType.DESCRIPTION: 16.0,
}


@dataclass(slots=True)
class _Cand:
    text: str
    score: float
    chunk_id: str
    start: int
    end: int


class LexicalSpanReader:
    """No model. Sentences scored by IDF-weighted query overlap, then a sub-span
    search inside the winner when the question wants a short answer.

    Confidence is the MARGIN between the best and runner-up spans, squashed to
    [0,1]. Margin separates correct from incorrect far better than an absolute
    score, which is what the conformal threshold needs.
    """
    name = "lexical"

    def __init__(self, max_span_words: int = 25, min_span_words: int = 1,
                 margin_scale: float = 6.0, prior_weight: float = 0.0,
                 answer_mode: str = "span"):
        self.max_span, self.min_span, self.scale = max_span_words, min_span_words, margin_scale
        # THE RETRIEVAL PRIOR. Passage.score carries the reranker's opinion about
        # which passage holds the answer, and this reader used to ignore it
        # completely -- every sentence in every candidate passage competed on
        # query overlap alone. Measured cost of ignoring it: the oracle over the
        # passage the reader CHOSE is 0.361 F1 while the oracle over ANY passage
        # is 0.715, i.e. half the achievable F1 is lost before span selection
        # even starts, in a candidate set where retrieval puts the gold passage
        # in the top 5 82% of the time.
        #
        # Blend multiplicatively on a [1-w, 1] ramp: at w=0 the prior is off, at
        # w=1 a passage the retriever scored 0 is fully suppressed. Multiplicative
        # rather than additive so a passage the retriever likes cannot rescue a
        # sentence with no lexical support at all.
        self.prior_weight = prior_weight

        # ANSWER MODE -- "span" trims to a sub-span of the winning sentence,
        # "sentence" returns the whole sentence.
        #
        # Two independent reasons "sentence" is the right default for SERVING:
        #
        #  1. IT MEASURES BETTER. The diagnosis put boundary loss at 0.120 F1 --
        #     oracle-within-the-chosen-sentence is 0.397 while the trimmed
        #     reader returns 0.276. The trimming is COSTING F1, because MS MARCO
        #     answers average 19.2 words, i.e. about a sentence. We were cutting
        #     sentences down to fragments to match targets that are sentences.
        #
        #  2. IT IS THE ONLY ONE THAT CAN BE SPOKEN. A trimmed span such as
        #     "को बीज से पके फल तक बढ़ने में ९०" starts mid-clause and stops
        #     mid-number. Through TTS that is not a slightly worse answer, it is
        #     an unusable one. Grounding is unaffected: a sentence of a retrieved
        #     passage is still a verbatim span of it, so guardrail gate 4 still
        #     passes by exact substring check.
        self.answer_mode = answer_mode if answer_mode in ("span", "sentence") else "span"

    # -- scoring ---------------------------------------------------------
    #
    # TWO DIFFERENT OBJECTIVES, and conflating them is the classic extractive-QA
    # mistake:
    #
    #   SENTENCE selection -> maximise QUERY OVERLAP. Find where the topic is
    #                         discussed.
    #   SPAN selection     -> the answer is the thing you DON'T know, so it
    #                         usually contains NONE of the query terms. Scoring
    #                         it by overlap picks the question back out of the
    #                         passage. Score by answer-type fit and proximity to
    #                         the query terms instead.

    def _score_tokens(self, span_tokens: list[str], q_tokens: set[str],
                      idf: dict[str, float], atype: str) -> float:
        if not span_tokens:
            return 0.0
        # Content words only: with stopwords included, every sentence matches
        # "how/does/the" equally and the length prior ends up choosing the answer.
        content = {t for t in q_tokens if t not in _STOP}
        if not content:
            content = q_tokens
        hit = sum(idf.get(t, 1.0) for t in set(span_tokens) & content)
        total = sum(idf.get(t, 1.0) for t in content) or 1.0
        cover = hit / total

        # Length prior: measured answers are ~12 words (median). Penalise spans
        # far from that, gently, so we neither return one word nor a paragraph.
        n = len(span_tokens)
        length_pen = math.exp(
            -((math.log(max(n, 1) / _TARGET_WORDS.get(atype, 12.0))) ** 2) / 2.0)

        prior = 1.0
        joined = " ".join(span_tokens)
        if atype == AnswerType.NUMERIC:
            prior = 1.35 if _DIGITS.search(joined) else 0.75
        elif atype in (AnswerType.PERSON, AnswerType.LOCATION, AnswerType.ENTITY):
            caps = sum(1 for t in span_tokens if t[:1].isupper())
            prior = 1.0 + 0.25 * min(caps, 3) / 3.0
        return cover * length_pen * prior

    def _score_subspan(self, span: list[str], span_at: int, sent: list[str],
                       q_tokens: set[str], atype: str) -> float:
        """Score a candidate answer span INSIDE the winning sentence."""
        n = len(span)
        if n == 0:
            return 0.0
        joined = " ".join(span)

        # 1. does it look like the requested answer type?
        if atype == AnswerType.NUMERIC:
            fit = 1.0 if _DIGITS.search(joined) else 0.05
        elif atype in (AnswerType.PERSON, AnswerType.LOCATION, AnswerType.ENTITY):
            caps = sum(1 for t in span if t[:1].isupper())
            fit = 0.20 + 0.80 * min(caps, 3) / 3.0
        else:
            fit = 1.0

        # 2. the answer is not the question -- penalise echoing query terms back
        overlap = len({_stem(_norm(t)) for t in span} & q_tokens) / n
        novelty = 1.0 - 0.70 * overlap

        # 3. answers sit near the query terms they answer
        anchors = [i for i, t in enumerate(sent) if _stem(_norm(t)) in q_tokens]
        if anchors:
            dist = min(abs(span_at - a) for a in anchors)
            prox = 1.0 / (1.0 + 0.12 * dist)
        else:
            prox = 0.6

        # 4. length prior, from the MEASURED median answer length per type
        #    (extractability.py over 42,000 real queries). Guessing "short answer
        #    = 4 words" produced spans like "fly 30 to 55" where the gold was
        #    "30 to 55 mph" -- the data says otherwise.
        target = _TARGET_WORDS.get(atype, 12.0)
        length = math.exp(-((math.log(max(n, 1) / target)) ** 2) / 2.0)

        # 5. phrase boundaries. Arbitrary word windows produce spans like
        #    "and dive at over" -- grammatical nonsense that tanks token-F1 and
        #    reads badly when spoken aloud. Real answers begin and end at phrase
        #    boundaries, so penalise spans that start or stop on a function word.
        boundary = 1.0
        if _norm(span[-1]) in _STOP:
            boundary *= 0.30
        if _norm(span[0]) in _STOP:
            boundary *= 0.60

        return fit * novelty * prox * length * boundary

    def read(self, query: str, passages: Sequence[Passage]) -> Span:
        if not passages:
            return Span("", 0.0, "")
        atype = detect_answer_type(query)
        texts = [p.text for p in passages]
        idf = _idf(texts)
        q_tokens = set(_norm_tokens(query))

        # Normalise the retrieval scores across THIS candidate set: reranker
        # scales differ per model and per query, so only the relative order is
        # meaningful.
        raw = [float(p.score) for p in passages]
        lo, hi = min(raw), max(raw)
        rng = hi - lo

        cands: list[_Cand] = []
        for p in passages:
            if self.prior_weight <= 0 or rng <= 0:
                prior_mult = 1.0
            else:
                rel = (float(p.score) - lo) / rng          # [0,1] within the set
                prior_mult = (1.0 - self.prior_weight) + self.prior_weight * rel
            words = words_of(p.text)
            norm_words = _norm_tokens(p.text)
            if len(norm_words) != len(words):        # normalisation split a token
                norm_words = [_stem(_norm(w)) or w.lower() for w in words]
            cursor = 0
            for sent in split_sentences(p.text):
                n = len(words_of(sent))
                if n == 0:
                    continue
                s, e = cursor, cursor + n
                cursor = e
                sc = self._score_tokens(norm_words[s:e], q_tokens, idf, atype) * prior_mult
                cands.append(_Cand(" ".join(words[s:e]), sc, p.chunk_id, s, e))

        if not cands:
            return Span("", 0.0, passages[0].chunk_id)
        cands.sort(key=lambda c: -c.score)
        best = cands[0]

        # For NUMERIC questions, prefer a competitive candidate that actually
        # contains a number. Query terms often do not appear near the figure
        # ("how fast does an eagle travel" vs "Eagles fly 30 to 55 mph"), so
        # pure overlap picks the wrong sentence.
        if atype == AnswerType.NUMERIC and not _DIGITS.search(best.text):
            withnum = [c for c in cands if _DIGITS.search(c.text)]
            if withnum and withnum[0].score >= 0.35 * best.score:
                best = withnum[0]

        # Short-answer types: tighten to a sub-span of the winning sentence.
        # MEASURED CONTEXT: gold answers average 10-18 words depending on type
        # (extractability.py), so a whole sentence is often already right. We only
        # tighten when a sub-span scores clearly better, and never for DESCRIPTION
        # -- 52% of traffic, mean answer 18.3 words, i.e. about a sentence.
        sentence_span = best
        if atype != AnswerType.DESCRIPTION and self.answer_mode == "span":
            p = next((x for x in passages if x.chunk_id == best.chunk_id), passages[0])
            words = words_of(p.text)[best.start:best.end]
            if 1 < len(words) <= 80:
                sub, sub_sc = None, 0.0
                for i in range(len(words)):
                    for L in range(self.min_span, min(self.max_span, len(words) - i) + 1):
                        sc = self._score_subspan(words[i:i + L], i, words, q_tokens, atype)
                        if sc > sub_sc:
                            sub_sc, sub = sc, (i, i + L)
                # Only accept a decisive sub-span; otherwise the sentence is safer.
                if sub and sub_sc >= 0.45:
                    i, j = sub
                    best = _Cand(" ".join(words[i:j]), best.score, best.chunk_id,
                                 best.start + i, best.start + j)

        # Confidence reflects SENTENCE selection -- did we find the right place --
        # because that is what the span's correctness actually depends on.
        runner = cands[1].score if len(cands) > 1 else 0.0
        conf = self._confidence(sentence_span.score, runner)
        return Span(best.text, conf, best.chunk_id, best.start, best.end)

    def _confidence(self, best: float, runner: float) -> float:
        """Margin -> [0,1]. Absolute scores are not comparable across queries;
        the gap to the runner-up is."""
        if best <= 0:
            return 0.0
        margin = (best - runner) / best          # relative margin in [0,1]
        strength = min(1.0, best)                 # did anything match at all
        return 1.0 / (1.0 + math.exp(-self.scale * (0.5 * margin + 0.5 * strength - 0.5)))

    # router.ReaderFn
    def __call__(self, query: str, passages: Sequence[Passage]) -> Span:
        return self.read(query, passages)


class NeuralSpanReader:
    """bge-m3 encoder + start/end heads. Falls back to lexical without a checkpoint.

    Kept deliberately thin: the router only needs (text, score, chunk_id), so the
    model can be swapped or trained later without touching anything downstream.
    """
    name = "neural"

    def __init__(self, model_path: str | None = None, checkpoint: str | None = None,
                 device: str = "cuda", max_len: int = 384,
                 fallback: LexicalSpanReader | None = None):
        self.fallback = fallback or LexicalSpanReader()
        self.model = None
        if not (model_path and checkpoint):
            return
        try:
            import torch
            from transformers import AutoModel, AutoTokenizer
            self.torch = torch
            self.tok = AutoTokenizer.from_pretrained(model_path)
            enc = AutoModel.from_pretrained(
                model_path, torch_dtype=torch.float16 if "cuda" in device else torch.float32)
            hidden = enc.config.hidden_size
            self.head = torch.nn.Linear(hidden, 2)
            state = torch.load(checkpoint, map_location="cpu")
            enc.load_state_dict(state["encoder"], strict=False)
            self.head.load_state_dict(state["head"])
            self.model = enc.to(device).eval()
            self.head = self.head.to(device).to(self.model.dtype).eval()
            self.device, self.max_len = device, max_len
        except Exception as exc:
            print(f"  reader: no checkpoint loaded ({exc}); using lexical fallback")
            self.model = None

    def read(self, query: str, passages: Sequence[Passage]) -> Span:
        if self.model is None or not passages:
            return self.fallback.read(query, passages)
        torch = self.torch
        best = None
        with torch.inference_mode():
            for p in passages[:4]:                      # deeper costs latency, not quality
                enc = self.tok(query, p.text, truncation="only_second",
                               max_length=self.max_len, return_offsets_mapping=True,
                               return_tensors="pt").to(self.device)
                offsets = enc.pop("offset_mapping")[0].tolist()
                h = self.model(**enc).last_hidden_state
                logits = self.head(h)[0].float()
                start_lp = torch.log_softmax(logits[:, 0], -1)
                end_lp = torch.log_softmax(logits[:, 1], -1)
                n = start_lp.shape[0]
                # best (i <= j) within a 40-token window
                scores = start_lp[:, None] + end_lp[None, :]
                mask = torch.triu(torch.ones(n, n, device=scores.device), 0) - \
                       torch.triu(torch.ones(n, n, device=scores.device), 41)
                scores = scores.masked_fill(mask <= 0, -1e9)
                flat = int(scores.argmax())
                i, j = flat // n, flat % n
                lp = float(scores.view(-1)[flat])
                if best is None or lp > best[0]:
                    a, b = offsets[i][0], offsets[j][1]
                    best = (lp, p.text[a:b], p.chunk_id, i, j)
        if best is None:
            return self.fallback.read(query, passages)
        lp, text, cid, i, j = best
        return Span(text.strip(), float(math.exp(min(0.0, lp))), cid, i, j)

    def __call__(self, query: str, passages: Sequence[Passage]) -> Span:
        return self.read(query, passages)