File size: 4,148 Bytes
0d06934 | 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 | """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
# results/wav33/pynemo_fp32.json and wav33.txt (the eval audio list) are not distributed; run from a dir that has them.
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]) # warmup / ANE load
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)}))
# long audio: groups of consecutive clips
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()
|