natural-memory-nm21/fit_answerability_threshold.py

"""Fit the answerability margin threshold on train families, evaluate on eval.

`probe_answerability_cosine` swept thresholds on the eval set itself, which is
in-sample and therefore optimistic.  This script does it properly:

1. build the same POSITIVE / NEGATIVE episode variants the head training uses
   (positives = answering facts present; negatives = only records that provably
   belong to *other* attribute families), from `train.jsonl`;
2. pick the margin threshold that refuses the most unanswerable episodes while
   holding false refusals at or below a budget;
3. apply that frozen threshold to `eval.jsonl`, whose attribute families never
   appear in train.

The margin is ``cos(query, best record) - cos(query, second best record)`` -- a
training-free, open-vocabulary signal, which matters because the shipped gate is
a closed 24-class classifier that disables itself outside its vocabulary.
"""

from __future__ import annotations

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

import torch

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

from V2_dpskw.make_answerability_data import FeatureBank, attribute_of_question, write_set


def load_episodes(path: Path):
    with path.open("r", encoding="utf-8") as handle:
        for line in handle:
            line = line.strip()
            if line:
                yield json.loads(line)


def variants(episode: dict):
    """(records, label, kind) triples for one episode.

    Three populations, deliberately kept separate because they are different
    questions:

    * ``answerable``      -- the answering fact is in the bank (label 1)
    * ``unknown``         -- an unknown-attribute episode as constructed: the bank
                             holds only other-family facts (label 0).  This is the
                             deployment case the gate exists for.
    * ``answer_removed``  -- an answerable episode with its answer stripped out,
                             leaving same-shaped records (label 0).  A stress case:
                             the records are drawn from the same episodes, so the
                             margin is only informative if the signal is real.
    """
    candidates = [str(c.get("text", "")) for c in (episode.get("candidates") or []) if isinstance(c, dict)]
    positives = [candidates[i] for i in (episode.get("positive_indices") or []) if 0 <= int(i) < len(candidates)]
    out = []
    if positives:
        out.append((write_set(episode, include_answer=True), 1, "answerable"))
        neg = write_set(episode, include_answer=False)
        if neg:
            out.append((neg, 0, "answer_removed"))
    else:
        records = write_set(episode, include_answer=True)
        if records:
            out.append((records, 0, "unknown"))
    return out


def margin_of(bank: FeatureBank, query: str, records: list[str]) -> float | None:
    qrow = bank.row(query)
    rows = [bank.row(t) for t in records]
    if qrow is None or len(rows) < 2 or any(r is None for r in rows):
        return None
    q = bank.key(qrow)
    q = q / (q.norm() + 1e-6)
    R = torch.stack([bank.key(r) for r in rows])
    R = R / (R.norm(dim=-1, keepdim=True) + 1e-6)
    sims = torch.sort(R @ q, descending=True).values
    return float(sims[0] - sims[1])


def collect(bank: FeatureBank, corpus: Path):
    margins, labels, kinds, cats = [], [], [], []
    for episode in load_episodes(corpus):
        category = str((episode.get("metadata") or {}).get("category", ""))
        for records, label, kind in variants(episode):
            m = margin_of(bank, str(episode.get("query", "")), records)
            if m is None:
                continue
            margins.append(m)
            labels.append(label)
            kinds.append(kind)
            cats.append(category)
    return margins, labels, kinds, cats


def auc(scores, labels) -> float:
    pairs = sorted(zip(scores, labels))
    n_pos = sum(labels)
    n_neg = len(labels) - n_pos
    if not n_pos or not n_neg:
        return float("nan")
    rank_sum_pos = 0.0
    i = 0
    while i < len(pairs):
        j = i
        while j < len(pairs) and pairs[j][0] == pairs[i][0]:
            j += 1
        avg_rank = (i + j - 1) / 2.0 + 1
        rank_sum_pos += avg_rank * sum(1 for k in range(i, j) if pairs[k][1])
        i = j
    return (rank_sum_pos - n_pos * (n_pos + 1) / 2.0) / (n_pos * n_neg)


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--bank", type=Path, default=Path(r"H:\Memory\nm_cache\nm_realistic_v2\feature_cache"))
    parser.add_argument("--train", type=Path, default=Path("data/realistic_v2/train.jsonl"))
    parser.add_argument("--eval", type=Path, default=Path("data/realistic_v2/eval.jsonl"))
    parser.add_argument("--false-refusal-budget", type=float, default=0.04)
    parser.add_argument("--json-out", type=Path, default=None)
    args = parser.parse_args()

    bank = FeatureBank(args.bank)
    tr_m, tr_l, tr_k, tr_c = collect(bank, args.train)
    ev_m, ev_l, ev_k, ev_c = collect(bank, args.eval)
    for name, m, l, k in (("train", tr_m, tr_l, tr_k), ("eval", ev_m, ev_l, ev_k)):
        counts = defaultdict(int)
        for kind in k:
            counts[kind] += 1
        print(f"{name}: {len(l)} variants {dict(counts)}")

    # AUC against the deployment population only: can the margin tell an
    # answerable episode from an unknown-attribute one?
    def sub(m, l, k, kinds):
        return ([x for x, kk in zip(m, k) if kk in kinds],
                [x for x, kk in zip(l, k) if kk in kinds])

    tr_md, tr_ld = sub(tr_m, tr_l, tr_k, {"answerable", "unknown"})
    ev_md, ev_ld = sub(ev_m, ev_l, ev_k, {"answerable", "unknown"})
    print(f"\nAUC (answerable vs unknown)  train: {auc(tr_md, tr_ld):.4f}   eval: {auc(ev_md, ev_ld):.4f}")
    tr_ms, tr_ls = sub(tr_m, tr_l, tr_k, {"answerable", "answer_removed"})
    ev_ms, ev_ls = sub(ev_m, ev_l, ev_k, {"answerable", "answer_removed"})
    print(f"AUC (answerable vs answer-removed)  train: {auc(tr_ms, tr_ls):.4f}   eval: {auc(ev_ms, ev_ld and ev_ls):.4f}")

    budget = args.false_refusal_budget
    tr_ans = [m for m, l in zip(tr_m, tr_l) if l == 1]
    tr_unk = [m for m, l, k in zip(tr_m, tr_l, tr_k) if k == "unknown"]
    best = None
    for step in range(0, 400):
        t = step / 1000.0
        fr = sum(1 for m in tr_ans if m < t) / max(1, len(tr_ans))
        ru = sum(1 for m in tr_unk if m < t) / max(1, len(tr_unk))
        if fr <= budget and (best is None or ru > best[1]):
            best = (t, ru, fr)
    if best is None:
        print("no threshold satisfies the false-refusal budget on train")
        return 1
    threshold, train_ru, train_fr = best
    print(f"\nfitted on TRAIN: margin >= {threshold:.3f}"
          f"  -> refused_unanswerable={train_ru * 100:.2f}%  false_refusal={train_fr * 100:.2f}%")

    ev_ans = [(m, c) for m, l, c in zip(ev_m, ev_l, ev_c) if l == 1]
    ev_unk = [(m, c) for m, l, k, c in zip(ev_m, ev_l, ev_k, ev_c) if k == "unknown"]
    ev_rem = [(m, c) for m, l, k, c in zip(ev_m, ev_l, ev_k, ev_c) if k == "answer_removed"]
    ru = sum(1 for m, _ in ev_unk if m < threshold) / max(1, len(ev_unk))
    fr = sum(1 for m, _ in ev_ans if m < threshold) / max(1, len(ev_ans))
    rr = sum(1 for m, _ in ev_rem if m < threshold) / max(1, len(ev_rem))
    print(f"applied to  EVAL: unknown refused={ru * 100:.2f}% ({len(ev_unk)})  "
          f"false_refusal={fr * 100:.2f}% ({len(ev_ans)})  answer_removed refused={rr * 100:.2f}% ({len(ev_rem)})")

    per_unknown = defaultdict(int)
    for m, c in ev_unk:
        if m < threshold:
            per_unknown[c] += 1
    print(f"\nrefused-by-category {dict(per_unknown)}")
    per_ans_refused = defaultdict(int)
    per_ans_total = defaultdict(int)
    for m, c in ev_ans:
        per_ans_total[c] += 1
        if m < threshold:
            per_ans_refused[c] += 1
    print(f"{'answerable category':<22}{'refused':>9}{'total':>7}")
    for c in sorted(per_ans_total):
        print(f"{c:<22}{per_ans_refused[c]:>9}{per_ans_total[c]:>7}")

    if args.json_out:
        args.json_out.write_text(json.dumps({
            "threshold": threshold,
            "train_auc": round(auc(tr_m, tr_l), 4),
            "eval_auc": round(auc(ev_m, ev_l), 4),
            "train_refused_unanswerable_pct": round(train_ru * 100, 2),
            "train_false_refusal_pct": round(train_fr * 100, 2),
            "eval_refused_unanswerable_pct": round(ru * 100, 2),
            "eval_false_refusal_pct": round(fr * 100, 2),
            "eval_refused_by_category": dict(per_unknown),
            "eval_answerable_refused": {c: [per_ans_refused[c], per_ans_total[c]] for c in sorted(per_ans_total)},
        }, ensure_ascii=False, indent=2), encoding="utf-8")
        print(f"\nwrote {args.json_out}")
    return 0


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