| |
| """Streaming zipformer (zh-en) board runtime for AX650 AXMODELs. |
| |
| Two input modes: |
| --wav FILE|DIR read 16 kHz mono PCM WAV, compute 80-dim kaldi fbank |
| (kaldi-native-fbank, same options as the C++ runtime) |
| --features FILE|DIR read precomputed fbank features (.npy, shape [T,80]) |
| |
| Runs encoder/decoder/joiner chunk by chunk (chunk=32, encoder input T=39) and |
| carries the 35 encoder states. Requires numpy + axengine; the WAV mode also |
| needs kaldi-native-fbank. |
| """ |
| from __future__ import annotations |
|
|
| import argparse |
| import glob |
| import json |
| import os |
| import time |
| import wave |
| from pathlib import Path |
| from typing import Dict, List |
|
|
| import numpy as np |
| from axengine import InferenceSession |
|
|
|
|
| def load_tokens(path: str) -> Dict[int, str]: |
| table: Dict[int, str] = {} |
| with open(path, encoding="utf-8") as f: |
| for line in f: |
| token, _, idx = line.rstrip("\n").rpartition(" ") |
| if token: |
| table[int(idx)] = token |
| return table |
|
|
|
|
| def tokens_to_text(tokens: List[int], table: Dict[int, str]) -> str: |
| return "".join(table.get(i, "") for i in tokens).replace("\u2581", " ").strip() |
|
|
|
|
| def read_wav(path: str) -> np.ndarray: |
| with wave.open(path, "rb") as w: |
| if w.getnchannels() != 1 or w.getframerate() != 16000 or w.getsampwidth() != 2: |
| raise ValueError( |
| f"{path}: need 16 kHz mono 16-bit PCM WAV " |
| f"(channels={w.getnchannels()}, rate={w.getframerate()}, " |
| f"sampwidth={w.getsampwidth()})" |
| ) |
| raw = w.readframes(w.getnframes()) |
| audio = np.frombuffer(raw, dtype=np.int16).astype(np.float32) / 32768.0 |
| return np.concatenate([audio, np.zeros(int(0.3 * 16000), dtype=np.float32)]) |
|
|
|
|
| def compute_fbank(audio: np.ndarray) -> np.ndarray: |
| try: |
| import kaldi_native_fbank as knf |
| except ImportError as exc: |
| raise SystemExit( |
| "WAV mode needs kaldi-native-fbank: pip install kaldi-native-fbank" |
| ) from exc |
| opts = knf.FbankOptions() |
| opts.frame_opts.samp_freq = 16000 |
| opts.frame_opts.dither = 0.0 |
| opts.frame_opts.frame_length_ms = 25.0 |
| opts.frame_opts.frame_shift_ms = 10.0 |
| opts.frame_opts.snip_edges = False |
| opts.frame_opts.window_type = "povey" |
| opts.mel_opts.num_bins = 80 |
| opts.mel_opts.low_freq = 20.0 |
| opts.mel_opts.high_freq = -400.0 |
| opts.energy_floor = 0.0 |
| fbank = knf.OnlineFbank(opts) |
| fbank.accept_waveform(16000, audio) |
| fbank.input_finished() |
| return np.stack([fbank.get_frame(i) for i in range(fbank.num_frames_ready)]) |
|
|
|
|
| class StreamingZipformer: |
| def __init__(self, model_dir: str, tokens: str, chunk: int = 32): |
| self.chunk = chunk |
| self.T = chunk + 7 |
| self.enc = InferenceSession(str(Path(model_dir) / "encoder.axmodel")) |
| self.dec = InferenceSession(str(Path(model_dir) / "decoder.axmodel")) |
| self.join = InferenceSession(str(Path(model_dir) / "joiner.axmodel")) |
| self.tokens = load_tokens(tokens) |
|
|
| self.enc_out_names = [o.name for o in self.enc.get_outputs()] |
| self.state_names = [i.name for i in self.enc.get_inputs() if i.name != "x"] |
| self.state_shapes = {i.name: list(i.shape) for i in self.enc.get_inputs()} |
| self.dec_in_name = self.dec.get_inputs()[0].name |
| self.dec_out_name = self.dec.get_outputs()[0].name |
| self.join_in_names = [i.name for i in self.join.get_inputs()] |
| self.join_out_name = self.join.get_outputs()[0].name |
| self.blank = 0 |
| self.context = 2 |
|
|
| def init_states(self) -> Dict[str, np.ndarray]: |
| states = {} |
| for name in self.state_names: |
| shape = self.state_shapes[name] |
| dtype = np.int32 if name.startswith("cached_len") else np.float32 |
| states[name] = np.zeros(shape, dtype=dtype) |
| return states |
|
|
| def reset(self) -> None: |
| self.states = self.init_states() |
| self.hyp: List[int] = [self.blank] * self.context |
| y = np.asarray([self.hyp[-self.context:]], dtype=np.int32) |
| self.dec_out = self.dec.run([self.dec_out_name], {self.dec_in_name: y})[0] |
|
|
| def run_encoder(self, x: np.ndarray) -> np.ndarray: |
| feed = {"x": np.ascontiguousarray(x, dtype=np.float32)} |
| feed.update(self.states) |
| outs = self.enc.run(self.enc_out_names, feed) |
| out = dict(zip(self.enc_out_names, outs)) |
| for name in self.state_names: |
| self.states[name] = out["new_" + name] |
| return out["encoder_out"] |
|
|
| def decode_features(self, feats: np.ndarray) -> dict: |
| self.reset() |
| n = feats.shape[0] |
| for start in range(0, n, self.chunk): |
| seg = feats[start:start + self.T] |
| if seg.shape[0] < self.T: |
| seg = np.concatenate( |
| [seg, np.zeros((self.T - seg.shape[0], 80), dtype=np.float32)], axis=0) |
| seg = np.ascontiguousarray(seg, dtype=np.float32) |
| enc_out = self.run_encoder(seg[None, :, :]) |
| for t in range(enc_out.shape[1]): |
| cur = enc_out[:, t:t + 1, :].reshape(1, -1).astype(np.float32) |
| logits = self.join.run( |
| [self.join_out_name], |
| {self.join_in_names[0]: cur, self.join_in_names[1]: self.dec_out}, |
| )[0] |
| sym = int(logits.reshape(-1).argmax()) |
| if sym != self.blank: |
| self.hyp.append(sym) |
| y = np.asarray([self.hyp[-self.context:]], dtype=np.int32) |
| self.dec_out = self.dec.run( |
| [self.dec_out_name], {self.dec_in_name: y})[0] |
| ids = self.hyp[self.context:] |
| return {"tokens": ids, "text": tokens_to_text(ids, self.tokens)} |
|
|
|
|
| def collect(path: str, suffix: str) -> List[str]: |
| p = Path(path) |
| if p.is_dir(): |
| return sorted(glob.glob(str(p / f"*{suffix}"))) |
| return [str(p)] |
|
|
|
|
| def main() -> None: |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--model-dir", required=True) |
| ap.add_argument("--tokens", required=True) |
| src = ap.add_mutually_exclusive_group(required=True) |
| src.add_argument("--wav", help="16 kHz mono PCM WAV file or directory") |
| src.add_argument("--features", help=".npy file or directory (shape [T,80])") |
| ap.add_argument("--out", default="", help="optional result JSON") |
| ap.add_argument("--chunk", type=int, default=32) |
| ap.add_argument("--stats", action="store_true", help="print load/inference timing") |
| args = ap.parse_args() |
|
|
| t0 = time.perf_counter() |
| asr = StreamingZipformer(args.model_dir, args.tokens, args.chunk) |
| load_seconds = time.perf_counter() - t0 |
| if args.wav: |
| files = collect(args.wav, ".wav") |
| audio = {os.path.basename(f): read_wav(f) for f in files} |
| feats = {k: compute_fbank(v) for k, v in audio.items()} |
| else: |
| files = collect(args.features, ".npy") |
| feats = {os.path.basename(f): np.load(f).astype(np.float32) for f in files} |
|
|
| result = {} |
| for name, f in feats.items(): |
| t1 = time.perf_counter() |
| r = asr.decode_features(f) |
| infer = time.perf_counter() - t1 |
| result[name] = r |
| print(f"{name}: {r['text']}") |
| if args.stats: |
| msg = f" load={load_seconds:.3f}s infer={infer:.3f}s" |
| if args.wav: |
| dur = len(audio[name]) / 16000.0 |
| msg += f" audio={dur:.2f}s rtf={infer / max(dur, 1e-9):.3f}" |
| print(msg) |
| if args.out: |
| Path(args.out).write_text( |
| json.dumps(result, ensure_ascii=False, indent=2), encoding="utf-8") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
