File size: 3,674 Bytes
19bb05c | 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 | """Minimal streaming inference for Muno459/fastconformer-quran-streaming (ONNX, cache-aware).
The ONNX takes 80-dim log-mel features (NOT raw audio) + carries cache tensors across chunks.
Pipeline: audio(16k mono) -> log-mel(80) -> fixed-global CMVN -> chunked encoder w/ cache -> CTC greedy.
EXACT shapes that fix the `Where`/`Add` broadcast errors (batch-first for ONNX):
cache_last_channel : float32 [1, 17, 70, 512] (init zeros) # 17 layers, 70 = left cache
cache_last_time : float32 [1, 17, 512, 8] (init zeros) # 8 = conv cache (kernel-1)
cache_last_channel_len : int64 [1] = [0] (init)
Carry the returned cache_*_next straight back in as the next step's cache_* (already batch-first).
Mel params (must match exactly if you reimplement on-device, e.g. iOS/Android):
sample_rate=16000, n_fft=512, win_length=400 (25ms hann), hop_length=160 (10ms),
n_mels=80 (librosa slaney-norm mel), power=2.0, log = ln(x + 2**-24), preemph=0.97, dither=0 (inference).
Then apply fixed-global CMVN from streaming_global_cmvn.npz: (mel - mean) / (std + 1e-5),
using the tlog_* constants for phone audio, clean_* for studio.
The simplest *correct* mel is NeMo's own preprocessor (one import); a hand-rolled mel must match the
params above or accuracy drops. Decode with the BPE tokenizer (tokenizer.model, blank id = 1024).
"""
import numpy as np, onnxruntime as ort, soundfile as sf, sentencepiece as spm
SESS = ort.InferenceSession("model.q8.onnx", providers=["CPUExecutionProvider"])
SP = spm.SentencePieceProcessor(model_file="tokenizer.model")
CMVN = np.load("streaming_global_cmvn.npz")
BLANK = 1024
N_LAYERS, D_MODEL, LEFT_CACHE, TIME_CACHE = 17, 512, 70, 8
CHUNK_MEL = 112 # mel frames per streaming step (8x subsampling -> ~14 output frames)
def log_mel(wav, sr=16000):
"""80-dim log-mel matching NeMo's FilterbankFeatures. Uses NeMo if available (exact),
else falls back to a torchaudio/librosa reimplementation with the documented params."""
import torch, torchaudio
wav = torch.tensor(np.asarray(wav, np.float32))
wav = torch.cat([wav[:1], wav[1:] - 0.97 * wav[:-1]]) # preemphasis
spec = torchaudio.transforms.MelSpectrogram(
sample_rate=sr, n_fft=512, win_length=400, hop_length=160,
n_mels=80, power=2.0, window_fn=torch.hann_window, norm="slaney", mel_scale="slaney")(wav)
return torch.log(spec + 2**-24).numpy() # (80, T)
def transcribe(path, kind="tlog"):
wav, sr = sf.read(path)
if wav.ndim > 1: wav = wav.mean(1)
mel = log_mel(wav, sr) # (80, T)
gm = CMVN[f"{kind}_mean"][:, None]; gs = CMVN[f"{kind}_std"][:, None]
mel = (mel - gm) / (gs + 1e-5)
clc = np.zeros((1, N_LAYERS, LEFT_CACHE, D_MODEL), np.float32)
clt = np.zeros((1, N_LAYERS, D_MODEL, TIME_CACHE), np.float32)
cll = np.zeros((1,), np.int64)
ids = []
for s in range(0, mel.shape[1], CHUNK_MEL):
chunk = mel[:, s:s + CHUNK_MEL][None].astype(np.float32) # (1,80,chunk)
length = np.array([chunk.shape[2]], np.int64)
logp, _, clc, clt, cll = SESS.run(None, {
"audio_signal": chunk, "length": length,
"cache_last_channel": clc, "cache_last_time": clt, "cache_last_channel_len": cll})
a = logp[0].argmax(-1)
prev = -1
for t in a:
t = int(t)
if t != prev and t != BLANK: ids.append(t)
prev = t
return SP.decode([i for i in ids])
if __name__ == "__main__":
import sys
print(transcribe(sys.argv[1], kind="tlog" if len(sys.argv) < 3 else sys.argv[2]))
|