| """Comma -> capital boundary errors of a joint truecase model, and a boundary-consistent decoding. |
| The mark of word i and the case of word i+1 are separate labels, so the model can write "hello, How are you". Where the argmax gives |
| (',' after word i, Cap on word i+1), compare three options with the marginal probabilities: keep (',' , Cap) * beta, ('.'/'?' , Cap), |
| (',' , lower). beta = 1 is the plain argmax. beta is chosen on held-out VALIDATION windows (val.parquet rows after the 6,000 used for |
| checkpoint selection), then the test sets are scored with every beta for the record. |
| python truecase/boundary_check.py OUT.json DATA_DIR EVAL_DIR MODEL [MODEL ...]""" |
| import json, os, sys |
| import numpy as np, pandas as pd, torch |
| sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))); sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "punct")) |
| from common import LABELS as PL, NOSPACE, lower_keep_len |
| from common_tc import unit_cases |
| from tc_metrics import chunks, score |
| from transformers import AutoModelForTokenClassification, AutoTokenizer |
|
|
| out_f, data, ev = sys.argv[1:4]; models = sys.argv[4:] |
| BETAS = [1.0, 0.8, 0.6, 0.5, 0.4, 0.3, 0.2, 0.1, 0.0] |
| COMMA, PER, Q = PL.index(","), PL.index("."), PL.index("?") |
|
|
| va = pd.read_parquet(f"{data}/val.parquet").sample(frac=1.0, random_state=0).iloc[6000:] |
| val_items = [dict(set="val", lang=l, text=t, ends=e.tolist(), gp=[PL[x] for x in lb], gc=unit_cases(t, e.tolist(), l in NOSPACE), nos=l in NOSPACE) |
| for t, e, lb, l in zip(va.text, va.ends, va.labs, va.lang)] |
| test_items = [it for it in chunks(ev) if it["set"] in ("europarl", "ted", "flores")] |
| print("val", len(val_items), "test", len(test_items), flush=True) |
|
|
|
|
| def probs(tok, model, items): |
| texts = [lower_keep_len(it["text"]) for it in items]; order = np.argsort([len(t) for t in texts]); res = [None] * len(items) |
| for s in range(0, len(order), 64): |
| idx = order[s:s + 64] |
| enc = tok([texts[i] for i in idx], truncation=True, max_length=512, padding=True, return_offsets_mapping=True, return_tensors="pt") |
| offs = enc.pop("offset_mapping").tolist() |
| with torch.no_grad(): |
| p = torch.softmax(model(**{k: v.cuda() for k, v in enc.items()}).logits.float(), -1).cpu().numpy() |
| for j, i in enumerate(idx): |
| at = {e: k for k, (b, e) in enumerate(offs[j]) if e > b} |
| P = np.zeros((len(items[i]["ends"]), 18), np.float32); P[:, 0] = 1.0 |
| for u, e in enumerate(items[i]["ends"]): |
| if e in at: |
| P[u] = p[j, at[e]] |
| res[i] = P |
| return res |
|
|
|
|
| def decode(P, beta): |
| j = P.argmax(1); pm, pc = (j // 3).tolist(), (j % 3).tolist() |
| if beta < 1: |
| M, C = P.reshape(-1, 6, 3).sum(2), P.reshape(-1, 6, 3).sum(1) |
| for i in range(len(pm) - 1): |
| if pm[i] == COMMA and pc[i + 1] == 1: |
| b = PER if M[i, PER] >= M[i, Q] else Q |
| z, x, y = M[i, COMMA] * C[i + 1, 1] * beta, M[i, b] * C[i + 1, 1], M[i, COMMA] * C[i + 1, 0] |
| if x >= z and x >= y: |
| pm[i] = b |
| elif y > z: |
| pc[i + 1] = 0 |
| return [PL[m] for m in pm], pc |
|
|
|
|
| def events(items, PP): |
| """raw argmax (',' , Cap) boundaries by gold: sentence boundary, comma + capital (name/I), comma/none + lower (wrong capital).""" |
| n = dict(total=0, gold_sentence_break=0, gold_comma_cap=0, gold_lower=0, other=0, words=0) |
| for it, P in zip(items, PP): |
| pm, pc = decode(P, 1.0); n["words"] += len(pm) |
| for i in range(len(pm) - 1): |
| if pm[i] == "," and pc[i + 1] == 1: |
| n["total"] += 1; g, gc = it["gp"][i], it["gc"][i + 1] |
| k = "gold_sentence_break" if g in (".", "?") and gc == 1 else "gold_comma_cap" if gc == 1 else "gold_lower" if gc == 0 else "other" |
| n[k] += 1 |
| return n |
|
|
|
|
| res = json.load(open(out_f)) if os.path.exists(out_f) else {} |
| for m in models: |
| tok = AutoTokenizer.from_pretrained(m); model = AutoModelForTokenClassification.from_pretrained(m, dtype=torch.bfloat16).cuda().eval() |
| PV, PT = probs(tok, model, val_items), probs(tok, model, test_items) |
| r = dict(events_val=events(val_items, PV), events_test={s: events([it for it in test_items if it["set"] == s], [P for it, P in zip(test_items, PT) if it["set"] == s]) |
| for s in ("europarl", "ted", "flores")}, val={}, test={}) |
| for b in BETAS: |
| dv = [decode(P, b) for P in PV]; dt = [decode(P, b) for P in PT] |
| r["val"][b] = score(val_items, [d[0] for d in dv], [d[1] for d in dv])["summary"] |
| r["test"][b] = score(test_items, [d[0] for d in dt], [d[1] for d in dt])["summary"] |
| v = r["val"][b]["val"]; print(m[-50:], "beta", b, "val", round(v["punct_f1"], 4), round(v["case_f1"], 4), |
| {s: (round(x["punct_f1"], 4), round(x["case_f1"], 4)) for s, x in r["test"][b].items()}, flush=True) |
| best = max(BETAS, key=lambda b: (r["val"][b]["val"]["punct_f1"] + r["val"][b]["val"]["case_f1"]) / 2); r["best_beta_val"] = best |
| print(m[-50:], "events", r["events_val"], r["events_test"], "best beta (val)", best, flush=True) |
| res[m] = r; json.dump(res, open(out_f, "w"), indent=1) |
|
|
