spb/wpai0002_uqo

1#!/usr/bin/env python32"""Decoding-side search (no training): for one checkpoint, generate greedy + K nucleus samples per validation pair, then score3candidate-selection rules against the originals — greedy, MBR with chrF++ utility (Eikema & Aziz 2020; Kumar & Byrne 2004),4MBR with the *source* added as a weighted pseudo-reference (identity anchoring), and an oracle (best candidate, upper bound).5Writes artifacts/metrics/<name>/decode_search.json and the selected outputs per rule (outputs_<system>__<rule>.jsonl) so they6can be scored by the regular evaluation runner. Exploratory: tuned on validation only; the chosen rule is then applied once to test."""7from __future__ import annotations89import argparse10import json11import sys12import time13from pathlib import Path1415import numpy as np16import torch1718ROOT = Path(__file__).resolve().parents[1]19sys.path.insert(0, str(ROOT / "src"))20from htr.evaluation.metrics import _chrf, factual  # noqa: E40221from htr.inference.pipeline import HTRReconstructor  # noqa: E40222from htr.parse.protect import protect, restore_placeholders  # noqa: E40223from htr.training.data import meta_prefix  # noqa: E40224from htr.data.taxonomy import era_of  # noqa: E40225from htr.utils.io import read_jsonl, write_json, write_jsonl  # noqa: E40226from htr.utils.provenance import run_header  # noqa: E402272829def chrf(h: str, r: str) -> float:30    return float(_chrf.sentence_score(h, [r]).score)313233def main() -> None:34    ap = argparse.ArgumentParser()35    ap.add_argument("--pairs", required=True)36    ap.add_argument("--splits", required=True)37    ap.add_argument("--split", default="val")38    ap.add_argument("--checkpoint", required=True)39    ap.add_argument("--system", required=True)40    ap.add_argument("--name", required=True, help="metrics dir name")41    ap.add_argument("--limit", type=int, default=1000)42    ap.add_argument("--k", type=int, default=8, help="samples per input")43    ap.add_argument("--temperature", type=float, default=0.7)44    ap.add_argument("--top-p", type=float, default=0.9)45    ap.add_argument("--batch-size", type=int, default=16)46    ap.add_argument("--device", default="auto")47    ap.add_argument("--seed", type=int, default=0)48    ap.add_argument("--out", default=None, help="predictions dir to write the selected outputs (optional)")49    a = ap.parse_args()50    rec = HTRReconstructor.from_pretrained(a.checkpoint, device=a.device)51    model, tok, dev = rec.model, rec.tok, rec.device52    splits = {s["doc_id"]: s["split"] for s in read_jsonl(a.splits)}53    pairs = [p for p in read_jsonl(a.pairs) if splits.get(p["doc_id"]) == a.split][: a.limit]54    print(f"{len(pairs)} pairs ({a.split}), K={a.k}, T={a.temperature}, p={a.top_p}", flush=True)55    torch.manual_seed(a.seed)56    pad = model.cfg.pad_id57    prepared = []58    for i, p in enumerate(pairs):59        m = p["metadata"]60        prot = protect(p["source_text"])61        src_text = (meta_prefix({"metadata": {**m, "era": m.get("era") or era_of(m.get("pub_year"))}}) if rec.text_meta_prefix else "") + prot.text62        prepared.append((i, p, prot, tok.encode(src_text, add_eos=True, max_len=model.cfg.max_len)))63    prepared.sort(key=lambda x: len(x[3]))64    cands: dict[int, list[str]] = {}65    t0 = time.time()66    for b in range(0, len(prepared), a.batch_size):67        chunk = prepared[b : b + a.batch_size]68        L = max(len(x[3]) for x in chunk)69        src = torch.tensor([x[3] + [pad] * (L - len(x[3])) for x in chunk], device=dev)70        meta_t = None71        if model.cfg.arch in ("m2", "m3"):72            parts = [rec._meta_tensors(x[1]["metadata"]) for x in chunk]73            meta_t = {k: torch.cat([q[k] for q in parts]) for k in parts[0]}74        mt = min(model.cfg.max_len - 2, int(L * 1.6) + 16)75        with torch.no_grad():76            g = model.generate(src, meta_t, max_new_tokens=mt, beam=1)77            srcK = src.repeat_interleave(a.k, 0)78            metaK = {k: v.repeat_interleave(a.k, 0) for k, v in meta_t.items()} if meta_t else None79            smp = model.generate(srcK, metaK, max_new_tokens=mt, beam=1, sample=True, temperature=a.temperature, top_p=a.top_p)80        for j, (i, p, prot, _) in enumerate(chunk):81            outs = [restore_placeholders(tok.decode(g[j].tolist()), prot)]82            outs += [restore_placeholders(tok.decode(smp[j * a.k + k].tolist()), prot) for k in range(a.k)]83            cands[i] = outs84        if (b // a.batch_size) % 10 == 0:85            print(f"  {b + len(chunk)}/{len(prepared)} ({time.time() - t0:.0f}s)", flush=True)86    # ---- selection rules ----87    rules = {"greedy": None, "mbr": 0.0, "mbr_src0.5": 0.5, "mbr_src1": 1.0, "mbr_src2": 2.0}88    sel: dict[str, list[str]] = {r: [] for r in rules}89    sel["identity"], sel["oracle"], sel["sample1"] = [], [], []90    per_rule_scores: dict[str, list[float]] = {r: [] for r in list(rules) + ["identity", "oracle", "sample1"]}91    nrec: dict[str, list[float]] = {r: [] for r in per_rule_scores}92    for i, p in enumerate(pairs):93        C = cands[i]94        src_t, tgt = p["source_text"], p["target_text"]95        n = len(C)96        U = np.zeros((n, n))97        for x in range(n):98            for y in range(n):99                if x != y:100                    U[x, y] = chrf(C[x], C[y])101        base = U.sum(1) / (n - 1)102        to_src = np.array([chrf(c, src_t) for c in C])103        vs_tgt = np.array([chrf(c, tgt) for c in C])104        picks = {"greedy": 0, "sample1": 1, "oracle": int(vs_tgt.argmax())}105        for r, w in rules.items():106            if w is not None:107                picks[r] = int((base + w * to_src).argmax())108        for r, k in picks.items():109            sel[r].append(C[k])110            per_rule_scores[r].append(float(vs_tgt[k]))111            nrec[r].append(factual(C[k], src_t)["number_recall"])112        sel["identity"].append(src_t)113        per_rule_scores["identity"].append(chrf(src_t, tgt))114        nrec["identity"].append(1.0)115    rep = {**run_header({"name": a.name, "checkpoint": a.checkpoint, "system": a.system, "split": a.split, "n": len(pairs), "k": a.k, "temperature": a.temperature, "top_p": a.top_p, "seed": a.seed}), "rules": {}}116    ident = np.array(per_rule_scores["identity"])117    for r, v in per_rule_scores.items():118        v = np.array(v)119        d = v - ident120        rep["rules"][r] = {"chrf_mean": float(v.mean()), "delta_vs_identity": float(d.mean()), "delta_se": float(d.std(ddof=1) / np.sqrt(len(d))), "number_recall": float(np.mean(nrec[r]))}121        print(f"{r:12s} chrF++ {v.mean():6.2f}  Δ vs identity {d.mean():+.2f} ± {1.96 * d.std(ddof=1) / np.sqrt(len(d)):.2f}  number recall {np.mean(nrec[r]):.3f}", flush=True)122    out_dir = ROOT / "artifacts/metrics" / a.name123    write_json(out_dir / "decode_search.json", rep)124    if a.out:125        for r in ("mbr", "mbr_src0.5", "mbr_src1", "mbr_src2"):126            write_jsonl(Path(a.out) / f"outputs_{a.system}__{r}.jsonl", [{"pair_id": p["pair_id"], "doc_id": p["doc_id"], "system": f"{a.system}__{r}", "output": sel[r][i]} for i, p in enumerate(pairs)])127    print("wrote", out_dir / "decode_search.json")128129130if __name__ == "__main__":131    main()132