"""asr4all ExecuTorch example: transcribe with an `executorch/xnnpack*` bundle. pip install executorch numpy python asr4all_executorch.py xnnpack_int8 clip.wav # punctuated + cased python asr4all_executorch.py xnnpack_int8 clip.wav --tier 2 # a different latency point python asr4all_executorch.py xnnpack_int8 clip.wav --stream # chunk by chunk python asr4all_executorch.py xnnpack_int8 --list-tiers PCEC is the model's second pass over the recognised tokens: punctuation, capitalisation and error correction. It runs a little behind the acoustic transcript; `--no-pcec` skips it. """ from __future__ import annotations import argparse import json import sys import wave from pathlib import Path import numpy as np import torch SAMPLE_RATE = 16_000 # Streaming method tensors, by position (fixed by the export): # in 0 = pcm, 1 = filter cache, 2 = attn cache # out 0 = logits, 1 = filter cache', 2 = attn cache', 3 = hidden (only with PCEC) T_PCM = T_LOGITS = 0 T_FILTER = 1 T_ATTN = 2 T_HIDDEN = 3 # ExecuTorch scalar-type ids -> torch dtypes. The caches are not all float: the PCEC # `seen` counter is int64, and zero-filling it as float fails set_inputs with # "expected Long but was Float". Always take the dtype the method declares. DTYPES = {4: torch.int64, 5: torch.float16, 6: torch.float32} # Blank-frame counts separating the pause buckets PCEC was trained on. PAUSE_THRESHOLDS = (1, 2, 4, 8, 16, 32, 64, 128) def _runtime(): """The ExecuTorch runtime, imported quietly. Importing it logs a torch pytree deprecation notice, warns that the Python bindings are experimental, and (through ExecuTorch's own `importlib.resources.open_text` call) raises a standard-library deprecation. None is about this script or the model, so they are held back here and nothing else is.""" import logging import warnings logging.getLogger("torch.utils._pytree").setLevel(logging.ERROR) with warnings.catch_warnings(): warnings.filterwarnings("ignore", message="This API is experimental") warnings.filterwarnings("ignore", category=DeprecationWarning, module="executorch") warnings.filterwarnings( "ignore", message="open_text is deprecated", category=DeprecationWarning ) from executorch.runtime import Runtime return Runtime def load_wav(path: str) -> np.ndarray: """16 kHz mono 16-bit .wav -> float32 in [-1, 1].""" with wave.open(path) as w: if w.getframerate() != SAMPLE_RATE or w.getnchannels() != 1: raise ValueError(f"{path}: want {SAMPLE_RATE} Hz mono") buf = np.frombuffer(w.readframes(w.getnframes()), dtype=np.int16) return buf.astype(np.float32) / 32768.0 def _pause_bucket(blanks: int) -> int: return sum(1 for t in PAUSE_THRESHOLDS if blanks >= t) class _Pcec: """ The token-space second pass for punctuation, capitalization and error correction """ def __init__(self, meta: dict, rung: dict, model, hidden_dim: int): self.meta, self.model = meta, model self.name, self.flush_name = rung["name"], rung["flush"] self.n_cache = int(rung.get("cache_tensors", 2)) sm = model.metadata(self.name) fm = model.metadata(self.flush_name) self.commit = int(sm.input_tensor_meta(0).sizes()[1]) flush_slots = int(fm.output_tensor_meta(0).sizes()[1]) # the model's own settle lag in tokens, read off the flush width self.lookahead = (flush_slots - 1) // 2 - self.commit n_feat = sm.num_inputs() - self.n_cache self.wants_acoustic = n_feat >= 3 self.wants_keep = n_feat == 4 self.dim = hidden_dim self._cache_spec = [ (tuple(sm.input_tensor_meta(i).sizes()), DTYPES[sm.input_tensor_meta(i).dtype()]) for i in range(sm.num_inputs() - self.n_cache, sm.num_inputs()) ] self.reset() @property def lag(self) -> tuple[int, int, int]: """(commit, lookahead, settle) in tokens, ie how far decode_pcec trails the audio.""" return (self.commit, self.lookahead, self.lookahead + self.commit - 1) def reset(self) -> None: self.cache = [torch.zeros(s, dtype=d) for s, d in self._cache_spec] self.ids: list[int] = [] self.pause: list[int] = [] self.acoustic: list[np.ndarray] = [] self.keep: list[float] = [] # per token, 1.0 = written by the word spotter self.slots: list[int] = [] self.consumed = 0 self._open = None # [id, hidden sum, frames, leading blanks] self._blanks = 0 self._prev = -1 def pool(self, frame_ids: list[int], hidden: np.ndarray, blank: int, keep=None) -> None: """One committed chunk of frames -> pooled tokens. `keep` (per frame, from a biasing decoder) pools to a per-token max: a token any of whose frames is marked is kept.""" for t, fid in enumerate(frame_ids): kf = float(keep[t]) if keep is not None else 0.0 if fid != self._prev and fid != blank: # a new token starts here if self._open: self._emit() self._open = [fid, hidden[t].astype(np.float64).copy(), 1, self._blanks, kf] self._blanks = 0 elif self._open: # a repeat or a blank belongs to the open token self._open[1] += hidden[t] self._open[2] += 1 self._open[4] = max(self._open[4], kf) self._prev = fid if fid == blank: self._blanks += 1 def flush(self) -> None: """Close the token still held back.""" if self._open: self._emit() def _emit(self) -> None: tok, acc, n, run, kf = self._open self._open = None self.ids.append(tok) self.acoustic.append((acc / n).astype(np.float32)) self.pause.append(_pause_bucket(run)) self.keep.append(kf) def drain(self) -> None: """Run every whole commit that is ready; leftovers stay buffered.""" k = self.commit while len(self.ids) >= k: outs = self.model.load_method(self.name).execute(self._feats(k, k) + self.cache) # The delay line holds `lookahead` tokens; until that many have gone through it, # its head is still warm-up from before the utterance. Drop those slots. skip = min(2 * k, max(0, 2 * self.lookahead - 2 * self.consumed)) self.slots += outs[0][0].tolist()[skip : 2 * k] self.cache = [t.clone() for t in outs[-self.n_cache :]] del self.ids[:k], self.pause[:k], self.acoustic[:k], self.keep[:k] self.consumed += k def tail(self) -> list[int]: """The last partial commit plus whatever is still inside the model.""" n = len(self.ids) feats = self._feats(self.commit, n) feats.append(torch.tensor(n, dtype=torch.long)) # n_real, before the cache outs = self.model.load_method(self.flush_name).execute(feats + self.cache) valid = 2 * self.lookahead + 2 * n + 1 skip = max(0, 2 * self.lookahead - 2 * self.consumed) return outs[0][0].tolist()[skip:valid] def _feats(self, k: int, n: int) -> list[torch.Tensor]: """A full block of `k` tokens, padded past the `n` real ones.""" ids = torch.full((1, k), int(self.meta["pad_in"]), dtype=torch.long) pause = torch.full((1, k), int(self.meta["pause_blackout"]), dtype=torch.long) for i in range(n): ids[0, i] = self.ids[i] # real_pause is a training decision, invisible in the graph # a model trained with the channel blacked out must be fed the blackout value if self.meta.get("real_pause"): pause[0, i] = self.pause[i] feats = [ids, pause] if self.wants_acoustic: ac = torch.zeros(1, k, self.dim) for i in range(n): ac[0, i] = torch.from_numpy(self.acoustic[i]) feats.append(ac) if self.wants_keep: # `keep`: per input token, how strongly PCEC should preserve it (0..1). A decoder that # biases recognition toward user vocabulary sets 1 on the words it inserted so PCEC # does not revert them. With no such decoder every value is 0 and PCEC runs as it # always did, which is what this script does unless set_dictionary() is used. kp = torch.zeros(1, k) for i in range(n): kp[0, i] = self.keep[i] feats.append(kp) return feats def render(self, slots) -> str: return render_slots(self.meta, slots) def render_slots(meta: dict, slots) -> str: """ Lattice slots -> punctuated, cased text. """ vocab, blank, ctrl = meta["vocab"], meta["blank_id"], meta["ctrl"] case = {ctrl["cap"]: "cap", ctrl["allcaps"]: "upper"} punct = {ctrl["period"]: ".", ctrl["comma"]: ",", ctrl["question"]: "?"} words: list[list] = [] # [pieces, case, punct] pending, prev = "", -1 for i in slots: if i in (prev, blank): # CTC collapse prev = i continue prev = i if i in case: pending = case[i] elif i in punct: if words: words[-1][2] = punct[i] elif i < len(vocab) and vocab[i]: if vocab[i].startswith(" ") or not words: words.append([[vocab[i]], pending, ""]) pending = "" else: words[-1][0].append(vocab[i]) out, sentence_start = [], True for pieces, cas, pnc in words: w = "".join(pieces).strip() if not w: continue if cas == "upper": w = w.upper() elif cas == "cap" or sentence_start: w = w[:1].upper() + w[1:] sentence_start = pnc in (".", "?") out.append(w + pnc) return " ".join(out) class Engine: """Load a bundle, pick a tier and decode.""" def __init__(self, bundle_dir: str): d = Path(bundle_dir) self.meta = json.loads((d / "metadata.json").read_text()) self.model = _runtime().get().load_program(str(d / self.meta["files"]["encoder"])) self.vocab, self.blank = self.meta["vocab"], self.meta["blank_id"] self.spf = int(self.meta["samples_per_frame"]) # ---- wake VAD ----------------------------------------------------------------------- # Present only when the bundle ships it. TWO methods compute the SAME scores; which one is # faster INVERTS by ISA: # x86 12900K (AVX-512) `vad` 0.263 ms `vad_dft` 0.648 ms # Pixel 10 (ARM NEON) `vad` 1.125 ms `vad_dft` 0.651 ms # `metadata.vad` is a list of methods; ARM takes `vad_dft` when shipped, anything else the # first wake entry, which is the conv frontend every other method here already uses. self.vad = None wake = [v for v in self.meta.get("vad") or [] if v["kind"] == "wake"] if wake: import platform as _plat arm = _plat.machine().lower() in ("arm64", "aarch64") vm = next((v for v in wake if arm and v["method"] == "vad_dft"), wake[0]) method = self.model.load_method(vm["method"]) self.vad = { "method": method, "name": vm["method"], # geometry from the model, like the tiers below "samples": int(method.metadata.input_tensor_meta(0).sizes()[-1]), "hop_s": float(vm["hop_s"]), "context_s": float(vm["context_s"]), } # Geometry comes from the model, not the metadata. Metadata carries names and the # vocab, and the schema deliberately does not restate shapes. self.tiers = [] for t in self.meta["tiers"]: mm = self.model.metadata(t["name"]) window = int(mm.input_tensor_meta(T_PCM).sizes()[-1]) commit = int(mm.output_tensor_meta(T_LOGITS).sizes()[1]) n_out = mm.num_outputs() self.tiers.append({ "name": t["name"], "pcec": t.get("pcec"), "window": window, "commit": commit, "lookahead": window // self.spf - commit, "n_out": n_out, "dim": int(mm.output_tensor_meta(T_HIDDEN).sizes()[-1]) if n_out > 3 else 0, }) # fmt: skip self.tiers.sort(key=lambda t: t["lookahead"] + t["commit"]) # the window: the dial self.spotter = None # personal dictionary, see set_dictionary() self.set_tier([t["name"] for t in self.tiers].index(self.meta["default_tier"])) def detect_speech(self, pcm): """Per-window speech probability, or None if this bundle ships no VAD. `pcm` is a float32 tensor of exactly `vad["samples"]` samples (the traced shape). Each returned score covers `context_s` of audio and consecutive scores are `hop_s` apart, so score w begins at w*hop_s seconds. Windows carry no state, so a caller may chunk freely -- but must keep a pre-roll of at least `context_s`, since the first score of a chunk describes the chunk's own opening window and speech starting just before it would be cut. """ if self.vad is None: return None return self.vad["method"].execute([pcm])[0] def _keep_only_methods_of(self, tier: dict) -> None: """Hold only the selected tier's methods in memory. Each loaded method carries its own packed copy of the weights (about one trunk's worth), so loading every tier would cost several times the program size; methods load on first use in about 60 ms, which is what a tier switch costs. The runtime caches every method it has loaded, so the ones this tier does not use are dropped from that cache here.""" cache = getattr(self.model, "_methods", None) if not isinstance(cache, dict): return keep = {tier["name"]} if tier.get("pcec"): keep |= {tier["pcec"]["name"], tier["pcec"]["flush"]} if getattr(self, "vad", None): keep.add(self.vad["name"]) for name in cache: if name not in keep: cache[name] = None # the runtime's "not loaded yet" marker def set_tier(self, idx: int | str) -> None: """Select a tier by index (0 = lowest latency) or by name (`stream_c256r32`).""" names = [t["name"] for t in self.tiers] if isinstance(idx, str): if idx.lstrip("-").isdigit(): idx = int(idx) elif idx in names: idx = names.index(idx) else: raise ValueError(f"unknown tier {idx!r}: have {names}") if not 0 <= idx < len(self.tiers): raise IndexError(f"tier {idx}: have 0..{len(self.tiers) - 1}") self.idx = idx t = self.tiers[idx] self._keep_only_methods_of(t) self.pcec = ( _Pcec(self.meta["pcec"], t["pcec"], self.model, t["dim"]) if self.meta.get("pcec") and t.get("pcec") else None ) self.reset() def reset(self) -> None: """Start a new utterance. Zero caches are the initialized state.""" mm = self.model.metadata(self.tiers[self.idx]["name"]) self.filter_cache, self.attn_cache = ( torch.zeros( tuple(mm.input_tensor_meta(i).sizes()), dtype=DTYPES[mm.input_tensor_meta(i).dtype()], ) for i in (T_FILTER, T_ATTN) ) self.pcm = np.zeros(0, dtype=np.float32) self.next_frame = 0 self.ids: list[int] = [] if self.pcec: self.pcec.reset() if self.spotter: self.spotter.reset() def pcec_lag(self): """(commit, lookahead, settle) in tokens, or None without PCEC.""" return self.pcec.lag if self.pcec else None def set_dictionary(self, words, **spotter_kw) -> None: """Contextual biasing hook: favour user vocabulary (names, jargon) where the audio nearly supports it, and mark those words so PCEC preserves them. It needs a word-spotting module (`word_spotter.py` beside this file) that is NOT part of this release; without it this method raises ImportError. Pass an empty list to switch biasing off, which is the default: PCEC then receives keep = 0 for every token and behaves exactly as before.""" if not words: # biasing off: needs nothing beyond this script self.spotter = None self.reset() return # imported as a sibling file, so this script depends on nothing beyond its own directory _here = str(Path(__file__).resolve().parent) if _here not in sys.path: sys.path.insert(0, _here) from word_spotter import build_spotter merges = self.meta.get("merges") if words and not merges: raise ValueError( "this bundle's metadata carries no `merges`; re-export it to use a dictionary" ) self.spotter = build_spotter(words, self.vocab, self.blank, merges or [], **spotter_kw) self.reset() def decode_stream(self, chunk: np.ndarray) -> str: """Feed audio, commit every complete window, return the transcript so far.""" t = self.tiers[self.idx] self.pcm = np.concatenate([self.pcm, chunk]) window_frames = t["window"] // self.spf while (self.next_frame + window_frames) * self.spf <= len(self.pcm): self._step(t) return self._text() def finish(self) -> str: """Flush the tail (zero-padding the missing lookahead) and close PCEC.""" t = self.tiers[self.idx] while self.next_frame * self.spf < len(self.pcm): self._step(t) if self.spotter: self._commit(*self.spotter.flush()) if self.pcec: self.pcec.flush() self.pcec.drain() return self._text() def decode_pcec(self, decode_tail: bool = False) -> str: """Punctuated + cased text. It trails the acoustic transcript by `pcec_lag()`; pass decode_tail=True at an endpoint to also render what is still inside the model.""" if not self.pcec: return self._text() slots = self.pcec.slots + (self.pcec.tail() if decode_tail else []) return self.pcec.render(slots) def decode_offline(self, clip: np.ndarray, apply_pcec: bool = True) -> str: """A whole clip: the streaming loop run to completion.""" self.reset() self.decode_stream(clip) self.finish() return self.decode_pcec(decode_tail=True) if (apply_pcec and self.pcec) else self._text() def _step(self, t: dict) -> None: """One method call: window + caches -> committed frames, updated caches.""" start = self.next_frame * self.spf # Zero-pad rather than shorten: the window is a fixed shape. No front trim. # The window starts at the commit boundary and left context comes from the caches. win = np.zeros(t["window"], dtype=np.float32) hi = min(len(self.pcm), start + t["window"]) if hi > start: win[: hi - start] = self.pcm[start:hi] outs = self.model.load_method(t["name"]).execute( [torch.from_numpy(win)[None, :], self.filter_cache, self.attn_cache] ) # outputs live in the method's arena, so copy before the next call self.filter_cache = outs[T_FILTER].clone() self.attn_cache = outs[T_ATTN].clone() valid = min(t["commit"], len(self.pcm) // self.spf - self.next_frame) lp = outs[T_LOGITS][0, :valid].float().numpy() frame_ids = lp.argmax(-1).tolist() hidden = ( outs[T_HIDDEN][0, :valid].float().numpy() if (self.pcec and t["n_out"] > T_HIDDEN) else None ) if self.spotter: # hold frames until no dictionary word can still rewrite them frame_ids, hidden, keep = self.spotter.feed(lp, frame_ids, hidden) self._commit(frame_ids, hidden, keep) else: self._commit(frame_ids, hidden) self.next_frame += t["commit"] def _commit(self, frame_ids: list[int], hidden, keep=None) -> None: """Frames that are final: into the transcript, and into PCEC (with per-frame keep marks when a biasing decoder supplies them).""" self.ids += frame_ids if self.pcec and hidden is not None and len(frame_ids): self.pcec.pool(frame_ids, hidden, self.blank, keep) self.pcec.drain() def _text(self) -> str: """CTC collapse of the accumulated frame ids: drop repeats, drop blanks.""" out, prev = [], -1 for i in self.ids: if i != prev and i != self.blank and i < len(self.vocab): out.append(self.vocab[i]) prev = i return " ".join("".join(out).split()) def load_engine(bundle_dir: str) -> Engine: return Engine(bundle_dir) def main() -> None: ap = argparse.ArgumentParser(description="asr4all ExecuTorch example") ap.add_argument("bundle_dir", help="e.g. executorch/xnnpack_int8") ap.add_argument("audio", nargs="?", help="16 kHz mono wav") ap.add_argument( "--tier", default=None, help="tier name (e.g. stream_c256r32) or index (0 = lowest latency)" ) ap.add_argument("--no-pcec", action="store_true", help="plain lowercase decode") ap.add_argument("--stream", action="store_true", help="feed in chunks, print each result") ap.add_argument("--list-tiers", action="store_true") args = ap.parse_args() e = load_engine(args.bundle_dir) if args.tier is not None: e.set_tier(args.tier) if e.pcec is not None and getattr(e.pcec, "wants_keep", False): print( "# note: this bundle's PCEC takes a `keep` input (per token, for contextual biasing); " "this script feeds it 0 everywhere, so PCEC behaves as it always did.", file=sys.stderr, ) if args.list_tiers: print(f"{'#':>2} {'tier':16} {'commit':>7} {'lookahead':>10} {'first result':>13} pcec") for i, t in enumerate(e.tiers): ttft = (t["commit"] + t["lookahead"]) * 40 / 1000 print(f"{i:>2} {t['name']:16} {t['commit']:>7} {t['lookahead']:>10} " f"{ttft:>10.2f} s {(t.get('pcec') or {}).get('name', '-')}") # fmt: skip return clip = load_wav(args.audio) if args.stream: for i in range(0, len(clip), 8000): # 0.5 s at a time, like a capture loop print(e.decode_stream(clip[i : i + 8000])) print(e.finish()) # the tail, which never got its full lookahead if e.pcec and not args.no_pcec: print("pcec:", e.decode_pcec(decode_tail=True)) return print(e.decode_offline(clip, apply_pcec=not args.no_pcec)) if __name__ == "__main__": main()