asr4all-m-plus / executorch /asr4all_executorch.py

asr4all_executorch.py: hold back ExecuTorch's own open_text deprecation warning

3e15eea verified 26 days ago

23.2 kB

"""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()

	"""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()