File size: 3,841 Bytes
9456b7b | 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 | """Evaluate 1-800-BAD-CODE punctuation models (punctuators package, ONNX; input = lowercased text without punctuation;
they also true-case and split sentences) on the same chunks and metric as punct/eval_punct.py (lowercased input only).
Their output is re-normalised with punct/common.units and aligned to the input words (case-insensitive).
python punct/eval_pcs.py OUT.json EVAL_DIR MODEL [MODEL ...] (run in a private venv with punctuators installed)
"""
import difflib, json, os, sys, time
import numpy as np, pandas as pd
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from common import LABELS, NOSPACE, units, lower_keep_len
from punctuators.models import PunctCapSegModelONNX
out_f, ev = sys.argv[1:3]; models = sys.argv[3:]
KREDOR = {"eng_Latn", "deu_Latn", "fra_Latn", "spa_Latn", "bul_Cyrl", "ita_Latn", "pol_Latn", "nld_Latn", "ces_Latn", "por_Latn",
"slk_Latn", "slv_Latn"}
items = []
for p in sorted(os.listdir(ev)):
if not p.endswith(".parquet"):
continue
name, df = p[:-8], pd.read_parquet(f"{ev}/{p}")
for text, lang in zip(df.text, df.lang):
nos = lang.startswith(("zho", "jpn", "yue")) or lang in NOSPACE
ul = "cmn_Hani" if lang.startswith(("zho", "yue")) else lang
c, e, l = units(text, ul)
cap = 250 if nos else 150
for a in range(0, len(e), cap):
off = e[a - 1] + (0 if nos else 1) if a else 0
ee = e[a:a + cap]
items.append((name, lang, ul, c[off:ee[-1]], [x - off for x in ee], l[a:a + cap], nos))
print("chunks", len(items), flush=True)
res = json.load(open(out_f)) if os.path.exists(out_f) else {}
for mname in models:
m = PunctCapSegModelONNX.from_pretrained(mname); t0 = time.time()
texts = [lower_keep_len(it[3]) for it in items]; outs = []
for s in range(0, len(texts), 64):
outs += m.infer(texts[s:s + 64]); print(s, f"{time.time() - t0:.0f}s", flush=True) if s % 6400 == 0 else None
cnt, miss = {}, 0
for it, o in zip(items, outs):
name, lang, ul, clean, ends, gold, nos = it
oc, oe, ol = units(("" if nos else " ").join(o), ul)
gw = [clean[(ends[i - 1] if i else 0):x].strip().lower() for i, x in enumerate(ends)]
pw = [oc[(oe[i - 1] if i else 0):x].strip().lower() for i, x in enumerate(oe)]
pred = ["0"] * len(gw)
for tag, a1, a2, b1, b2 in difflib.SequenceMatcher(None, gw, pw, autojunk=False).get_opcodes():
if tag == "equal":
pred[a1:a2] = ol[b1:b2]
else:
miss += a2 - a1
c = cnt.setdefault((name, lang), np.zeros((6, 6), np.int64))
for g, q in zip(gold, pred):
c[LABELS.index(g), LABELS.index(q)] += 1
r = {}
for (sname, lang), c in cnt.items():
tp = sum(c[i, i] for i in range(1, 6)); fp = c[:, 1:].sum() - tp; fn = c[1:, :].sum() - tp
r.setdefault(sname, {})[lang] = {"lower": dict(f1=2 * tp / max(1, 2 * tp + fp + fn), conf=c.tolist(),
**{LABELS[i]: 2 * c[i, i] / max(1, c[i, :].sum() + c[:, i].sum()) for i in range(1, 6)})}
summ = {"_unaligned_words": miss}
for sname, langs in r.items():
for sub, ls in [("all", list(langs)), ("kredor12", [l for l in langs if l in KREDOR])]:
if ls:
summ[f"{sname}/lower/{sub}"] = float(np.mean([langs[l]["lower"]["f1"] for l in ls]))
pooled = sum(np.array(langs[l]["lower"]["conf"]) for l in langs)
for i in (1, 2, 3):
summ[f"{sname}/lower/pooled_{LABELS[i]}"] = float(2 * pooled[i, i] / max(1, pooled[i, :].sum() + pooled[:, i].sum()))
r["_summary"] = summ; res[mname] = r
print(mname, f"{time.time() - t0:.0f}s", json.dumps({k: round(v, 4) for k, v in summ.items()}), flush=True)
json.dump(res, open(out_f, "w"))
|