validate_gold.py

#!/usr/bin/env python3
"""
validate_gold.py — measure lemma accuracy against a hand-annotated gold treebank.

Why this exists.  Every check on the disambiguation so far has routed through
either the system's own output or a hand audit of 50 flips.  This measures it
against independent scholarly annotation, at a scale no hand audit reaches, and
without needing anyone to read Greek.

Gold sources (both hand-annotated, both derived from the same Perseus editions
this project parses, so the texts genuinely correspond):

  AGDT 2.1   github.com/PerseusDL/treebank_data  (v2.1/Greek/texts/*.tb.xml)
             per-token CTS URN in @cite.  Plato Euthyphro; all of Homer, Hesiod,
             Aeschylus, Sophocles — the verse gold.
  Gorman     github.com/perseids-publications/gorman-trees  (public/xml/*.xml)
             @subdoc on <sentence> carries the citation (e.g. "17a").
             Plato Apology; Aristotle Politics 1-2.

What it reports, following the convention in the Ancient Greek lemmatization
literature (Vatri & McGillivray 2020; Kostkan et al. 2023): top-1 exact-match
lemma accuracy, broken out for ALL tokens, AMBIGUOUS tokens (>1 candidate lemma
in the cache — the only ones disambiguation can affect), and UNKNOWN tokens
(absent from the cache).  With --base it also reports accuracy AFTER the
facing-translation re-ranking, so the gain is measured on gold, not on itself.

    python -u validate_gold.py --gold gold/plato-apology.xml --work apology
    python -u validate_gold.py --gold gold/plato-apology.xml --work apology \
                               --base http://127.0.0.1:8000
    python -u validate_gold.py --gold gold/tlg0059.tlg001...tb.xml --work euthyphro

Alignment is per citation unit (Stephanus section / Bekker column / book.line),
then a sequence match on accent-stripped forms inside the unit.  The alignment
rate is printed first and is the gate: below ~90% the tokenizers diverge
materially (usually elision or crasis) and the accuracy numbers are not
trustworthy until that is fixed.
"""

# Windows consoles default to cp1252, which raises UnicodeEncodeError the moment
# this script prints a Greek lemma.  Force UTF-8 on stdout/stderr rather than
# relying on the caller to set PYTHONUTF8.
import sys as _sys
for _s in (_sys.stdout, _sys.stderr):
    try:
        if (getattr(_s, "encoding", "") or "").lower().replace("-", "") != "utf8":
            _s.reconfigure(encoding="utf-8")
    except Exception:
        pass

import os
import re
import sys
import json
import argparse
import difflib
import unicodedata as ud
from collections import defaultdict, Counter

APOS = "\u02bc\u2019'\u1fbd\u0313\u00b4"


def norm_form(s):
    """Compare surface forms ignoring accents, breathings, case and apostrophes —
    the things the two tokenizations disagree about."""
    s = ud.normalize("NFD", s or "")
    s = "".join(c for c in s if not ud.combining(c))
    s = "".join(c for c in s if c not in APOS)
    return s.lower().strip()


def norm_lemma(s):
    """Lemma comparison key: NFC, strip the homograph digit suffix (λέγω1 → λέγω),
    drop accents.  Accents are dropped because gold and Morpheus disagree on
    precomposed oxia vs tonos and on where to accent a citation form."""
    s = (s or "").strip()
    s = re.sub(r"\d+$", "", s)
    s = ud.normalize("NFD", s)
    s = "".join(c for c in s if not ud.combining(c))
    return s.lower()


# --- gold parsing -----------------------------------------------------------
def parse_gold(path):
    """[(locus, form, lemma, postag)] for word tokens, punctuation dropped."""
    x = open(path, encoding="utf-8", errors="replace").read()
    out = []
    # split into sentences so we can read @subdoc (Gorman) per sentence
    for sm in re.finditer(r"<sentence\b([^>]*)>(.*?)</sentence>", x, re.S):
        attrs, body = sm.group(1), sm.group(2)
        sub = re.search(r'subdoc="([^"]*)"', attrs)
        sent_locus = sub.group(1).strip() if sub else ""
        for wm in re.finditer(r"<word\b([^>]*?)/?>", body):
            a = wm.group(1)
            def at(n):
                m = re.search(n + r'="([^"]*)"', a)
                return m.group(1) if m else ""
            postag = at("postag")
            if postag.startswith("u"):            # punctuation
                continue
            form, lemma = at("form"), at("lemma")
            if not form or lemma.startswith("punc"):
                continue
            # AGDT: per-token CTS URN, passage after the final colon
            cite = at("cite")
            locus = sent_locus
            if cite and ":" in cite:
                locus = cite.rsplit(":", 1)[1].strip() or sent_locus
            out.append((locus, form, lemma, postag))
    return out


def load_work(data_dir, slug):
    """[(locus, seg_id, tok_i, surface, morph_key)] for word tokens."""
    doc = json.load(open(os.path.join(data_dir, f"{slug}.json"), encoding="utf-8"))
    out = []
    for s in doc["segments"]:
        for t in s["tokens"]:
            if t.get("kind") == "word":
                out.append((s["locus"], s["id"], t["i"], t["t"], t.get("w")))
    return out, doc


# --- alignment --------------------------------------------------------------
def align(mine, gold, mode="auto"):
    """Match my tokens to gold tokens.

    Two strategies, because which one wins is not predictable in advance:
      per-citation — group by Stephanus/Bekker/book.line, match inside each unit.
                     Fails when the two texts disagree about where a section
                     starts (Gorman's Apology sections are offset from Perseus's,
                     and Gorman writes "21α" with a Greek alpha).
      global       — one sequence match over the whole work.  Absorbs section
                     offsets and differing section labels automatically.
    In `auto` we run both and keep whichever aligns more tokens."""
    gloci = {g[0] for g in gold if g[0]}
    mloci = {m[0] for m in mine}
    shared = gloci & mloci

    def match(mseq, gseq, into):
        a = [norm_form(x[3]) for x in mseq]
        b = [norm_form(x[1]) for x in gseq]
        sm = difflib.SequenceMatcher(a=a, b=b, autojunk=False)
        for op, i1, i2, j1, j2 in sm.get_opcodes():
            if op != "equal":
                continue
            for k in range(i2 - i1):
                into.append((mseq[i1 + k], gseq[j1 + k]))

    def by_locus():
        out = []
        gby, mby = defaultdict(list), defaultdict(list)
        for g in gold:
            gby[g[0]].append(g)
        for m in mine:
            mby[m[0]].append(m)
        for loc in mby:
            if loc in gby:
                match(mby[loc], gby[loc], out)
        return out

    def globally():
        out = []
        match(mine, gold, out)
        return out

    stats = Counter(gold_loci=len(gloci), my_loci=len(mloci), shared_loci=len(shared))
    if mode == "locus":
        pairs, used = by_locus(), "per-citation"
    elif mode == "global":
        pairs, used = globally(), "global sequence"
    else:
        pl, pg = by_locus(), globally()
        if len(pg) >= len(pl):
            pairs, used = pg, "global sequence"
        else:
            pairs, used = pl, "per-citation"
        stats["alt_rate"] = min(len(pg), len(pl))
    stats["mode"] = used
    return pairs, stats


# --- candidate handling -----------------------------------------------------
def grouped_candidates(cache, key):
    """Cache distribution collapsed on the digit-suffix homograph split."""
    d = cache.get(key) or {}
    if not isinstance(d, dict) or not d:
        return []
    g = defaultdict(float)
    for lem, p in d.items():
        try:
            g[re.sub(r"\d+$", "", lem)] += float(p)
        except (TypeError, ValueError):
            continue
    return sorted(g.items(), key=lambda kv: -kv[1])


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--gold", required=True)
    ap.add_argument("--work", required=True, help="slug in data/, e.g. apology")
    ap.add_argument("--data", default="data")
    ap.add_argument("--base", default=None,
                    help="reader base URL; enables the after-re-ranking numbers")
    ap.add_argument("--cache", default=None, help="path to lemma_cache.json")
    ap.add_argument("--mode", choices=["auto","locus","global"], default="auto",
                    help="alignment strategy (auto tries both, keeps the better)")
    ap.add_argument("--errors", type=int, default=15, help="sample errors to print")
    ap.add_argument("--gold-cache", default=None,
                    help="lemma_gold.json overlay from build_gold_cache.py; consulted"
                         " only where the Morpheus cache misses, as the reader does")
    ap.add_argument("--prior", default=None,
                    help="lemma_prior.json from build_lemma_prior.py; adds a prior column")
    ap.add_argument("--prior-gate", type=float, default=10.0,
                    help="promote on frequency ratio >= this (measured sweet spot: 10)")
    args = ap.parse_args()

    gold = parse_gold(args.gold)
    mine, doc = load_work(args.data, args.work)
    print(f"gold   {os.path.basename(args.gold)}: {len(gold):,} word tokens")
    print(f"mine   {args.work}: {len(mine):,} word tokens "
          f"({len(doc['segments'])} segments)")

    pairs, st = align(mine, gold, args.mode)
    rate = len(pairs) / max(1, len(mine))
    print(f"\nALIGNMENT")
    print(f"  citation units: gold {st['gold_loci']}, mine {st['my_loci']}, "
          f"shared {st['shared_loci']}  (strategy chosen: {st['mode']})")
    if st.get('alt_rate'):
        print(f"  the rejected strategy aligned {st['alt_rate']:,} "
              f"({100*st['alt_rate']/max(1,len(mine)):.1f}%) — kept the better one")
    print(f"  aligned {len(pairs):,}/{len(mine):,} of my tokens = {100*rate:.1f}%")
    if rate < 0.90:
        print("  ⚠ below 90% — tokenizations diverge (usually elision/crasis).")
        print("    Accuracy below is computed on the aligned subset only and may be biased.")
    else:
        print("  ✓ alignment is sound; accuracy below is trustworthy")

    cache_path = args.cache or os.path.join(args.data, "lemma_cache.json")
    cache = json.load(open(cache_path, encoding="utf-8"))
    print(f"  lemma cache: {len(cache):,} forms")

    # hand-annotated overlay, merged the way the reader merges it: cache wins,
    # overlay fills only the gaps
    gold_overlay, overlay_meta = {}, {}
    if args.gold_cache:
        try:
            gj = json.load(open(args.gold_cache, encoding="utf-8"))
            overlay_meta = gj.get("meta", {}) or {}
            added = 0
            for form, entry in (gj.get("forms", {}) or {}).items():
                if cache.get(form):
                    continue
                lem = (entry or {}).get("lemmas") or {}
                if lem:
                    cache[form] = dict(lem)
                    gold_overlay[form] = True
                    added += 1
            contrib = {s_["slug"] for s_ in overlay_meta.get("sources", [])}
            print(f"  gold overlay: +{added:,} forms filled "
                  f"(from {', '.join(sorted(contrib)) or '?'})")
            if args.work in contrib:
                print(f"  ⚠ CIRCULAR: '{args.work}' contributed to this overlay — its"
                      f" own gold lemmas were injected into the cache, so accuracy"
                      f" below is meaningless. Rebuild with --exclude {args.work}.")
        except Exception as e:
            print(f"  ! could not load gold overlay: {e!r}")

    # optional: gold-derived frequency prior
    prior, prior_meta = {}, {}
    if args.prior:
        try:
            pj = json.load(open(args.prior, encoding="utf-8"))
            prior = pj.get("freq", {}) or {}
            prior_meta = pj.get("meta", {}) or {}
            contrib = {s_["slug"] for s_ in prior_meta.get("sources", [])}
            print(f"  frequency prior: {len(prior):,} lemmas from "
                  f"{prior_meta.get('total_tokens',0):,} gold tokens "
                  f"(gate {args.prior_gate:g}x)")
            if args.work in contrib:
                print(f"  ⚠ CIRCULAR: '{args.work}' contributed to this prior — the prior")
                print(f"    column below is measured on its own training data and is NOT valid.")
                print(f"    Rebuild with:  build_lemma_prior.py --gold ... --exclude {args.work}")
            else:
                print(f"  ✓ '{args.work}' did not contribute to the prior — held-out, valid")
        except Exception as e:
            print(f"  ! could not load prior: {e!r}")

    # optional: after-re-ranking, from the live endpoint
    disamb = {}
    if args.base:
        try:
            import urllib.request
            u = f"{args.base}/api/disambiguate?id={args.work}"
            d = json.loads(urllib.request.urlopen(u, timeout=300).read())
            if d.get("available"):
                disamb = d.get("by_token") or {}
                print(f"  re-ranking overlay: {len(disamb):,} tokens "
                      f"({d.get('n_flips',0):,} flips)")
            else:
                print("  re-ranking overlay unavailable (no embedding edges for this work)")
        except Exception as e:
            print(f"  ! could not fetch re-ranking overlay: {e!r}")

    # --- score ---
    rows = []
    for (loc, seg_id, ti, surf, key), (gloc, gform, glem, gpos) in pairs:
        gl = norm_lemma(glem)
        cands = grouped_candidates(cache, key)
        known = bool(cands)
        amb = len(cands) > 1
        top_before = norm_lemma(cands[0][0]) if cands else None
        # frequency prior: promote the most frequent candidate, but only when the
        # evidence is lopsided (a small edge is noise; a 10x edge is a real signal).
        top_prior = top_before
        if prior and cands:
            best, bf = None, -1.0
            for c_, _p in cands:
                f = float(prior.get(norm_lemma(c_), 0))
                if f > bf:
                    best, bf = c_, f
            cur = float(prior.get(top_before or "", 0))
            if best is not None and bf >= max(1.0, cur) * args.prior_gate:
                top_prior = norm_lemma(best)

        tok = f"{seg_id}#{ti}"
        e = disamb.get(tok)
        if e and e.get("reranked"):
            # context-specific evidence outranks a global prior where it exists
            top_after = norm_lemma(e["reranked"][0].get("lemma"))
        else:
            top_after = top_before
        top_both = norm_lemma(e["reranked"][0].get("lemma")) if (e and e.get("reranked")) else top_prior
        rows.append({
            "surface": surf, "gold": gl, "known": known, "amb": amb,
            "before": top_before, "after": top_after,
            "prior": top_prior, "both": top_both,
            "ok_before": top_before == gl, "ok_after": top_after == gl,
            "ok_prior": top_prior == gl, "ok_both": top_both == gl,
            "in_cands": any(norm_lemma(c) == gl for c, _ in cands),
            "touched": bool(e), "locus": loc,
        })

    def acc(sel, field):
        s = [r for r in rows if sel(r)]
        if not s:
            return 0.0, 0
        return sum(1 for r in s if r[field]) / len(s), len(s)

    show_prior = bool(prior)
    print(f"\nTOP-1 LEMMA ACCURACY (exact match on the aligned subset)")
    if show_prior:
        print(f"  {'subset':22}{'n':>8}{'baseline':>10}{'+facing':>9}"
              f"{'+prior':>9}{'+both':>9}{'Δ both':>9}")
    else:
        print(f"  {'subset':22}{'n':>8}{'before':>10}{'after':>10}{'Δ':>8}")
    for label, sel in [("all tokens", lambda r: True),
                       ("known (in cache)", lambda r: r["known"]),
                       ("AMBIGUOUS (>1 cand)", lambda r: r["amb"]),
                       ("unambiguous", lambda r: r["known"] and not r["amb"]),
                       ("unknown (not in cache)", lambda r: not r["known"])]:
        b, n = acc(sel, "ok_before")
        a, _ = acc(sel, "ok_after")
        if show_prior:
            pr, _ = acc(sel, "ok_prior")
            bo, _ = acc(sel, "ok_both")
            d = f"{100*(bo-b):+.2f}" if n else "—"
            print(f"  {label:22}{n:8,}{100*b:9.2f}%{100*a:8.2f}%"
                  f"{100*pr:8.2f}%{100*bo:8.2f}%{d:>9}")
        else:
            d = f"{100*(a-b):+.2f}" if n else "—"
            print(f"  {label:22}{n:8,}{100*b:9.2f}%{100*a:9.2f}%{d:>8}")

    ceiling, n_amb = acc(lambda r: r["amb"], "in_cands")
    print(f"\n  ceiling on ambiguous tokens (gold lemma present among candidates): "
          f"{100*ceiling:.2f}%")
    print(f"  → {100*(1-ceiling):.2f}% of ambiguous tokens cannot be got right by any"
          f" re-ranking;\n    the correct lemma is simply absent from the cache.")

    # Separate a RANKING error (the gold lemma is among my candidates but not
    # ranked first — disambiguation can fix this) from an INVENTORY disagreement
    # (the gold lemma is not among my candidates at all — Gorman and Morpheus
    # simply disagree about the headword, e.g. gold lemmatises ὑμῶν under σύ
    # where Morpheus offers ὑμός).  Only the first is a fault of the ranking.
    fair = [r for r in rows if r["amb"] and r["in_cands"]]
    if fair:
        fb = sum(1 for r in fair if r["ok_before"]) / len(fair)
        fa = sum(1 for r in fair if r["ok_after"]) / len(fair)
        print(f"\n  RANKING accuracy on ambiguous tokens whose gold lemma is among my"
              f" candidates\n    (n={len(fair):,}) — the subset disambiguation can actually"
              f" fix: {100*fb:.2f}% → {100*fa:.2f}%")
    unamb_bad = [r for r in rows if r["known"] and not r["amb"] and not r["ok_before"]]
    n_unamb = sum(1 for r in rows if r["known"] and not r["amb"])
    if n_unamb:
        print(f"  INVENTORY disagreement on unambiguous tokens: "
              f"{len(unamb_bad):,}/{n_unamb:,} = {100*len(unamb_bad)/n_unamb:.2f}%")
        print(f"    (single candidate, still != gold — a headword-convention gap,"
              f" not a ranking error)")

    if prior:
        moved = [r for r in rows if r["amb"] and r["prior"] != r["before"]]
        if moved:
            pf = sum(1 for r in moved if r["ok_prior"] and not r["ok_before"])
            pb = sum(1 for r in moved if r["ok_before"] and not r["ok_prior"])
            print(f"\n  the PRIOR moved {len(moved):,} ambiguous tokens: "
                  f"fixed {pf}, broke {pb}, net {pf-pb:+}  "
                  f"(precision {100*pf/max(1,pf+pb):.1f}%)")
        overlap = sum(1 for r in rows if r["amb"] and r["touched"]
                      and r["prior"] != r["before"])
        print(f"  prior and facing-text overlap on {overlap:,} tokens "
              f"— the rest are disjoint, so the two layers mostly add rather than compete")

    touched = [r for r in rows if r["touched"]]
    if touched:
        tb = sum(1 for r in touched if r["ok_before"]) / len(touched)
        ta = sum(1 for r in touched if r["ok_after"]) / len(touched)
        fixed = sum(1 for r in touched if r["ok_after"] and not r["ok_before"])
        broke = sum(1 for r in touched if r["ok_before"] and not r["ok_after"])
        print(f"\n  on the {len(touched):,} tokens the re-ranking actually TOUCHED:")
        print(f"    accuracy {100*tb:.2f}% → {100*ta:.2f}%   "
              f"fixed {fixed}, broke {broke}, net {fixed-broke:+}")

    # report against the BEST available ranking, not the facing-only one, or the
    # sample misrepresents what is actually still broken.
    best_field = "ok_both" if prior else "ok_after"
    best_val = "both" if prior else "after"
    bad = [r for r in rows if r["amb"] and not r[best_field]]
    if bad and args.errors:
        print(f"\n  sample of remaining errors on ambiguous tokens "
              f"({len(bad):,} still wrong after {'prior + facing' if prior else 'facing'}):")
        step = max(1, len(bad) // args.errors)
        for r in bad[::step][:args.errors]:
            mark = "in-cands" if r["in_cands"] else "NOT-IN-CACHE"
            print(f"    {r['locus']:>7}  {r['surface']:14} gold={r['gold']:14} "
                  f"got={r[best_val] or '—':14} [{mark}]")


if __name__ == "__main__":
    main()