| """Haas-ja-v1.0 end to end on Apple silicon: NeMo-exact log-mel -> CoreML int8 encoder (ANE, fixed 15 s window, |
| overlapping chunks for longer audio) -> TDT greedy decoding on the CPU, with either the CoreML decoder/joint (default; what the iOS app runs) or the |
| fp32 ONNX decoder_joint (--onnx-decoder). |
| |
| Usage: python haas_coreml.py [--onnx-decoder] <wav> [<wav> ...] |
| """ |
| import os, sys, time |
| HERE = os.path.dirname(os.path.abspath(__file__)) |
| ROOT = os.path.dirname(HERE) |
| |
| ONNX_DIR = os.environ.get("HAAS_ONNX_DIR", os.path.join(ROOT, "onnx")) |
| sys.path.insert(0, HERE) |
| import numpy as np, torch, coremltools as ct |
| from haas_ort import DURATIONS, HOP, MAX_SYMBOLS, N_FFT, WIN, SR, BLANK, Tokenizer, greedy_tdt, load_wav, sessions |
|
|
| WINDOW = 1504 |
| CONTEXT = 208 |
| SUB = 8 |
|
|
|
|
| class Frontend: |
| """NeMo AudioToMelSpectrogramPreprocessor as in the checkpoint: bf16 window/filterbank, seq_len = len // hop.""" |
|
|
| def __init__(self, path=os.path.join(ROOT, "swift/Resources/frontend.bin")): |
| b = torch.from_numpy(np.fromfile(path, dtype=np.float32)) |
| self.window, self.fb = b[:400], b[400:].reshape(128, 257) |
|
|
| def __call__(self, audio): |
| x = torch.tensor(audio, dtype=torch.float32)[None] |
| seq_len = x.shape[1] // HOP |
| x = torch.cat([x[:, :1], x[:, 1:] - 0.97 * x[:, :-1]], dim=1) |
| spec = torch.stft(x, N_FFT, HOP, WIN, window=self.window, center=True, pad_mode="constant", return_complex=True) |
| spec = torch.view_as_real(spec).pow(2).sum(-1) |
| mel = torch.log(torch.matmul(self.fb, spec) + 2 ** -24) |
| m = mel[:, :, :seq_len] |
| mean = m.mean(-1, keepdim=True) |
| std = torch.sqrt(((m - mean) ** 2).sum(-1, keepdim=True) / (seq_len - 1)) + 1e-5 |
| return ((m - mean) / std).numpy(), seq_len |
|
|
|
|
| def chunks(L, window=WINDOW, ctx=CONTEXT): |
| """Split L mel frames into (start, keep_from, keep_to) windows; kept ranges tile [0, L) and are multiples of SUB.""" |
| if L <= window: |
| return [(0, 0, L)] |
| out, keep = [], 0 |
| while keep < L: |
| last = -(-(L - window) // SUB) * SUB |
| start = max(0, min(keep - ctx, last)) |
| end_keep = L if start >= last else start + window - ctx |
| out.append((start, keep, end_keep)) |
| keep = end_keep |
| return out |
|
|
|
|
| class Encoder: |
| def __init__(self, path=os.path.join(ROOT, "encoder_T1504_w8.mlpackage"), compute_units="CPU_AND_NE"): |
| self.m = ct.models.MLModel(path, compute_units=getattr(ct.ComputeUnit, compute_units)) |
|
|
| def run(self, mel, L): |
| """mel [1,128,>=L] normalized -> encoder output [1,1024,ceil(L/8)].""" |
| 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] |
| y = self.m.predict({"audio_signal": x, "length": np.array([n], np.int32)})["encoded"] |
| outs.append(y[:, :, (k0 - start) // SUB:(k1 - start + SUB - 1) // SUB]) |
| return np.concatenate(outs, axis=2) |
|
|
|
|
| class CoreMLDecoder: |
| """TDT greedy decoding with decoder / joint_enc / joint .mlpackages (convert_decoder.py), same rule as greedy_tdt.""" |
|
|
| def __init__(self, compute_units="CPU_ONLY"): |
| cu = getattr(ct.ComputeUnit, compute_units) |
| self.dec = ct.models.MLModel(os.path.join(ROOT, "decoder.mlpackage"), compute_units=cu) |
| self.jenc = ct.models.MLModel(os.path.join(ROOT, "joint_enc.mlpackage"), compute_units=cu) |
| self.joint = ct.models.MLModel(os.path.join(ROOT, "joint.mlpackage"), compute_units=cu) |
| self.width = self.dec.get_spec().description.input[1].type.multiArrayType.shape[-1] |
|
|
| def project(self, enc): |
| """encoder output [1,1024,T] -> f [T,640], in 188-frame windows as joint_enc is fixed-shape.""" |
| T, n = enc.shape[2], 188 |
| fs = [] |
| for s in range(0, T, n): |
| x = np.zeros((1, 1024, n), np.float32); x[:, :, :min(n, T - s)] = enc[:, :, s:s + n] |
| fs.append(self.jenc.predict({"encoded": x})["f"][0, :min(n, T - s)]) |
| return np.concatenate(fs) |
|
|
| def __call__(self, enc, T): |
| f = self.project(enc) |
| rec = np.zeros((2, 1, self.width), np.float32); conv = np.zeros_like(rec) |
|
|
| def pred(tok, rec, conv): |
| o = self.dec.predict({"token": np.array([tok], np.int32), "rec": rec, "conv": conv}) |
| return o["g"], o["rec_out"], o["conv_out"] |
|
|
| g, rec_n, conv_n = pred(BLANK, rec, conv) |
| tokens, t, last_ts, count = [], 0, -1, 0 |
| while t < T: |
| while True: |
| o = self.joint.predict({"f": f[t:t + 1], "g": g}) |
| k, d = int(o["token_id"][0]), DURATIONS[int(o["duration_id"][0])] |
| if k == BLANK and d == 0: |
| d = 1 |
| t_label, t = t, t + d |
| if k != BLANK or t >= T: |
| break |
| if k != BLANK: |
| tokens.append(k) |
| count = count + 1 if t_label == last_ts else 1 |
| last_ts = t_label |
| g, rec_n, conv_n = pred(k, rec_n, conv_n) |
| if count >= MAX_SYMBOLS and last_ts == t: |
| t += 1 |
| return tokens |
|
|
|
|
| class Haas: |
| def __init__(self, compute_units="CPU_AND_NE", onnx_decoder=False): |
| self.front = Frontend() |
| self.enc = Encoder(compute_units=compute_units) |
| self._onnx_dec = None |
| self.cml_dec = None if onnx_decoder else CoreMLDecoder() |
| self.tok = Tokenizer(os.path.join(ROOT, "tokenizer.model")) |
|
|
| @property |
| def dec(self): |
| if self._onnx_dec is None: |
| _, self._onnx_dec = sessions(os.path.join(ONNX_DIR, "encoder.onnx"), os.path.join(ONNX_DIR, "decoder_joint.onnx"), 4) |
| return self._onnx_dec |
|
|
| def decode(self, enc): |
| if self.cml_dec is not None: |
| return self.cml_dec(enc, enc.shape[2]) |
| return greedy_tdt(self.dec, enc, enc.shape[2]) |
|
|
| def transcribe(self, audio): |
| mel, L = self.front(audio) |
| enc = self.enc.run(mel, L) |
| ids = self.decode(enc) |
| return [int(i) for i in ids], self.tok.decode(ids) |
|
|
|
|
| if __name__ == "__main__": |
| args = sys.argv[1:] |
| h = Haas(onnx_decoder="--onnx-decoder" in args) |
| for p in [a for a in args if not a.startswith("--")]: |
| a = load_wav(p) |
| t0 = time.time(); _, text = h.transcribe(a); dt = time.time() - t0 |
| print(f"{os.path.basename(p)}\t{len(a) / SR:.1f}s\tRTF {dt / (len(a) / SR):.4f}\t{text}") |
|
|
