meter/predict.py

"""Inference + benchmark glue for a trained MeterModel.

  # Norma benchmark, both tasks, scored in-process (--norma-source hf|git, default hf):
  python -m meter.predict --model $STOICHEIA_DATA/runs/meter_joint/best.pt --norma \
      [--pred-out data/norma_preds.jsonl] [--norma-source hf]
  # work-split scanner dev/test:
  python -m meter.predict --model ... --scan-split
  # production:
  python -m meter.predict --model ... --macronize in.txt out.txt
  python -m meter.predict --model ... --scan in.txt out.txt

--norma scores both tasks directly via meter.norma_score (acc, balanced acc, per-class
F1 for macronize; acc, balanced acc, boundary-F1, weight acc for syllabify), split into
dev/test, and additionally dumps macron predictions/probabilities in the old
predictions-jsonl format for anyone still using the legacy ensemble/scorer scripts.
"""
from __future__ import annotations

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

import numpy as np
import torch

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from meter.backbone import load_backbone
from meter.dataset import batch_rows, encode_plain, load_records, pack_records
from meter.marks import (ambiguous_mask, bracketize, enforce_circumflex_heavy,
                         insert_marks, merge_vowelless_syllables,
                         parse_macron_line, parse_scan_line)
from meter.model import MeterConfig, MeterModel
from meter.norma_data import add_norma_source_arg
from meter.norma_score import scan_metrics
from meter.train import SCAN_DEV_WORKS, SCAN_TEST_WORKS


def load_model(path, device, attn="sdpa"):
    sd = torch.load(path, map_location="cpu")
    encoder, _ = load_backbone(os.path.expandvars(sd["cfg"]["ckpt"]), device, attn)
    mcfg = MeterConfig(**sd["mcfg"])
    model = MeterModel(encoder, mcfg).to(device)
    model.load_state_dict(sd["model"])
    model.eval()
    return model, sd


@torch.no_grad()
def predict_records(model, records, T, micro, device):
    """-> per record (mac_argmax, scan_argmax, mac_P(long), scan_probs) with None
    entries preserved. mac_P(long) is P(class 0) per letter; scan_probs is (n,4).
    scan_argmax has two deterministic corrections applied, in order: (1)
    merge_vowelless_syllables -- a predicted syllable span with no vowel gets
    folded into the preceding one, keeping its own (usually already-correct)
    weight; (2) enforce_circumflex_heavy -- a circumflexed syllable is always
    heavy. Every consumer (scoring, --scan production output, Viterbi) gets both
    fixed rules for free."""
    live = [i for i, r in enumerate(records) if r is not None]
    rows, skipped = pack_records(records, T, live)
    if skipped:
        print(f"  WARNING: {skipped} records longer than T={T} skipped", file=sys.stderr)
    out = [None] * len(records)
    for i in range(0, len(rows), micro):
        chunk = rows[i:i + micro]
        batch = batch_rows(chunk, records, T, device=device, with_slots=True)
        with torch.autocast("cuda", dtype=torch.bfloat16, enabled=device.type == "cuda"):
            o = model({k: v for k, v in batch.items() if k != "slots"})
        mac_p = torch.softmax(o["mac"].float(), -1).cpu().numpy()
        scan_p = torch.softmax(o["scan"].float(), -1).cpu().numpy()
        pm = mac_p.argmax(-1)
        ps = scan_p.argmax(-1)
        for b, slots in enumerate(batch["slots"]):
            for ri, c in slots:
                n = len(records[ri])
                scan_argmax = merge_vowelless_syllables(records[ri].chars, ps[b, c:c + n])
                scan_argmax = enforce_circumflex_heavy(records[ri].dia, scan_argmax)
                out[ri] = (pm[b, c:c + n], scan_argmax,
                           mac_p[b, c:c + n, 0], scan_p[b, c:c + n])
    return out


def run_norma(model, sd, device, micro, pred_out, norma_source="hf"):
    from meter.norma_data import load_norma
    from meter.norma_score import MAC_LONG, mac_metrics
    T = sd["T"]
    norma = load_norma(norma_source)

    # -------- macronize: scored directly (in-process; no external scorer needed)
    probs_rows = []
    with open(pred_out, "w", encoding="utf-8") as f:
        for split_name, rows in [("dev", norma["dev"]), ("test", norma["test"])]:
            mac_rows = [d for d in rows if d["task"] == "macronize"]
            items = [(d["source"], *parse_macron_line(d["text"])) for d in mac_rows]
            recs = [encode_plain(p) for _, p, _ in items]
            preds = predict_records(model, recs, T, micro, device)
            pairs, by_source = [], defaultdict(list)
            for (srcname, plain, gold), pr in zip(items, preds):
                g = sorted(gold.items())
                pred_labels = [(k, int(pr[0][k]) if pr is not None else MAC_LONG) for k, _ in g]
                f.write(json.dumps(dict(
                    split=split_name, source=srcname,
                    gold=[[k, "", v] for k, v in g],
                    pred=[[k, "", pv] for k, pv in pred_labels]), ensure_ascii=False) + "\n")
                if pr is not None:
                    probs_rows.append(dict(split=split_name, source=srcname, seq=[
                        [v, float(pr[2][k])] for k, v in g]))
                for (k, gv), (_, pv) in zip(g, pred_labels):
                    pairs.append((gv, pv))
                    by_source[srcname].append((gv, pv))
            overall = mac_metrics(pairs)
            per_source = {s: mac_metrics(ps) for s, ps in by_source.items()}
            bals = [m["bal_acc"] for m in per_source.values() if m and m["bal_acc"] is not None]
            macro = round(sum(bals) / len(bals), 4) if bals else None
            print(f"macron {split_name}: {json.dumps(dict(**overall, macro_bal_acc=macro))}"
                  if overall else f"macron {split_name}: (empty)")
    probs_out = str(pred_out).replace(".jsonl", "") + "_probs.json"
    json.dump(probs_rows, open(probs_out, "w"))
    print(f"macron P(long) dump -> {probs_out}")
    print(f"macron preds -> {pred_out}")

    # -------- syllabify: test only (Norma has no syllabify dev rows, on hf or git)
    syl_items, syl_recs = [], []
    for d in norma["test"]:
        if d["task"] != "syllabify":
            continue
        parsed = parse_scan_line(d["text"])
        if parsed is None:
            continue
        plain, gold = parsed
        rec = encode_plain(plain)
        if rec is None:
            continue
        syl_items.append((d["source"], gold))
        syl_recs.append(rec)
    preds = predict_records(model, syl_recs, T, micro, device)
    print("\n=== norma syllabify (test) ===")
    pairs, by_src = [], defaultdict(list)
    for (srcname, gold), pr in zip(syl_items, preds):
        if pr is None:
            continue
        golds = np.zeros(len(pr[1]), dtype=np.int64)
        for k, g in gold.items():
            golds[k] = g
        rec_pairs = list(zip(golds.tolist(), pr[1].tolist()))
        pairs += rec_pairs
        by_src[srcname] += rec_pairs
    m = scan_metrics(pairs)
    bals = [b["bal_acc"] for s in sorted(by_src) if (b := scan_metrics(by_src[s]))]
    macro = round(sum(bals) / len(bals), 4) if bals else None
    print(f"  {json.dumps(dict(**m, macro_bal_acc=macro))}")


def run_scan_split(model, sd, device, micro):
    T = sd["T"]
    enc_dir = Path(os.path.expandvars(sd["cfg"]["encoded"]))
    recs, works = load_records(enc_dir / "scan_corpus.npz")
    for split, wanted in (("dev", SCAN_DEV_WORKS), ("test", SCAN_TEST_WORKS)):
        sel = [(r, w) for r, w in zip(recs, works) if w in wanted]
        preds = predict_records(model, [r for r, _ in sel], T, micro, device)
        pairs, by_work = [], defaultdict(list)
        for (r, w), pr in zip(sel, preds):
            if pr is None:
                continue
            p = list(zip(r.y_scan.tolist(), pr[1].tolist()))
            p = [(g, q) for g, q in p if g != -100]
            pairs += p
            by_work[w] += p
        print(f"\n=== scan {split} (whole verses, by work) ===")
        print(f"  all: {json.dumps(scan_metrics(pairs))}")
        for w in sorted(by_work):
            print(f"  {w}: {json.dumps(scan_metrics(by_work[w]))}")


def run_viterbi(model, sd, src, device, micro, theta, norma_source="hf"):
    """Meter-constrained decoding: exact-line + char accuracy, raw argmax vs
    gated Viterbi (meter forced from the corpus label when a grammar exists,
    else auto-detected; Norma syllabify is always auto)."""
    import numpy as np
    from meter.viterbi import METER_MAP, gated_auto, gated_decode
    T = sd["T"]

    def decode_set(name, items):
        """items: (group, meter_name|None, gold {ord: lab}, record)"""
        recs = [r for _, _, _, r in items]
        preds = predict_records(model, recs, T, micro, device)
        agg = defaultdict(lambda: np.zeros(6, np.int64))
        # [lines, exact_raw, exact_vit, char_ok_raw, char_ok_vit, chars] per group
        applied = 0
        for (group, mname, gold, rec), pr in zip(items, preds):
            if pr is None:
                continue
            n = len(pr[3])
            golds = np.zeros(n, np.int64)
            for k, g in gold.items():
                golds[k] = g
            logp = np.log(np.clip(pr[3], 1e-9, 1.0))
            raw = pr[1]   # pre-computed argmax, with enforce_circumflex_heavy already applied
            grammar = METER_MAP.get(mname or "")
            if grammar:
                vit, ok = gated_decode(logp.tolist(), grammar, theta)
            else:
                _, vit, ok = gated_auto(logp.tolist(), theta)
            applied += ok
            vit = np.asarray(vit)
            for g in (group, "ALL"):
                a = agg[g]
                a[0] += 1
                a[1] += int((raw == golds).all())
                a[2] += int((vit == golds).all())
                a[3] += int((raw == golds).sum())
                a[4] += int((vit == golds).sum())
                a[5] += n
        print(f"\n=== viterbi {name} (theta={theta}, applied {applied}/"
              f"{agg['ALL'][0]}) ===")
        for g in sorted(agg, key=lambda x: (x != "ALL", x)):
            a = agg[g]
            print(f"  {g}: lines={a[0]} exact raw={a[1]/a[0]:.3f} "
                  f"vit={a[2]/a[0]:.3f} | char raw={a[3]/a[5]:.4f} vit={a[4]/a[5]:.4f}")

    # ---- work-split dev/test from the scanner corpus (grouped by grammar)
    from meter.marks import parse_scan_line
    for split, works in (("scan-dev", SCAN_DEV_WORKS), ("scan-test", SCAN_TEST_WORKS)):
        items = []
        for line in open(src / "data/scanner/corpus_v3.tsv", encoding="utf-8"):
            parts = line.rstrip("\n").split("\t")
            if len(parts) != 3 or parts[0] not in works:
                continue
            parsed = parse_scan_line(parts[2])
            if parsed is None:
                continue
            rec = encode_plain(parsed[0])
            if rec is None:
                continue
            grammar = METER_MAP.get(parts[1], f"auto({parts[1] or 'lyric'})")
            items.append((grammar, parts[1], parsed[1], rec))
        decode_set(split, items)

    # ---- Norma syllabify (auto meter; test only -- Norma has no syllabify dev rows)
    from meter.norma_data import load_norma
    items = []
    for d in load_norma(norma_source)["test"]:
        if d["task"] != "syllabify":
            continue
        parsed = parse_scan_line(d["text"])
        if parsed is None:
            continue
        rec = encode_plain(parsed[0])
        if rec is None:
            continue
        items.append((d["source"], None, parsed[1], rec))
    decode_set("norma-syllabify", items)


def run_file(model, sd, device, micro, mode, infile, outfile):
    T = sd["T"]
    raw = [l.rstrip("\n") for l in open(infile, encoding="utf-8")]
    plains = [parse_macron_line(l)[0] for l in raw]   # strips any existing marks
    recs = [encode_plain(p) if p.strip() else None for p in plains]
    preds = predict_records(model, recs, T, micro, device)
    with open(outfile, "w", encoding="utf-8") as f:
        for plain, rec, pr in zip(plains, recs, preds):
            if rec is None or pr is None:
                f.write(plain + "\n")
                continue
            if mode == "macronize":
                amb = ambiguous_mask(rec.chars, rec.boundary, rec.dia)
                labels = {int(i): int(pr[0][i]) for i in np.flatnonzero(amb)}
                f.write(insert_marks(plain, labels) + "\n")
            else:
                labels = {i: int(c) for i, c in enumerate(pr[1]) if c > 0}
                f.write(bracketize(plain, labels) + "\n")
    print(f"{mode}: {len(raw)} lines -> {outfile}")


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--model", required=True)
    ap.add_argument("--attn", default="sdpa")
    ap.add_argument("--micro", type=int, default=16)
    ap.add_argument("--norma", action="store_true")
    ap.add_argument("--pred-out", default=None)
    ap.add_argument("--scan-split", action="store_true")
    ap.add_argument("--viterbi", action="store_true")
    ap.add_argument("--theta", type=float, default=0.1)
    ap.add_argument("--macronize", nargs=2, metavar=("IN", "OUT"))
    ap.add_argument("--scan", nargs=2, metavar=("IN", "OUT"))
    add_norma_source_arg(ap)
    a = ap.parse_args()
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model, sd = load_model(a.model, device, a.attn)
    src = Path(os.path.expandvars(os.environ.get("MACRONIZER_SRC",
                              "$MACRONIZER_SRC")))
    if a.norma:
        pred_out = a.pred_out or (Path(sd["cfg"]["out_dir"]) / "norma_pred.jsonl")
        run_norma(model, sd, device, a.micro, pred_out, a.norma_source)
    if a.scan_split:
        run_scan_split(model, sd, device, a.micro)
    if a.viterbi:
        run_viterbi(model, sd, src, device, a.micro, a.theta, a.norma_source)
    if a.macronize:
        run_file(model, sd, device, a.micro, "macronize", *a.macronize)
    if a.scan:
        run_file(model, sd, device, a.micro, "scan", *a.scan)


if __name__ == "__main__":
    main()