File size: 4,599 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 | """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)
# The Haas checkpoint stores the preprocessor's window and filterbank in bfloat16 and NeMo computes with those,
# so round ours the same way to match the official features.
_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)) # symmetric Hann, as torch.hann_window(periodic=False)
_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 # NeMo get_seq_len: (len + 2*(n_fft//2) - n_fft) // hop
x = np.concatenate([x[:1], x[1:] - np.float32(0.97) * x[:-1]])
x = np.pad(x, N_FFT // 2) # center=True, pad_mode="constant"
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 # [257, T], power 2.0
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)
|