| """Evaluate an omniASR CTC model (+ optional fine-tuned checkpoint) with the duplex decoding (greedy CTC + Latin-script mask). |
| ~/venv_omni/bin/python eval_asr.py --tag NAME [--card omniASR_CTC_1B_v2] [--ckpt /scratch/asr/runs/A/best.pt] [--sets duplex,realcs,synthtest,ace] |
| --hyp_only MANIFEST : just transcribe a manifest (used for the synthetic-data CER filter) -> MANIFEST.hyp.jsonl |
| Writes /scratch/asr/eval/<tag>.json (scores) and <tag>_hyps.jsonl.""" |
| import argparse, json, os, re, unicodedata, time, numpy as np, torch, torch.nn.functional as F |
| import soundfile as sf |
| from scipy.signal import resample_poly |
| from fairseq2.models.hub import load_model |
| from fairseq2.data.tokenizers.hub import load_tokenizer |
| from fairseq2.nn.batch_layout import BatchLayout |
| import jiwer |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--tag", default="x"); ap.add_argument("--card", default="omniASR_CTC_1B_v2"); ap.add_argument("--ckpt", default="") |
| ap.add_argument("--sets", default="duplex,realcs,synthtest,ace"); ap.add_argument("--hyp_only", default=""); ap.add_argument("--latin", type=int, default=1) |
| a = ap.parse_args() |
| dev = torch.device("cuda") |
| model = load_model(a.card, device=dev, dtype=torch.bfloat16).eval() |
| if a.ckpt: |
| ck = torch.load(a.ckpt, map_location="cuda", weights_only=False); print("load", a.ckpt, model.load_state_dict(ck["model"]), ck.get("step"), ck.get("dev")) |
| tok = load_tokenizer(a.card); dec = tok.create_decoder(skip_special_tokens=True) |
| V = model.final_proj.weight.shape[0] |
| def _ok(sx): |
| for c in sx: |
| cat = unicodedata.category(c) |
| if cat[0] == "M": |
| if not ("\u0300" <= c <= "\u036f"): return False |
| elif c.isalpha() and "LATIN" not in unicodedata.name(c, ""): return False |
| return True |
| BAD = torch.tensor([i for i in range(V) if not _ok(str(dec(torch.tensor([i]))))], dtype=torch.long, device=dev) if a.latin else None |
| print("latin mask blocks", 0 if BAD is None else len(BAD), "of", V) |
|
|
| def load(r): |
| if r.get("wav"): |
| y, sr = sf.read(r["wav"], dtype="float32") |
| if y.ndim > 1: y = y.mean(1) |
| return y if sr == 16000 else resample_poly(y, 16000, sr).astype(np.float32) |
| m = np.memmap(r["pcm"], dtype=np.int16, mode="r"); y = m[r["off24"]: r["off24"] + r["n24"]].astype(np.float32) / 32768 |
| return resample_poly(y, 2, 3).astype(np.float32) |
|
|
| @torch.no_grad() |
| def transcribe(items, bs=8): |
| out, lat = [], [] |
| order = sorted(range(len(items)), key=lambda i: items[i].get("dur", 0)); res = [None] * len(items) |
| for s in range(0, len(order), bs): |
| idx = order[s: s + bs]; ys = [torch.from_numpy(load(items[i])) for i in idx] |
| ys = [F.layer_norm(y, y.shape) for y in ys]; L = [len(y) for y in ys] |
| X = torch.zeros(len(ys), max(L)); [X[k, : len(y)].copy_(y) for k, y in enumerate(ys)] |
| X = X.to(dev, torch.bfloat16); bl = BatchLayout(X.shape, seq_lens=L, device=dev) |
| lg, blo = model(X, bl) |
| if BAD is not None: lg[..., BAD] = -float("inf") |
| pr = lg.argmax(-1) |
| for k, i in enumerate(idx): |
| sq = pr[k, : blo.seq_lens[k]]; m = torch.ones_like(sq, dtype=torch.bool); m[1:] = sq[1:] != sq[:-1] |
| res[i] = str(dec(sq[m])) |
| return res |
|
|
| def n(s): |
| s = unicodedata.normalize("NFC", s.lower()); s = re.sub(r"[^\w\s']", " ", s); return re.sub(r"\s+", " ", s).strip() |
| def fold(s): return "".join(c for c in unicodedata.normalize("NFD", s) if unicodedata.category(c) != "Mn") |
| def score(refs, hyps): |
| R = [n(r) or "x" for r in refs]; H = [n(h) or "x" for h in hyps] |
| return dict(n=len(R), wer=round(jiwer.wer(R, H), 4), cer=round(jiwer.cer(R, H), 4), |
| cer_nodiac=round(jiwer.cer([fold(r) for r in R], [fold(h) for h in H]), 4)) |
|
|
| if a.hyp_only: |
| items = [json.loads(l) for l in open(a.hyp_only)]; hy = transcribe(items, 16) |
| with open(a.hyp_only + ".hyp.jsonl", "w") as f: |
| for r, h in zip(items, hy): |
| f.write(json.dumps(dict(id=r["id"], hyp=h, cer=round(jiwer.cer(n(r["text"]) or "x", n(h) or "x"), 4)), ensure_ascii=False) + "\n") |
| print("HYP_DONE", len(items)); raise SystemExit |
|
|
| os.makedirs("/scratch/asr/eval", exist_ok=True); res = {}; allh = [] |
| sets = a.sets.split(",") |
| if "duplex" in sets: |
| R = [json.loads(l) for l in open("/scratch/duplex/asr_eval/list.jsonl")]; H = transcribe(R) |
| for lg in ["eng", "pcm", "ibo", "yor", "hau", "cs"]: |
| ix = [i for i, r in enumerate(R) if r["lang"] == lg]; res[f"duplex_{lg}"] = score([R[i]["text"] for i in ix], [H[i] for i in ix]) |
| allh += [dict(set="duplex", id=r["id"], ref=r["text"], hyp=h) for r, h in zip(R, H)] |
| if "realcs" in sets: |
| R = [json.loads(l) for l in open("/scratch/asr/data/realcs_test.jsonl")]; H = transcribe(R) |
| res["realcs_all"] = score([r["text"] for r in R], H) |
| for k in ["natural_conversational", "insertional"]: |
| ix = [i for i, r in enumerate(R) if r["cs_kind"] == k]; res[f"realcs_{k}"] = score([R[i]["text"] for i in ix], [H[i] for i in ix]) |
| hit = tot = 0 |
| for r, h in zip(R, H): |
| hw = n(h).split() |
| for w in r["cs_eng"]: |
| tot += 1; hit += w in hw |
| res["realcs_eng_word_recall"] = round(hit / max(1, tot), 4) |
| allh += [dict(set="realcs", id=r["id"], ref=r["text"], hyp=h) for r, h in zip(R, H)] |
| if "synthtest" in sets and os.path.exists("/scratch/asr/synth/test.jsonl"): |
| R = [json.loads(l) for l in open("/scratch/asr/synth/test.jsonl")]; H = transcribe(R); res["synthtest"] = score([r["text"] for r in R], H) |
| allh += [dict(set="synthtest", id=r["id"], ref=r["text"], hyp=h) for r, h in zip(R, H)] |
| if "ace" in sets and os.path.exists("/scratch/asr/ace/list.jsonl"): |
| R = [json.loads(l) for l in open("/scratch/asr/ace/list.jsonl")]; H = transcribe(R, 1); res["ace"] = score([r["text"] for r in R], H) |
| allh += [dict(set="ace", id=r["id"], ref=r["text"], hyp=h) for r, h in zip(R, H)] |
| json.dump(dict(tag=a.tag, card=a.card, ckpt=a.ckpt, scores=res), open(f"/scratch/asr/eval/{a.tag}.json", "w"), indent=1) |
| with open(f"/scratch/asr/eval/{a.tag}_hyps.jsonl", "w") as f: |
| for h in allh: f.write(json.dumps(h, ensure_ascii=False) + "\n") |
| print(json.dumps(res, indent=1)); print("EVAL_DONE", a.tag) |
|
|
