File size: 5,297 Bytes
b11a6e7 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 | """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)
|