| """Haas-ja-v1.0 ONNX inference without NeMo: NeMo-compatible log-mel + TDT greedy decoding.""" |
| import librosa |
| import numpy as np |
| import onnxruntime as ort |
| import sentencepiece as spm |
| import soundfile as sf |
|
|
| SR, N_FFT, WIN, HOP, N_MELS, PAD_TO = 16000, 512, 400, 160, 128, 16 |
| DURATIONS = [0, 1, 2, 3, 4] |
| MAX_SYMBOLS = 10 |
| BLANK = 3072 |
|
|
|
|
| def _bf16(x): |
| """Round float32 values to bfloat16 precision (round to nearest even), kept as float32.""" |
| b = np.asarray(x, np.float32).view(np.uint32) |
| b = (b + 0x7FFF + ((b >> 16) & 1)) & 0xFFFF0000 |
| return b.astype(np.uint32).view(np.float32) |
|
|
|
|
| |
| |
| _FB = _bf16(librosa.filters.mel(sr=SR, n_fft=N_FFT, n_mels=N_MELS, fmin=0, fmax=SR / 2, norm="slaney")) |
| _WINDOW = _bf16(np.hanning(WIN).astype(np.float32)) |
| _WINDOW_PAD = np.pad(_WINDOW, ((N_FFT - WIN) // 2, (N_FFT - WIN + 1) // 2)) |
|
|
|
|
| def logmel(audio): |
| """NeMo AudioToMelSpectrogramPreprocessor (eval: no dither), batch 1. Returns ([1,128,T_pad], valid_len).""" |
| x = np.asarray(audio, np.float32) |
| seq_len = len(x) // HOP |
| x = np.concatenate([x[:1], x[1:] - np.float32(0.97) * x[:-1]]) |
| x = np.pad(x, N_FFT // 2) |
| n = 1 + (len(x) - N_FFT) // HOP |
| frames = np.lib.stride_tricks.as_strided(x, (n, N_FFT), (x.strides[0] * HOP, x.strides[0])) |
| spec = np.fft.rfft(frames * _WINDOW_PAD, axis=1) |
| power = (spec.real ** 2 + spec.imag ** 2).astype(np.float32).T |
| mel = np.log(_FB @ power + np.float32(2 ** -24)) |
| m = mel[:, :seq_len] |
| mean = m.mean(-1, keepdims=True) |
| std = np.sqrt(((m - mean) ** 2).sum(-1, keepdims=True) / (seq_len - 1)) + np.float32(1e-5) |
| mel = (mel - mean) / std |
| mel[:, seq_len:] = 0 |
| T = mel.shape[-1] |
| if T % PAD_TO: |
| mel = np.pad(mel, ((0, 0), (0, PAD_TO - T % PAD_TO))) |
| return mel[None].astype(np.float32), seq_len |
|
|
|
|
| def load_wav(path): |
| audio, sr = sf.read(path, dtype="float32") |
| if audio.ndim > 1: |
| audio = audio.mean(1) |
| if sr != SR: |
| audio = librosa.resample(audio, orig_sr=sr, target_sr=SR) |
| return audio |
|
|
|
|
| def sessions(enc_path, dec_path, threads=4): |
| so = ort.SessionOptions() |
| so.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL |
| so.intra_op_num_threads = threads |
| p = ["CPUExecutionProvider"] |
| return ort.InferenceSession(enc_path, so, providers=p), ort.InferenceSession(dec_path, so, providers=p) |
|
|
|
|
| def _ftype(sess, name): |
| t = next(i.type for i in sess.get_inputs() if i.name == name) |
| return np.float16 if "float16" in t else np.float32 |
|
|
|
|
| def encode(enc_sess, mel, mel_len): |
| ft = _ftype(enc_sess, "audio_signal") |
| enc, enc_len = enc_sess.run(None, {"audio_signal": mel.astype(ft), "length": np.array([mel_len], np.int64)}) |
| return enc, int(enc_len[0]) |
|
|
|
|
| def greedy_tdt(dec_sess, enc, T, num_layers=2, width=20480): |
| """Same rule as real-weights/scripts/verify_onnx.py greedy_tdt_onnx (NeMo label-looping).""" |
| ft = _ftype(dec_sess, "encoder_outputs") |
| st = _ftype(dec_sess, "input_states_1") |
| enc = enc.astype(ft) |
| s1 = np.zeros((num_layers, 1, width), st) |
| s2 = np.zeros_like(s1) |
| last = BLANK |
| tokens, t, last_ts, count = [], 0, -1, 0 |
|
|
| def step(frame): |
| out = dec_sess.run(None, { |
| "encoder_outputs": enc[:, :, frame:frame + 1], |
| "targets": np.array([[last]], np.int32), "target_length": np.array([1], np.int32), |
| "input_states_1": s1, "input_states_2": s2, |
| }) |
| logits = out[0][0, 0, 0] |
| k = int(logits[:-len(DURATIONS)].argmax()) |
| d = DURATIONS[int(logits[-len(DURATIONS):].argmax())] |
| if k == BLANK and d == 0: |
| d = 1 |
| return k, d, out[2], out[3] |
|
|
| while t < T: |
| while True: |
| k, d, n1, n2 = step(t) |
| 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 |
| s1, s2, last = n1, n2, k |
| if count >= MAX_SYMBOLS and last_ts == t: |
| t += 1 |
| return tokens |
|
|
|
|
| class Tokenizer: |
| def __init__(self, path): |
| self.sp = spm.SentencePieceProcessor(model_file=path) |
|
|
| def decode(self, ids): |
| return self.sp.decode(ids) |
|
|
