Blame for code/analyze.py

"""Collect metrics for E1-E4 from prediction files into results/*.json and LaTeX table rows.

Run locally after syncing /home/jovyan/humonen/runic_access/preds -> ~/runic_access/preds.
Missing prediction files are skipped, so the script can be re-run as jobs finish.
"""
import json
import os
import sys

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import metrics  # noqa: E402
from loopcut import cut_loops  # noqa: E402

ROOT = os.path.expanduser("~/runic_access")
P = f"{ROOT}/preds"
OUT = f"{ROOT}/results"
os.makedirs(OUT, exist_ok=True)


LEAKFREE = json.load(open(f"{ROOT}/data/leakfree.json"))
KEEP = set(LEAKFREE["keep_test"]) | set(LEAKFREE["keep_val"])


def keep_row(r):
    """Test/val rows whose inscription is not linked to any training pair (relink.py)."""
    if r["id"].startswith(("syn", "dif")):
        return True
    rid = r.get("base_id") or r.get("row_id") or r["id"]
    return rid in KEEP


PHOTO_TEST = {json.loads(l)["image"] for l in open(f"{ROOT}/splits_server/ocr_realf_test.jsonl")} | \
    {json.loads(l)["image"] for l in open(f"{ROOT}/splits_server/ocr_realf_val.jsonl")}
ONLY_PHOTOS = os.environ.get("ONLY_PHOTOS") == "1"  # default: all images; domains are reported separately
CREAM = json.load(open(f"{ROOT}/data/cream_share.json"))
RED = json.load(open(f"{ROOT}/data/red_share.json"))
# test images of the card domain checked by hand: photographs of objects and manuscript pages show the find itself
CARD_LAB = json.load(open(f"{ROOT}/data/card_labels.json"))["labels"]


def domain(r):
    """Image domain of a real test image: Gamla runor field photograph, RuneS photograph, RuneS archive card."""
    if r["source"] == "ch":
        return "gr"
    img = r.get("image_full", r["image"])
    if img in CARD_LAB:
        return "rdb_photo" if CARD_LAB[img] in ("photo", "manuscript") else "card"
    return "card" if CREAM.get(img, 0) > 0.2 else "rdb_photo"


def load(name):
    path = f"{P}/{name}.jsonl"
    if not os.path.exists(path):
        return None
    rows = [r for r in metrics.read(path) if keep_row(r)]
    if ONLY_PHOTOS and rows and "image" in rows[0] and rows[0]["image"].startswith(("images/", "crops/")):
        rows = [r for r in rows if r.get("image_full", r["image"]) in PHOTO_TEST]
    return rows


def fmt_ci(v, ci):
    return f"{v:.1f} [{ci[0]:.1f}, {ci[1]:.1f}]"


res = {}

# ---------------- E1 recognition ----------------
e1 = {}
OCR_TAGS = ["ocr_test_zs08", "ocr_test_zs4b", "ocr_test_ocr4b_syn", "ocr_test_ocr4b_diff", "ocr_test_ocr4b_real",
            "ocr_test_ocr4b_mix", "ocr_test_ocr4b_diffmix", "ocr_test_ocr4b_mix_hr", "ocr_test_ocr08_mix",
            "ocr_test_ocr4b_crop", "ocr_test_ocr4b_mix_oncrop"]
for tag in OCR_TAGS:
    rows = load(tag)
    if not rows:
        continue
    raw = metrics.ocr_scores(rows)
    rows = [dict(r, pred=cut_loops(r.get("pred_raw", r["pred"]))) for r in rows]  # repetition truncation
    e1[tag] = {"all": metrics.ocr_scores(rows), "raw": raw}
    for dom in ("gr", "rdb_photo", "card"):
        sub = [r for r in rows if domain(r) == dom]
        if sub:
            e1[tag][dom] = metrics.ocr_scores(sub)
    # output length diagnostics: runaway enumerations that loopcut does not catch inflate corpus-level CER
    per_image = sorted(100 * metrics.cer_parts(r["target"], r["pred"])[0] / metrics.cer_parts(r["target"], r["pred"])[1]
                       for r in rows)
    e1[tag]["median_cer"] = per_image[len(per_image) // 2] if len(per_image) % 2 else \
        (per_image[len(per_image) // 2 - 1] + per_image[len(per_image) // 2]) / 2
    e1[tag]["len_ratio"] = sum(len(r["pred"]) for r in rows) / sum(len(r["target"]) for r in rows)
    e1[tag]["over2x"] = 100 * sum(len(r["pred"]) > 2 * len(r["target"]) for r in rows) / len(rows)
    painted = [r for r in rows if RED.get(r.get("image_full", r["image"]), 0) > 0.01]
    if painted:
        e1[tag]["painted"] = metrics.ocr_scores(painted)
    base = tag.replace("ocr_test_", "")
    for vtag, vname in ((f"ocr_synval_{base}", "proc_syn"), (f"ocr_difval_{base}", "diff_syn")):
        vr = load(vtag)
        if vr:
            e1[tag][vname] = metrics.ocr_scores([dict(x, pred=cut_loops(x["pred"])) for x in vr])
    lens = {"short": [r for r in rows if len(r["target"]) <= 15], "medium": [r for r in rows if 15 < len(r["target"]) <= 40],
            "long": [r for r in rows if len(r["target"]) > 40]}
    for k, v in lens.items():
        e1[tag][k] = metrics.ocr_scores(v)
    for src in ("ch", "rdb"):
        sub = [r for r in rows if r["source"] == src]
        if sub:
            e1[tag][src] = metrics.ocr_scores(sub)
import collections
import random as _random

_train_mt = metrics.read(f"{ROOT}/splits_local/mt_train.jsonl")
COMMON = {w for w, _ in collections.Counter(w for r in _train_mt for w in r["src"].split()).most_common(150)}


def rare_word_recall(rows, shuffle=False):
    preds = [r["pred"] for r in rows]
    if shuffle:
        _random.Random(0).shuffle(preds)
    hit = tot = 0
    for r, p in zip(rows, preds):
        ref = {w for w in r["target"].split() if w not in COMMON and len(w) >= 4}
        tot += len(ref)
        hit += len(ref & set(p.split()))
    return 100 * hit / max(tot, 1)


for tag in list(e1):
    rows = [dict(r, pred=cut_loops(r.get("pred_raw", r["pred"]))) for r in load(tag)]
    e1[tag]["rare_recall"] = rare_word_recall(rows)
    e1[tag]["rare_recall_shuffled"] = rare_word_recall(rows, shuffle=True)

PRIOR = json.load(open(f"{ROOT}/results/prior_baselines.json"))
ocr_test_ref = [r for r in metrics.read(f"{ROOT}/splits_server/ocr_real_test.jsonl") if keep_row(r)
                and (not ONLY_PHOTOS or r["image"] in PHOTO_TEST)]
const_rows = [dict(r, pred=PRIOR["ocr_constant"]) for r in ocr_test_ref]
e1["ocr_test_constant"] = {"all": metrics.ocr_scores(const_rows)}
for dom in ("gr", "rdb_photo", "card"):
    if any(domain(r) == dom for r in const_rows):  # ONLY_PHOTOS leaves the card domain empty
        e1["ocr_test_constant"][dom] = metrics.ocr_scores([r for r in const_rows if domain(r) == dom])
for tag in ["ocr_synval_ocr4b_syn", "ocr_synval_ocr4b_mix", "ocr_difval_ocr4b_diff"]:
    rows = load(tag)
    if rows:
        e1[tag] = {"all": metrics.ocr_scores([dict(r, pred=cut_loops(r["pred"])) for r in rows])}
res["E1"] = e1

# ---------------- E2 translation ----------------
e2 = {}
for tag in ["mt_test_zs08", "mt_test_zs4b", "mt_test_retrieval", "mt_test_mt08", "mt_test_mt4b", "mt_test_mt4b_noise"]:
    rows = load(tag)
    if not rows:
        continue
    e2[tag] = {}
    for lang in ("en", "sv"):
        e2[tag][lang] = metrics.mt_scores([r for r in rows if r["lang"] == lang])
    e2[tag]["strict"] = metrics.mt_scores([r for r in rows if (r.get("max_ratio_to_train") or 0) < 80])
    e2[tag]["all"] = metrics.mt_scores(rows)
import re as _re

_NAME_STOP = {"May", "God", "Christ", "The", "This", "These", "He", "She", "They", "His", "Her", "And", "In", "I", "A",
              "Holy", "Mary", "Lord", "Jesus", "Amen", "Saint", "St", "Gud", "Han", "Hon", "De", "Och", "Den", "Denna",
              "Detta", "Dessa", "Kristus"}


def _names(t):
    return {w for w in _re.findall(r"\b[A-ZÁÉÍÓÚÝÞÆØÖÄÅ][\w'-]+", t or "") if w not in _NAME_STOP}


def name_recall(rows):
    hit = tot = 0
    for r in rows:
        ref = _names(r["tgt"])
        tot += len(ref)
        hit += len(ref & _names(r.get("pred")))
    return 100 * hit / max(tot, 1)


for tag in list(e2):
    e2[tag]["name_recall"] = name_recall(load(tag))
pairs = [("mt_test_mt4b", "mt_test_retrieval"), ("mt_test_mt4b", "mt_test_mt08"), ("mt_test_mt4b", "mt_test_mt4b_noise"),
         ("mt_test_mt4b_noise", "mt_test_mt4b")]
e2["paired"] = {}
for a, b in pairs:
    ra, rb = load(a), load(b)
    if ra and rb:
        e2["paired"][f"{a}>{b}"] = metrics.paired_bootstrap_chrf(ra, rb)
res["E2"] = e2

# ---------------- E3 end-to-end ----------------
TRAIN_VOCAB = {w for r in _train_mt for w in r["src"].split()}


def in_vocab(texts):
    """Share of words (tokens) that occur in the training transliterations."""
    ws = [w for t in texts for w in (t or "").split()]
    return 100 * sum(w in TRAIN_VOCAB for w in ws) / max(len(ws), 1)


def ref_words_kept(pairs, select=lambda w: True):
    """Share of reference words (optionally a subset) that occur verbatim in the corrupted or recognized transliteration."""
    hit = tot = 0
    for ref, hyp in pairs:
        hw = set((hyp or "").split())
        ws = [w for w in ref.split() if select(w)]
        tot += len(ws)
        hit += sum(w in hw for w in ws)
    return 100 * hit / max(tot, 1)


e3 = {}
ocr_tag = os.environ.get("OCR_TAG", "ocr4b_mixcut")
GOLD_TAG = "ocr4b_mixcut"  # translations of reference transliterations do not depend on the recognizer


def load_e2e(tr, src):
    rows = load(f"e2e_{ocr_tag}_{tr}_{src}")
    if rows is None and src == "fromgold":
        rows = load(f"e2e_{GOLD_TAG}_{tr}_{src}")
    return rows


for tr in ("mt4b", "mt4b_noise"):
    for src in ("fromgold", "fromocr"):
        rows = load_e2e(tr, src)
        if not rows:
            continue
        tag = f"e2e_{ocr_tag}_{tr}_{src}"
        e3[tag] = {"all": metrics.mt_scores(rows)}
        for lang in ("en", "sv"):
            e3[tag][lang] = metrics.mt_scores([r for r in rows if r["lang"] == lang])
ocr_rows = load(f"ocr_test_{ocr_tag}")
if ocr_rows:
    # retrieval from the recognizer output and from the reference transliteration
    import subprocess
    for field, name in (("pred", "fromocr"), ("target", "fromgold")):
        out = f"{P}/e2e_{ocr_tag}_retrieval_{name}.jsonl"
        tmp = f"{P}/_e2e_input.jsonl"
        with open(tmp, "w") as f:
            for r in ocr_rows:
                f.write(json.dumps(dict(r, src=r[field] or ""), ensure_ascii=False) + "\n")
        subprocess.run([sys.executable, f"{ROOT}/code/retrieval_baseline.py", f"{ROOT}/splits_local/mt_train.jsonl", tmp, out],
                       check=True)
        rows = [r for r in metrics.read(out) if keep_row(r)]
        e3[f"e2e_{ocr_tag}_retrieval_{name}"] = {"all": metrics.mt_scores(rows)}
    const_mt = [dict(r, pred=PRIOR[f"mt_constant_{r['lang']}"]) for r in ocr_rows]
    e3[f"e2e_{ocr_tag}_constant"] = {"all": metrics.mt_scores(const_mt)}
    # CER of the recogniser on the same photographs, per language
    e3["ocr_cer"] = {lang: metrics.ocr_scores([r for r in ocr_rows if r["lang"] == lang]) for lang in ("en", "sv")}
    e3["ocr_cer"]["all"] = metrics.ocr_scores(ocr_rows)
    a, b = load(f"e2e_{ocr_tag}_mt4b_noise_fromocr"), load(f"e2e_{ocr_tag}_mt4b_fromocr")
    if a and b:
        e3["paired_noise_vs_clean_fromocr"] = metrics.paired_bootstrap_chrf(a, b)
    # subsets: recognition CER of each photograph (analysis only, needs the reference) and image domain
    img_cer = {r["id"]: 100 * metrics.cer_parts(r["target"], r["pred"] or "")[0] / metrics.cer_parts(r["target"], r["pred"] or "")[1]
               for r in ocr_rows}
    by_id = {r["id"]: r for r in ocr_rows}

    def cer_bin(c):
        return "cer<25" if c < 25 else "cer25-50" if c < 50 else "cer50-75" if c < 75 else "cer>=75"

    subsets = {}
    for i, c in img_cer.items():
        subsets.setdefault(cer_bin(c), set()).add(i)
        subsets.setdefault(domain(by_id[i]), set()).add(i)
    systems = {"mt4b_fromgold": load_e2e("mt4b", "fromgold"), "mt4b_noise_fromgold": load_e2e("mt4b_noise", "fromgold"),
               "mt4b_fromocr": load_e2e("mt4b", "fromocr"), "mt4b_noise_fromocr": load_e2e("mt4b_noise", "fromocr"),
               "retrieval_fromocr": [r for r in metrics.read(f"{P}/e2e_{ocr_tag}_retrieval_fromocr.jsonl") if keep_row(r)],
               "constant": const_mt}
    # paired comparisons of translation from the photograph against the image-blind and lookup baselines
    for sysname in ("mt4b_fromocr", "mt4b_noise_fromocr"):
        for base in ("constant", "retrieval_fromocr"):
            if systems[sysname] and systems[base]:
                e3.setdefault("paired", {})[f"{sysname}>{base}"] = metrics.paired_bootstrap_chrf(systems[sysname], systems[base])
    e3["subsets"] = {}
    for name in ("cer<25", "cer25-50", "cer50-75", "cer>=75", "gr", "rdb_photo", "card"):
        ids = subsets.get(name, set())
        if not ids:
            continue
        v = {"n": len(ids), "ocr": metrics.ocr_scores([by_id[i] for i in sorted(ids)]),
             "in_vocab": in_vocab([by_id[i]["pred"] for i in ids]),
             "ref_words_kept": ref_words_kept([(by_id[i]["target"], by_id[i]["pred"]) for i in ids])}
        for sysname, rows in systems.items():
            if rows:
                v[sysname] = metrics.mt_scores([r for r in rows if r["id"] in ids], boot=False)
        e3["subsets"][name] = v
    # archive cards by what the crop shows (manual inspection, data/card_labels.json): whole image vs this pipeline
    whole_ocr = {r["id"]: r for r in load(f"ocr_test_{GOLD_TAG}") or []}
    whole_mt = {r["id"]: r for r in load(f"e2e_{GOLD_TAG}_mt4b_fromocr") or []}
    mt_by = {k: {r["id"]: r for r in systems[k] or []} for k in ("mt4b_fromocr", "mt4b_fromgold")}
    crop_shows = {"card_photo": "photo", "card_drawing": "drawing", "drawing": "drawing", "card_whole": "whole"}
    groups = {}
    for i in by_id:
        lab = CARD_LAB.get(by_id[i].get("image_full", by_id[i]["image"]))
        if lab in crop_shows and domain(by_id[i]) == "card":
            groups.setdefault(crop_shows[lab], []).append(i)
            groups.setdefault("any", []).append(i)
    e3["cards"] = {}
    for lab, ids in groups.items():
        ids = sorted(ids)
        e3["cards"][lab] = {"n": len(ids),
                            "cer_whole": metrics.ocr_scores([whole_ocr[i] for i in ids])["cer"] if whole_ocr else None,
                            "cer": metrics.ocr_scores([by_id[i] for i in ids])["cer"],
                            "chrfpp_whole": metrics.mt_scores([whole_mt[i] for i in ids], boot=False)["chrfpp"] if whole_mt else None,
                            "chrfpp": metrics.mt_scores([mt_by["mt4b_fromocr"][i] for i in ids], boot=False)["chrfpp"],
                            "chrfpp_gold": metrics.mt_scores([mt_by["mt4b_fromgold"][i] for i in ids], boot=False)["chrfpp"]}
    # reference-free selection: rank photographs by a confidence signal and score the most confident fraction
    from rapidfuzz.distance import Levenshtein as _Lev
    from scipy.stats import spearmanr

    def _nd(x, y):
        return _Lev.distance(x or "", y or "") / max(len(x or ""), len(y or ""), 1)

    signals = {}
    score_path = f"{P}/score_test_{ocr_tag}.jsonl"
    if os.path.exists(score_path):
        sc = {r["id"]: r["nll"] for r in metrics.read(score_path)}
        # decodes stopped by the generation limit (320 tokens, count_raw_tokens.py) never ended; rank them last
        ntok_path = f"{P}/ntok_ocr_test_{ocr_tag}.json"
        ntok = json.load(open(ntok_path)) if os.path.exists(ntok_path) else {}
        if all(sc.get(i) is not None for i in by_id if (by_id[i]["pred"] or "").strip()):
            signals["nll"] = {i: (sc.get(i) if sc.get(i) is not None else float("inf")) + (1e3 if ntok.get(i, 0) >= 320 else 0)
                              for i in by_id}
    others = {}
    for t in ("ocr4b_mix_hr", "ocr4b_real", "ocr4b_diffmix", "ocr08_mix"):
        rows_t = load(f"ocr_test_{t}")
        if rows_t:
            others[t] = {r["id"]: cut_loops(r.get("pred_raw", r["pred"]) or "") for r in rows_t}
    if "ocr4b_mix_hr" in others:
        signals["agree_hr"] = {i: _nd(by_id[i]["pred"], others["ocr4b_mix_hr"].get(i)) for i in by_id}
    if len(others) == 4:
        signals["agree4"] = {i: sum(_nd(by_id[i]["pred"], o.get(i)) for o in others.values()) / 4 for i in by_id}
    signals["oracle_cer"] = img_cer
    mt_ocr = {r["id"]: r for r in systems["mt4b_fromocr"] or []}
    mt_gold = {r["id"]: r for r in systems["mt4b_fromgold"] or []}
    e3["selection"] = {}
    for name, sig in signals.items():
        ids_sorted = sorted(by_id, key=lambda i: sig[i])
        v = {"spearman_cer": spearmanr([sig[i] for i in ids_sorted], [img_cer[i] for i in ids_sorted]).correlation}
        for frac in (0.1, 0.2, 0.3, 0.5):
            sel = ids_sorted[:round(frac * len(ids_sorted))]
            v[str(frac)] = {"n": len(sel), "cer": metrics.ocr_scores([by_id[i] for i in sel])["cer"],
                            "chrfpp_ocr": metrics.mt_scores([mt_ocr[i] for i in sel if i in mt_ocr], boot=False)["chrfpp"],
                            "chrfpp_gold": metrics.mt_scores([mt_gold[i] for i in sel if i in mt_gold], boot=False)["chrfpp"]}
        e3["selection"][name] = v
res["E3"] = e3

# ---------------- E4 noise curve ----------------
e4 = {}
act = f"{P}/../splits_server/mt_test_noise.jsonl.actual_cer.json"
if os.path.exists(act):
    e4["actual_cer"] = json.load(open(act))
for tr in ("mt4b", "mt4b_noise"):
    rows = load(f"noise_{tr}")
    if not rows:
        continue
    by = {}
    for r in rows:
        by.setdefault(r["level"], []).append(r)
    e4[tr] = {str(lvl): metrics.mt_scores(v, boot=False) for lvl, v in sorted(by.items())}
    if tr == "mt4b":
        e4["input"] = {str(lvl): {"in_vocab": in_vocab([r["src"] for r in v]),
                                  "ref_words_kept": ref_words_kept([(r["src_clean"], r["src"]) for r in v]),
                                  "common_kept": ref_words_kept([(r["src_clean"], r["src"]) for r in v], lambda w: w in COMMON),
                                  "other_kept": ref_words_kept([(r["src_clean"], r["src"]) for r in v], lambda w: w not in COMMON)}
                       for lvl, v in sorted(by.items())}
if ocr_rows:
    e4["recognizer_input"] = {"in_vocab": in_vocab([r["pred"] for r in ocr_rows]),
                              "common_kept": ref_words_kept([(r["target"], r["pred"]) for r in ocr_rows], lambda w: w in COMMON),
                              "other_kept": ref_words_kept([(r["target"], r["pred"]) for r in ocr_rows], lambda w: w not in COMMON),
                              "reference_in_vocab": in_vocab([r["target"] for r in ocr_rows])}
res["E4"] = e4

suffix = ("_photos" if ONLY_PHOTOS else "") + ("" if ocr_tag == "ocr4b_mixcut" else f"_{ocr_tag}")
json.dump(res, open(f"{OUT}/results{suffix}.json", "w"), indent=1, default=float)

# ---------------- console summary ----------------
for tag, v in e1.items():
    a = v["all"]
    extra = " ".join(f"{k}:{v[k]['cer']:.1f}" for k in ("proc_syn", "diff_syn", "gr", "rdb_photo", "card", "painted") if k in v)
    raw = f" (raw CER {v['raw']['cer']:.1f})" if "raw" in v else ""
    lens = " ".join(f"{k}:{v[k]['cer']:.0f}" for k in ("short", "medium", "long") if k in v)
    if "rare_recall" in v:
        lens += f" rare {v['rare_recall']:.1f}/{v['rare_recall_shuffled']:.1f}"
    diag = f" median {v['median_cer']:.1f} len {v['len_ratio']:.2f} >2x {v['over2x']:.1f}%" if "median_cer" in v else ""
    print(f"E1 {tag:28s} n={a['n']} CER {fmt_ci(a['cer'], a['cer_ci'])}{raw} WER {a['wer']:.1f} EM {a['exact']:.1f} {extra} {lens}{diag}")
for tag, v in e2.items():
    if tag == "paired":
        continue
    print(f"E2 {tag:22s} " + " | ".join(
        f"{k}: chrF++ {v[k]['chrfpp']:.1f} BLEU {v[k]['bleu']:.1f} (n={v[k]['n']}, empty={v[k]['empty']})" for k in ("en", "sv", "strict")))
for k, v in e2.get("paired", {}).items():
    print(f"E2 paired {k}: delta {v['delta']:.2f} p={v['p']:.3f}")
for tag, v in e3.items():
    if "chrfpp" in v.get("all", {}):
        print(f"E3 {tag:40s} chrF++ {v['all']['chrfpp']:.1f} BLEU {v['all']['bleu']:.1f} n={v['all']['n']}")
if "ocr_cer" in e3:
    print("E3 OCR CER on e2e photos:", {k: round(v["cer"], 1) for k, v in e3["ocr_cer"].items()})
for name, v in e3.get("subsets", {}).items():
    print(f"E3 subset {name:9s} n={v['n']:3d} CER {v['ocr']['cer']:.1f} invocab {v['in_vocab']:.1f} kept {v['ref_words_kept']:.1f} " +
          " ".join(f"{k}:{v[k]['chrfpp']:.1f}" for k in v if isinstance(v[k], dict) and "chrfpp" in v[k]))
for name, v in e3.get("selection", {}).items():
    print(f"E3 select {name:10s} rho {v['spearman_cer']:.3f} " + " ".join(
        f"{f}:n{v[f]['n']} CER {v[f]['cer']:.1f} chrF {v[f]['chrfpp_ocr']:.1f}/{v[f]['chrfpp_gold']:.1f}" for f in ("0.1", "0.2", "0.3", "0.5")))
for k, v in e3.get("paired", {}).items():
    print(f"E3 paired {k}: delta {v['delta']:.2f} p={v['p']:.3f}")
print("E4 input stats:", {k: {a: round(b, 1) for a, b in v.items()} for k, v in e4.get("input", {}).items()},
      e4.get("recognizer_input"))
for tr in ("mt4b", "mt4b_noise"):
    if tr in e4:
        print(f"E4 {tr}: " + ", ".join(f"{lvl}->{v['chrfpp']:.1f}" for lvl, v in e4[tr].items()))
if "actual_cer" in e4:
    print("E4 actual CER per level:", e4["actual_cer"])