src/extractability.py

#!/usr/bin/env python3
"""
Is the answer actually a SPAN in a retrieved passage?

WHY THIS DECIDES THE ARCHITECTURE
---------------------------------
On a MIG 2g.35gb slice (32 GB, 32 SMs), free-form generation of a 64-token answer
costs 640-1000 ms. The 200 ms budget is only reachable if most answers can be
EXTRACTED from a passage (~5-10 ms) rather than generated.

But MS MARCO answers are human-written and often ABSTRACTIVE -- a rewrite of the
passage, not a copy of it. Whether extraction works is an empirical question about
this corpus, not a design preference. This script answers it before we build.

FOUR MATCH LEVELS, strictest first:
  exact      normalised answer is a literal substring of the passage
  subseq     every answer token appears in the passage, in order (allows
             "$9,438" vs "$ 9,438", inserted words)
  overlap80  >= 80% of answer content tokens appear anywhere in the passage
  none       abstractive -- must be generated

DECISION RULE
  < 30% strict           ->  extraction not viable; the budget must come from a
                             smaller model + speculative decoding, and TTFT must
                             be reported separately from full completion
  >= 30% strict          ->  extraction is worth building. HOW to route is then
                             decided by whether query_type is discriminative:
                               * wide spread AND the biggest class is not the most
                                 extractive  ->  route on query_type
                               * otherwise                ->  route on the READER'S
                                 CONFIDENCE, because a type router would push the
                                 largest, most-extractable share of traffic down
                                 the slow generative path

MEASURED ON MSMARCO-XI (n=42,000): strict 58.6%, loose 88.9%.
DESCRIPTION -- 52% of answerable traffic -- is the MOST extractive type (62.7%),
not the least. Type-based routing is therefore the wrong axis here.

  python src/extractability.py --per-lang 3000
"""
from __future__ import annotations

import argparse
import json
import sys
from collections import Counter, defaultdict
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from src.schema_utils import default_root, iter_passages, norm_lang, load_report  # noqa: E402
# Re-exported so existing importers keep working. See src/textnorm.py for why
# this is NOT a `[^\w\s]` regex: that form deletes Indic vowel marks and splits
# every word at the gap, which silently corrupted every token-level metric here.
from src.textnorm import normalise, strip_punct  # noqa: E402,F401

NO_ANSWER = "no answer present"


def is_subsequence(needle: list[str], hay: list[str]) -> bool:
    it = iter(hay)
    return all(tok in it for tok in needle)


def classify(answer: str, passage: str) -> str:
    """Strictest matching level that holds."""
    a_n, p_n = normalise(answer), normalise(passage)
    if a_n and a_n in p_n:
        return "exact"

    a_p, p_p = normalise(answer, True), normalise(passage, True)
    if a_p and a_p in p_p:
        return "exact"

    a_tok, p_tok = a_p.split(), p_p.split()
    if not a_tok:
        return "none"
    if is_subsequence(a_tok, p_tok):
        return "subseq"

    p_set = set(p_tok)
    hits = sum(1 for t in a_tok if t in p_set)
    return "overlap80" if hits / len(a_tok) >= 0.80 else "none"


# WHICH PAIR? This is the decisive question for a multilingual system.
#   english     Eng_Answer  vs English_passages     (what we measured first)
#   translated  Answer      vs Translated_passages  (what the reader ACTUALLY sees)
# Translation is not word-for-word, so verbatim overlap can survive in English and
# vanish in the target language -- which would make extract-first viable for
# English and hopeless for the other 14. The reader runs on the TRANSLATED pair,
# so the translated number is the one that governs the architecture.
PAIRS = {
    # pair    -> (answer logical field, passage logical key)
    "english": ("answer_en", "text_en"),
    "translated": ("answer", "text"),
}


def run(root: Path, per_lang: int, langs_wanted: set[str] | None, pair: str = "english"):
    import polars as pl

    rep = load_report(root)
    fmap, pmap = rep["field_mapping"], rep["passage_mapping"]
    pcol = fmap["passages"]
    qid_c, lang_c, qt_c = fmap["query_id"], fmap.get("lang"), fmap.get("query_type")

    ans_field, psg_key = PAIRS[pair]
    ans_c = fmap.get(ans_field)
    p_key = pmap.get(psg_key)
    # the query column is only used to print readable examples
    q_c = fmap.get("query_en") if pair == "english" else fmap.get("query")
    sel_key = pmap.get("is_selected")

    if not (ans_c and p_key):
        raise SystemExit(f"pair '{pair}' needs fields {ans_field!r} and passage key "
                         f"{psg_key!r}; schema_report has {ans_c!r} / {p_key!r}")

    files = [f for f in rep["files"] if "val" in Path(f).name]
    cols = [c for c in (qid_c, lang_c, qt_c, q_c, ans_c, pcol) if c]
    cols = list(dict.fromkeys(cols))

    by_level: Counter = Counter()
    by_type: dict[str, Counter] = defaultdict(Counter)
    by_lang: dict[str, Counter] = defaultdict(Counter)
    gold_vs_any: Counter = Counter()
    examples: dict[str, list] = defaultdict(list)
    n_checked = 0

    for fp in files:
        lang_guess = norm_lang(Path(fp).name[:3])
        if langs_wanted and lang_guess not in langs_wanted:
            continue
        try:
            df = pl.read_parquet(fp, columns=cols, n_rows=per_lang * 3)
        except Exception as exc:
            print(f"  skip {Path(fp).name}: {exc}")
            continue

        lang = norm_lang(df[lang_c][0]) if lang_c and len(df) else lang_guess
        get = lambda c: df[c].to_list() if c and c in df.columns else [None] * len(df)  # noqa: E731
        kept = 0

        for q, ans, qt, plist in zip(get(q_c), get(ans_c), get(qt_c), df[pcol].to_list()):
            if kept >= per_lang:
                break
            if not isinstance(ans, str) or not ans.strip():
                continue
            # the no-answer sentinel is stored in English even in translated rows,
            # but the translated view may carry a rendered version -- drop both.
            if ans.strip().lower().startswith(NO_ANSWER):
                continue

            gold, others = [], []
            for _i, txt, txt_en, sel, _u in iter_passages(plist, pmap["text"],
                                                          pmap.get("text_en"), sel_key, None):
                use = txt_en if psg_key == "text_en" else txt
                if not isinstance(use, str) or not use.strip():
                    continue
                (gold if sel == 1 else others).append(use)
            if not gold:
                continue

            best_gold = min((classify(ans, p) for p in gold),
                            key=lambda L: ["exact", "subseq", "overlap80", "none"].index(L))
            best_any = min((classify(ans, p) for p in gold + others),
                           key=lambda L: ["exact", "subseq", "overlap80", "none"].index(L))

            by_level[best_gold] += 1
            by_lang[lang][best_gold] += 1
            if qt:
                by_type[qt][best_gold] += 1
            gold_vs_any[(best_gold != "none", best_any != "none")] += 1
            if len(examples[best_gold]) < 3:
                examples[best_gold].append((q, ans, gold[0][:180]))
            kept += 1
            n_checked += 1

        print(f"  {Path(fp).name:22s} {lang:3s} checked {kept:,}")

    return by_level, by_type, by_lang, gold_vs_any, examples, n_checked


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--root", type=Path, default=None)
    ap.add_argument("--per-lang", type=int, default=3000)
    ap.add_argument("--langs", default=None, help="comma-separated ISO-2; default all")
    ap.add_argument("--pair", choices=sorted(PAIRS), default="english",
                    help="english = Eng_Answer vs English_passages (reference); "
                         "translated = Answer vs Translated_passages (what the reader sees)")
    args = ap.parse_args()

    root = args.root.expanduser().resolve() if args.root else default_root()
    print(f"==> data root: {root}")
    print(f"==> pair     : {args.pair}")
    wanted = set(args.langs.split(",")) if args.langs else None

    ans_field, psg_key = PAIRS[args.pair]
    print(f"\n==> testing whether {ans_field} is a span of a gold {psg_key} passage")
    by_level, by_type, by_lang, gold_vs_any, examples, n = run(
        root, args.per_lang, wanted, args.pair)
    if not n:
        raise SystemExit("nothing checked — is the validation split present?")

    order = ["exact", "subseq", "overlap80", "none"]
    print(f"\n{'='*64}\nMATCH LEVEL vs the GOLD passage   (n={n:,})\n{'='*64}")
    cum = 0
    for lvl in order:
        c = by_level[lvl]
        cum += c
        print(f"  {lvl:10s} {c:>8,}  {100*c/n:5.1f}%   cumulative {100*cum/n:5.1f}%")

    extractive = by_level["exact"] + by_level["subseq"]
    loose = extractive + by_level["overlap80"]
    print(f"\n  STRICT  (exact+subseq)          : {100*extractive/n:5.1f}%")
    print(f"  LOOSE   (+overlap80)            : {100*loose/n:5.1f}%")

    print(f"\n{'='*64}\nBY LANGUAGE (strict extractive share)\n{'='*64}")
    print(f"  {'lang':6s}{'n':>8}{'exact':>8}{'subseq':>8}{'strict%':>9}{'loose%':>8}")
    print("  " + "-"*47)
    lang_strict: dict[str, float] = {}
    for lg in sorted(by_lang):
        c = by_lang[lg]
        tot = sum(c.values())
        if not tot:
            continue
        s = c["exact"] + c["subseq"]
        lo = s + c["overlap80"]
        lang_strict[lg] = 100 * s / tot
        print(f"  {lg:6s}{tot:>8,}{c['exact']:>8,}{c['subseq']:>8,}"
              f"{100*s/tot:>8.1f}%{100*lo/tot:>7.1f}%")
    if len(lang_strict) >= 2:
        lo_lg = min(lang_strict, key=lang_strict.get)
        hi_lg = max(lang_strict, key=lang_strict.get)
        print(f"\n  spread across languages: {lang_strict[hi_lg]-lang_strict[lo_lg]:.1f} pp "
              f"({lo_lg} {lang_strict[lo_lg]:.1f}% .. {hi_lg} {lang_strict[hi_lg]:.1f}%)")

    print(f"\n{'='*64}\nBY QUERY TYPE (strict extractive share)\n{'='*64}")
    print(f"  {'type':14s}{'n':>8}{'exact':>8}{'subseq':>8}{'strict%':>9}{'loose%':>8}")
    print("  " + "-"*55)
    for qt in sorted(by_type, key=lambda k: -sum(by_type[k].values())):
        c = by_type[qt]
        tot = sum(c.values())
        s = c["exact"] + c["subseq"]
        l = s + c["overlap80"]
        print(f"  {qt:14s}{tot:>8,}{c['exact']:>8,}{c['subseq']:>8,}"
              f"{100*s/tot:>8.1f}%{100*l/tot:>7.1f}%")

    print(f"\n{'='*64}\nEXAMPLES\n{'='*64}")
    for lvl in order:
        for qen, ans, psg in examples[lvl][:2]:
            print(f"\n  [{lvl}]")
            print(f"    Q: {str(qen)[:80]}")
            print(f"    A: {str(ans)[:100]}")
            print(f"    P: {psg[:150]}...")

    print(f"\n{'='*64}\nVERDICT\n{'='*64}")
    pct = 100 * extractive / n

    # Is query_type actually discriminative enough to route on? Check before
    # recommending it -- on MS MARCO the biggest class is also the MOST
    # extractive, which makes type-based routing actively harmful.
    rates = {qt: (c["exact"] + c["subseq"]) / max(1, sum(c.values()))
             for qt, c in by_type.items()}
    sizes = {qt: sum(c.values()) for qt, c in by_type.items()}
    routable = False
    if len(rates) >= 2:
        biggest = max(sizes, key=sizes.get)
        spread = 100 * (max(rates.values()) - min(rates.values()))
        biggest_is_best = rates[biggest] == max(rates.values())
        routable = spread >= 30 and not biggest_is_best
        print(f"  query_type spread : {spread:.1f} pp "
              f"({min(rates, key=rates.get)} {100*min(rates.values()):.1f}% .. "
              f"{max(rates, key=rates.get)} {100*max(rates.values()):.1f}%)")
        print(f"  largest class     : {biggest} ({100*sizes[biggest]/n:.0f}% of sample), "
              f"extractive {100*rates[biggest]:.1f}%"
              + ("  <- also the BEST" if biggest_is_best else ""))

    print(f"\n  strict extractive : {pct:.1f}%")
    print(f"  loose  (+overlap) : {100*loose/n:.1f}%")

    if pct < 30:
        print("\n  -> EXTRACTION NOT VIABLE. Answers are abstractive. Budget must come\n"
              "     from a smaller model + speculative decoding; report TTFT separately.")
    elif routable:
        print("\n  -> HYBRID ROUTER ON query_type is justified: the spread is wide and\n"
              "     the largest class is not the most extractive.")
    else:
        print("\n  -> ROUTE ON EXTRACTION CONFIDENCE, NOT query_type.")
        print("     The type spread is too narrow, and/or the biggest class is also the\n"
              "     most extractive -- so a type router would send the largest, most\n"
              "     extractable share of traffic down the SLOW generative path.")
        print("     Instead: always run the reader, accept its span when the score clears\n"
              "     a threshold, else fall back. Same machinery as the abstention head,\n"
              "     and the threshold is calibrated on the calib split.")
        if 100 * loose / n - pct > 20:
            print(f"\n     NOTE: {100*(loose-extractive)/n:.1f}% land in overlap80 -- the answer tokens ARE\n"
                  "     in the passage but reordered. Those are recoverable by a SHORT\n"
                  "     constrained generation over passage vocabulary, not free-form\n"
                  "     decoding. Cap generation length near the observed answer median.")

    if args.pair == "translated":
        print("\n  THIS IS THE NUMBER THAT GOVERNS THE READER. The reader scores spans of\n"
              "  Translated_passages against Answer. If this is far below the English\n"
              "  figure, the lexical reader is not underperforming -- it is being asked\n"
              "  for a span that does not exist, and no amount of tuning recovers it.")

    out = root / "results" / f"extractability_{args.pair}.json"
    out.parent.mkdir(parents=True, exist_ok=True)
    out.write_text(json.dumps({
        "pair": args.pair,
        "answer_field": ans_field,
        "passage_key": psg_key,
        "n_checked": n,
        "by_level": dict(by_level),
        "by_query_type": {k: dict(v) for k, v in by_type.items()},
        "by_language": {k: dict(v) for k, v in by_lang.items()},
        "strict_pct_by_language": {k: round(v, 2) for k, v in lang_strict.items()},
        "strict_pct": round(pct, 2),
        "loose_pct": round(100 * loose / n, 2),
    }, indent=2, ensure_ascii=False))
    print(f"\n==> wrote {out}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())