| """Validate haas_coreml.py end to end. |
| 1) wav33 (all <= 15 s, one window): token match with results/wav33/pynemo_fp32.json, RTF. |
| 2) long audio (wav33 concatenated with 0.4 s silence, ~45-55 s): CER against the joined per-clip NeMo texts for |
| CoreML chunked, fp32 ONNX with the same chunking, and fp32 ONNX on the whole utterance (no chunking). |
| Writes results/coreml_e2e.json.""" |
| import json, os, sys, time, unicodedata |
| sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) |
| import numpy as np |
| from rapidfuzz.distance import Levenshtein |
| from haas_coreml import ONNX_DIR, SR, SUB, WINDOW, Haas, chunks, load_wav, sessions |
|
|
| |
| h = Haas() |
| enc_onnx, _ = sessions(os.path.join(ONNX_DIR, "encoder.onnx"), os.path.join(ONNX_DIR, "decoder_joint.onnx"), 4) |
|
|
|
|
| def onnx_encode(mel, L, chunked): |
| def run(x, n): |
| y, _ = enc_onnx.run(None, {"audio_signal": x, "length": np.array([n], np.int64)}) |
| return y |
| if not chunked: |
| return run(np.ascontiguousarray(mel[:, :, :L]), L) |
| outs = [] |
| for start, k0, k1 in chunks(L): |
| n = min(WINDOW, L - start) |
| x = np.zeros((1, 128, WINDOW), np.float32); x[:, :, :n] = mel[:, :, start:start + n] |
| outs.append(run(x, n)[:, :, (k0 - start) // SUB:(k1 - start + SUB - 1) // SUB]) |
| return np.concatenate(outs, axis=2) |
|
|
|
|
| def norm(t): |
| t = unicodedata.normalize("NFKC", t).lower() |
| return "".join(c for c in t if unicodedata.category(c)[0] not in "PZS") |
|
|
|
|
| def decode(enc): |
| return [int(i) for i in __import__("haas_coreml").greedy_tdt(h.dec, enc, enc.shape[2])] |
|
|
|
|
| def main(): |
| ref = json.load(open("results/wav33/pynemo_fp32.json")) |
| paths = open("wav33.txt").read().split() |
| items = [(os.path.splitext(os.path.basename(p))[0], load_wav(p)) for p in paths] |
| h.transcribe(items[0][1]) |
| out = {"wav33": {}, "long": {}} |
| match, t_all, dur = 0, 0.0, 0.0 |
| for k, a in items: |
| t0 = time.time(); ids, text = h.transcribe(a); t_all += time.time() - t0; dur += len(a) / SR |
| match += ids == ref[k]["ids"] |
| out["wav33"][k] = {"ids": ids, "text": text, "match": ids == ref[k]["ids"]} |
| if ids != ref[k]["ids"]: |
| print("DIFF", k, "\n ref:", ref[k]["text"], "\n hyp:", text) |
| print(json.dumps({"wav33_match": f"{match}/{len(items)}", "e2e_rtf": round(t_all / dur, 4), "audio_s": round(dur, 1)})) |
|
|
| |
| sil = np.zeros(int(0.4 * SR), np.float32) |
| groups, cur = [], [] |
| for k, a in items: |
| cur.append((k, a)) |
| if sum(len(x) for _, x in cur) / SR > 45: |
| groups.append(cur); cur = [] |
| if cur: |
| groups.append(cur) |
| tot = {k: [0, 0] for k in ["coreml_chunked", "onnx_chunked", "onnx_full"]} |
| for gi, g in enumerate(groups): |
| a = np.concatenate([np.concatenate([x, sil]) for _, x in g]) |
| ref_text = "".join(ref[k]["text"] for k, _ in g) |
| mel, L = h.front(a) |
| t0 = time.time(); _, c_text = h.transcribe(a); dt = time.time() - t0 |
| texts = {"coreml_chunked": c_text, "onnx_chunked": h.tok.decode(decode(onnx_encode(mel, L, True))), |
| "onnx_full": h.tok.decode(decode(onnx_encode(mel, L, False)))} |
| row = {"keys": [k for k, _ in g], "audio_s": round(len(a) / SR, 1), "n_chunks": len(chunks(L)), |
| "rtf": round(dt / (len(a) / SR), 4), "ref_text": ref_text, "texts": texts} |
| for name, t in texts.items(): |
| e = Levenshtein.distance(norm(t), norm(ref_text)); row[f"{name}_cer"] = round(100 * e / len(norm(ref_text)), 2) |
| tot[name][0] += e; tot[name][1] += len(norm(ref_text)) |
| out["long"][f"group{gi}"] = row |
| print(json.dumps({k: v for k, v in row.items() if "text" not in k})) |
| summ = {k: f"{100 * v[0] / v[1]:.2f}% ({v[0]}/{v[1]})" for k, v in tot.items()} |
| print(json.dumps({"long_cer_vs_joined_clip_refs": summ})) |
| out["summary"] = summ |
| json.dump(out, open("results/coreml_e2e.json", "w"), ensure_ascii=False, indent=1) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
