File size: 7,061 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 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 | """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) # 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}")
|