Add Pocket (300M, CPU) variant: int8/fp32 ONNX, PyTorch checkpoint, vocab, CPU inference script, results; README variants table
19eb87a verified | #!/usr/bin/env python3 | |
| """Dynamic ASR Pocket (300M, CPU): offline transcription with onnxruntime and greedy CTC decoding. | |
| pip install onnxruntime numpy soundfile scipy | |
| python transcribe_pocket.py clip.wav [more.wav ...] # int8 model (default) | |
| python transcribe_pocket.py clip.wav --model model_fp32.onnx # fp32 model | |
| Input: any wav; it is mixed to mono and resampled to 16 kHz. Clips longer than ~20 s are cut at quiet points. | |
| Model files are looked up next to this script: model_int8.onnx (or model_fp32.onnx) and vocab.json. | |
| Code licence: Apache-2.0. Model weights: CC-BY-SA-4.0. | |
| """ | |
| import argparse, json, os, time | |
| import numpy as np | |
| HERE = os.path.dirname(os.path.abspath(__file__)) | |
| SR = 16000 | |
| class PocketASR: | |
| def __init__(self, model=None, vocab=None, threads=4): | |
| import onnxruntime as ort | |
| so = ort.SessionOptions() | |
| so.intra_op_num_threads = threads | |
| so.inter_op_num_threads = 1 | |
| so.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL | |
| self.sess = ort.InferenceSession(model or os.path.join(HERE, "model_int8.onnx"), so, | |
| providers=["CPUExecutionProvider"]) | |
| v = json.load(open(vocab or os.path.join(HERE, "vocab.json"), encoding="utf-8")) | |
| self.pieces = v["pieces"] # index 0 = CTC blank; 1-3 = special tokens (never emitted) | |
| def logprobs(self, wav): | |
| # input "wav": float32 [1, T] at 16 kHz -> output "logprobs": [1, frames, vocab] | |
| return self.sess.run(None, {"wav": np.asarray(wav, dtype=np.float32)[None]})[0][0] | |
| def greedy(self, lp): | |
| out, prev = [], -1 | |
| for p in lp.argmax(-1): | |
| if p != prev and p > 3: | |
| out.append(self.pieces[p]) | |
| prev = p | |
| return " ".join("".join(out).split()) | |
| def transcribe(self, wav): | |
| if len(wav) < SR // 10: | |
| return "" | |
| return self.greedy(self.logprobs(wav)) | |
| def load_wav(path): | |
| import soundfile as sf | |
| y, sr = sf.read(path, dtype="float32", always_2d=True) | |
| y = y.mean(1) | |
| if sr != SR: | |
| from math import gcd | |
| from scipy.signal import resample_poly | |
| g = gcd(sr, SR) | |
| y = resample_poly(y, SR // g, sr // g).astype(np.float32) | |
| return y | |
| def transcribe_long(asr, y, win=20.0): | |
| """Audio over ~20 s: cut at the quietest 100 ms in the last 4 s of each window, transcribe pieces, join.""" | |
| if len(y) <= win * SR: | |
| return asr.transcribe(y) | |
| out, s = [], 0 | |
| while s < len(y): | |
| e = min(len(y), s + int(win * SR)) | |
| if e < len(y): | |
| lo = s + int((win - 4) * SR) | |
| fr = y[lo:e][: (e - lo) // 1600 * 1600].reshape(-1, 1600) | |
| e = lo + int(np.argmin((fr ** 2).mean(1))) * 1600 + 800 | |
| out.append(asr.transcribe(y[s:e])) | |
| s = e | |
| return " ".join(t for t in out if t) | |
| if __name__ == "__main__": | |
| ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) | |
| ap.add_argument("wav", nargs="+") | |
| ap.add_argument("--model", default=None, help="path to model_int8.onnx (default) or model_fp32.onnx") | |
| ap.add_argument("--vocab", default=None) | |
| ap.add_argument("--threads", type=int, default=4) | |
| a = ap.parse_args() | |
| asr = PocketASR(a.model, a.vocab, a.threads) | |
| for p in a.wav: | |
| y = load_wav(p) | |
| t = time.perf_counter() | |
| txt = transcribe_long(asr, y) | |
| dt = time.perf_counter() - t | |
| print(f"{p}\t{txt}\t[{len(y) / SR:.1f} s audio, {dt * 1000:.0f} ms, RTF {dt / max(len(y) / SR, 1e-3):.3f}]") | |