"""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] [ ...] """ import os, sys, time HERE = os.path.dirname(os.path.abspath(__file__)) ROOT = os.path.dirname(HERE) # repo root: the .mlpackages and tokenizer.model # fp32 ONNX encoder/decoder_joint (from Sanami/haas-ja-onnx), only needed for --onnx-decoder and validate.py 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 # noqa: F401 WINDOW = 1504 # mel frames per encoder call (15.04 s) -> 188 encoder frames CONTEXT = 208 # mel frames of context on each side of a chunk that are encoded but dropped (2.08 s) SUB = 8 # encoder subsampling 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)) # window [400] then filterbank [128, 257] 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 # normalized over the whole utterance, as NeMo 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 # first aligned start whose window reaches L 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): # ONNX session, used by validate.py for the fp32 reference paths 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}")