open-standards/aac/aac-rs/tools/asr/build_asr.py

#!/usr/bin/env python3
"""Build the listener's model files and their catalogue.

The listener ("Hearing your partner") turns what the person's partner says
into text on the device: NVIDIA Nemotron 3.5 ASR Streaming 0.6B writes the
words down, and Silero VAD finds where speech starts and ends. Both run on
onnxruntime-web's WebGPU execution provider. Their graphs are soniqo's fp16
ONNX export and Silero's own, edited by the asr-proto prototype so every node
runs on WebGPU (README.md says how, and where those tools live). This script
takes those edited files, checks each against the size and SHA-256 pinned
here (FILES), names each by its contents, and writes the catalogue. Every
file is at most 64 MiB, so the app's model store can check each one in
memory on a tablet and Cloudflare can cache it.

Outputs (schema "aac-asr-index/1", read by crates/aac-web/src/listen.rs):

  ../aac-asr-packs/                     beside the open-standards checkout ($AAC_ASR_DIR
                                        overrides; --out too): what the site builds copy
                                        to /app/asr/ after tools/asr/check_packs.py
    index.json                          pointer {schema, current}; serve no-store
    index.<sha12>.json                  the immutable catalogue
    LICENSES.md                         the licences and attribution
    licences/<name>.<sha12>.txt         the licence texts
    <model-id>/<file>.<sha12>.<ext>     the graphs, their weights in parts, the vocabulary
  data/asr/index.json                   committed copy of the catalogue
  data/asr/LICENSES.md                  committed copy of LICENSES.md

The languages, their status, prompt slots and the evidence for each come
from tools/asr/languages.json. A language is ready only when we measured it
on this app's own runtime (tools/asr/eval) at most MAX_ERROR_RATE percent.

Usage (from aac/aac-rs):
  python3 tools/asr/build_asr.py --source DIR     # DIR: the edited files (asr-proto's build/model)
  python3 tools/asr/build_asr.py                  # the same, from $AAC_ASR_SOURCE or tools/asr/.cache/model
  python3 tools/asr/build_asr.py --out DIR --no-commit
                                                  # a trial build elsewhere, data/asr left alone

Standard library only (Python 3.9 or later).
"""

import argparse
import hashlib
import json
import os
import re
import sys
from datetime import datetime, timezone
from pathlib import Path

HERE = Path(__file__).resolve().parent
ROOT = HERE.parents[1]  # aac/aac-rs
SOURCE_DEFAULT = Path(os.environ.get("AAC_ASR_SOURCE") or HERE / ".cache" / "model")
OUT_DEFAULT = Path(os.environ.get("AAC_ASR_DIR") or ROOT.parents[2] / "aac-asr-packs")
COMMITTED = ROOT / "data" / "asr"
SCHEMA = "aac-asr-index/1"
# The runtime the graphs are for: crates/aac-web/src/listen.rs RUNTIME_ID.
RUNTIME = "[email protected]"
PART_BYTES = 64 * 1024 * 1024
# A language is ready only with our measured error rate at most this (percent
# of words, or of characters for Chinese, Japanese and Korean): listen.rs
# MAX_ERROR_RATE.
MAX_ERROR_RATE = 15.0
PROMPT_SLOTS = 128
# The day the sources were read: the catalogue's date, and every file's
# modification time, so a rebuild of the same sources is byte- and
# time-identical (rsync then sends nothing).
RETRIEVED = "2026-10-01"
EPOCH = int(datetime.fromisoformat(RETRIEVED).replace(tzinfo=timezone.utc).timestamp())

MODEL_ID = "nemotron-3.5-asr-streaming-0.6b"
MODEL = {
    "name": "Nemotron 3.5 ASR Streaming 0.6B",
    "source": ("https://huggingface.co/soniqo/Nemotron-3.5-ASR-Streaming-Multilingual-0.6B-ONNX-FP16/tree/"
               "76daabfd0aaf5ec6ef1e6640eae3b364af6c9970 (NVIDIA's "
               "https://huggingface.co/nvidia/nemotron-3.5-asr-streaming-0.6b/tree/"
               "ea30d66debe3740a08b573244286791d423d6b3e), with "
               "https://github.com/snakers4/silero-vad/tree/v6.2.3 (silero_vad_16k_sequence.onnx); "
               "graphs edited to run on WebGPU"),
    # GPU memory while it listens, as measured (asr-proto: every buffer the
    # runtime asked the GPU for, Chrome on a Mac): 1.2 GiB loaded with the
    # VAD, up to 1.3 GB while it runs.
    "gpu_memory_bytes": 1_300_000_000,
    "min_memory_gb": 4,
}

# Every file, as asr-proto's tools name it (build/model/manifest.json), by
# size and SHA-256.
FILES = {
    "encoder_mel.onnx": (1749211, "0bb7c36cde260069c7e1ea2c802a3421fddf3bfd57eaa1739de653d7fee8005e"),
    "encoder.00.bin": (66916352, "cd3330f47afc7fc74e1b6fd1c359cf1ea42b04440d59f43d20d9c49143011826"),
    "encoder.01.bin": (65306624, "4560f03a6df0d5e6c78b0f4ca2b67058977d28d89cc2ffe727eacfc603639cb4"),
    "encoder.02.bin": (63483904, "814d178a255dec7a98538a741f656594ac509127cdcbeabcab2f904c12bb8491"),
    "encoder.03.bin": (65310720, "932a8c0ce8cf969e086dfce29d9210dfb8651d102acf7f1e7ac79f3df2d81188"),
    "encoder.04.bin": (65306624, "489a5bca9a56d7f76c947622b45ae1acfd248877573172cfdda34783a47286ae"),
    "encoder.05.bin": (63483904, "9235a5cc5c91d9f36ea29c5d862af3c60172b593715d500c19862e32f84b43a8"),
    "encoder.06.bin": (65310720, "c1ed0b1cd7736e533f651ec1dd47a46502f42eda55270fb497a1b6671505f024"),
    "encoder.07.bin": (65306624, "ae430eed432d5a43324e8508285a254996fef7be1cd4e7f8a3027cc37804e6ac"),
    "encoder.08.bin": (63483904, "27e0965650ac8bcea44a42b8cad98173f7df07f4d69c01df17a393588bcd9efa"),
    "encoder.09.bin": (65310720, "db56f6f3d5430b02adc0c95e7ef9914f89577a3a512d7115feac4cd0a8a5ab03"),
    "encoder.10.bin": (65306624, "fac3d9388bc1e85a633972390ff43b1b7bc75f19ca2274728174ad87b70877cd"),
    "encoder.11.bin": (63483904, "9e3b9006dbd3bf39bb6abbdbf0420b768c0eecdf6fc0eda871132aed65a29af1"),
    "encoder.12.bin": (65310720, "28c56605ebb2fa4793cae2fe23203bff292926421b1af6d38c1654f735ee3a14"),
    "encoder.13.bin": (65306624, "c1771281a15b7ae323d6ecdd05d99ab31b77cd05e2a3c14b292b58f3a132b727"),
    "encoder.14.bin": (63483904, "f2b92139ee2f3f8495d93b0f05cb3454c577e5ce114f5fbfd62a6d0d0449ce07"),
    "encoder.15.bin": (65310720, "6f96b19f25da1311741f9201d542def9879e5b0a2fa39365ae141e268ae6ce0d"),
    "encoder.16.bin": (65306624, "008698bea20628ef4247e3f6d1e93662e0d945cf5a0786f454907deb0282ec8f"),
    "encoder.17.bin": (63483904, "267b6a6d76f596f60fde37e8ffbe19da4ea4bc9fdee80fa9c97e2b9872f9db4d"),
    "encoder.18.bin": (25700352, "747068f82cbfb62cbf596ff8fcac1b02a7caf00adbc52bff65d3215984562e24"),
    "decoder_joint.onnx": (6619, "352f24fbd875c0d1c4fc41feca29eef4c2a71d2e8ffd53474174cddf571bdb77"),
    "decoder_joint.00.bin": (48781376, "82d9d90267b5ee8441cdb4ed58e5441a988b0d1acbd5aa5fe84a1a7769716d37"),
    "silero_vad_320ms.onnx": (1262546, "1eec45a9a632b30e9aa8077b01adb31a4c7bd5b1caa2eba363efae0e68e48e2a"),
    "vocab.json": (236127, "fd726dc25cc0c0b46403675127e831bab56ca56b8c7ae3ee7781f5c6d7692f4f"),
}
# The three graphs: the encoder (log-mel front end + cache-aware encoder),
# the step (prediction net + joint + ArgMax), the VAD; each with the
# external-data files it names, in order.
GRAPHS = {
    "encoder": ("encoder_mel.onnx", [f"encoder.{i:02d}.bin" for i in range(19)]),
    "step": ("decoder_joint.onnx", ["decoder_joint.00.bin"]),
    "vad": ("silero_vad_320ms.onnx", []),
}
VOCAB = "vocab.json"

# The licences, verbatim, by SHA-256: OpenMDW-1.1 from github.com/OpenMDW/OpenMDW
# (1.1/LICENSE.OpenMDW-1.1 at b26b32b3), Silero VAD's LICENSE at commit 5cd79456 (v6.2.3).
LICENCES = {
    "OpenMDW-1.1": (HERE / "licences" / "OpenMDW-1.1.txt",
                    "2ab44b68365473c112f5092211a38f231cb23e50de68b75a13369adbd76a74df"),
    "MIT-silero-vad": (HERE / "licences" / "MIT-silero-vad.txt",
                       "2e63e9a38b6e8fc0c7bc37ce174caca1862870856c6daf5697cfb785e925520b"),
}
CREDITS = [
    {
        "what": "speech recognition",
        "attribution": ("NVIDIA Nemotron 3.5 ASR Streaming 0.6B by NVIDIA, in soniqo's fp16 ONNX export, its graphs "
                        "edited to run on WebGPU"),
        "licence": "OpenMDW-1.1",
        "licence_url": "https://openmdw.ai/license/1-1/",
        "licence_key": "OpenMDW-1.1",
    },
    {
        "what": "finding speech",
        "attribution": "Silero VAD v6.2.3 by the Silero Team, its graph edited to run on WebGPU",
        "licence": "MIT",
        "licence_url": "https://github.com/snakers4/silero-vad/blob/v6.2.3/LICENSE",
        "licence_key": "MIT-silero-vad",
    },
]


class BuildError(Exception):
    pass


def log(msg):
    print(msg, flush=True)


def sha256_bytes(data):
    return hashlib.sha256(data).hexdigest()


def sha256_file(path):
    h = hashlib.sha256()
    with open(path, "rb") as f:
        for chunk in iter(lambda: f.read(1 << 22), b""):
            h.update(chunk)
    return h.hexdigest()


def dumps(obj):
    return (json.dumps(obj, indent=1, ensure_ascii=False) + "\n").encode("utf-8")


def languages():
    doc = json.loads((HERE / "languages.json").read_text(encoding="utf-8"))
    langs = doc["languages"]
    for code, l in langs.items():
        if not re.fullmatch(r"[a-z]{2,3}", code):
            raise BuildError(f"languages.json: code {code!r}")
        if not l.get("name") or not l.get("note"):
            raise BuildError(f"languages.json: {code} needs a name and a note")
        slots = [l["slot"]] if "slot" in l else []
        slots += list(l.get("locales", {}).values())
        if any(not (isinstance(s, int) and 0 <= s < PROMPT_SLOTS) for s in slots):
            raise BuildError(f"languages.json: {code} has a prompt slot out of range")
        if l["status"] == "ready":
            if l.get("model") != MODEL_ID:
                raise BuildError(f"languages.json: {code} is ready with an unknown model {l.get('model')!r}")
            if "slot" not in l:
                raise BuildError(f"languages.json: {code} is ready without a prompt slot")
            if l.get("measured") is None or l["measured"] > MAX_ERROR_RATE:
                raise BuildError(f"languages.json: {code} is ready without our measured error rate of at most {MAX_ERROR_RATE}")
            if not l.get("source"):
                raise BuildError(f"languages.json: {code} is ready without saying where its error rate comes from")
        elif l["status"] != "not-yet":
            raise BuildError(f"languages.json: {code} status {l['status']!r}")
    return dict(sorted(langs.items()))


def check_sources(source):
    """Every pinned file in `source`, with its size and SHA-256 checked."""
    found = {}
    for name, (size, digest) in FILES.items():
        path = source / name
        if not path.is_file():
            raise BuildError(f"{path} is missing")
        if path.stat().st_size != size:
            raise BuildError(f"{path} is {path.stat().st_size} bytes; the pin is {size}")
        got = sha256_file(path)
        if got != digest:
            raise BuildError(f"{path} has sha256 {got}; the pin is {digest}")
        if size > PART_BYTES:
            raise BuildError(f"{name} is larger than 64 MiB")
        found[name] = path
    # Each graph names its weights files by these names (ONNX external data
    # locations are plain strings in the graph).
    for graph, data in GRAPHS.values():
        blob = found[graph].read_bytes()
        for d in data:
            if d.encode() not in blob:
                raise BuildError(f"{graph} does not name {d}")
    return found


def published_name(name, digest):
    stem, dot, ext = name.rpartition(".")
    return f"{MODEL_ID}/{stem}.{digest[:12]}.{ext}"


def plan(found):
    """Every file to write, as rel -> ("copy", src) or ("data", bytes), and the catalogue."""
    files = {}
    licence_rel = {}
    for key, (path, digest) in LICENCES.items():
        data = path.read_bytes()
        if sha256_bytes(data) != digest:
            raise BuildError(f"{path} is not the licence text pinned here")
        rel = f"licences/{key}.{digest[:12]}.txt"
        files[rel] = ("data", data)
        licence_rel[key] = rel

    def part(name):
        size, digest = FILES[name]
        rel = published_name(name, digest)
        files[rel] = ("copy", found[name])
        return {"url": rel, "bytes": size, "sha256": digest}

    model = {"name": MODEL["name"], "runtime": RUNTIME, "source": MODEL["source"], "bytes": sum(s for s, _ in FILES.values())}
    for key, (graph, data) in GRAPHS.items():
        model[key] = {"file": part(graph), "data": [{"name": d, **part(d)} for d in data]}
    model["vocab"] = part(VOCAB)
    model["credits"] = [
        {k: v for k, v in c.items() if k != "licence_key"} | {"licence_file": licence_rel[c["licence_key"]]}
        for c in CREDITS
    ]
    model["gpu_memory_bytes"] = MODEL["gpu_memory_bytes"]
    model["min_memory_gb"] = MODEL["min_memory_gb"]
    catalogue = {"schema": SCHEMA, "generated": RETRIEVED, "models": {MODEL_ID: model}, "languages": languages()}
    cat_bytes = dumps(catalogue)
    cat_name = f"index.{sha256_bytes(cat_bytes)[:12]}.json"
    files[cat_name] = ("data", cat_bytes)
    files["index.json"] = ("data", dumps({"schema": SCHEMA, "current": cat_name}))
    files["LICENSES.md"] = ("data", licences_md(catalogue).encode("utf-8"))
    return files, cat_bytes


def licences_md(catalogue):
    m = catalogue["models"][MODEL_ID]
    ready = [f"{l['name']} ({c})" for c, l in catalogue["languages"].items() if l["status"] == "ready"]
    out = [
        "# Hearing your partner: licences\n\n"
        "The listener in AuraSense AAC (aac-web) downloads one speech model and a voice activity detector, on "
        "request, and runs them on the person's device, on its graphics chip. What is heard there is turned into "
        "text there: nothing is recorded or sent anywhere.\n",
        f"\n## {m['name']}\n\n",
    ]
    for c in m["credits"]:
        out.append(
            f"- {c['attribution']}. Licence: {c['licence']} ({c['licence_url']}); the text is "
            f"`{c['licence_file']}`.\n"
        )
    out.append(
        f"- Source: {m['source']}.\n"
        f"- What was changed: the graphs only, so that every node runs on WebGPU in onnxruntime-web 1.30 "
        f"(shapes folded, the attention mask rebuilt in int32, the mel front end merged into the encoder, the "
        f"prediction net, joint and ArgMax fused into one step, the LSTMs unrolled, the VAD's first convolution "
        f"padded to 132 channels); the weights are NVIDIA's and Silero's own, in {len(FILES)} files of at most 64 MiB, "
        f"each named by its own SHA-256 ({m['bytes'] / 1e9:.2f} GB).\n"
        f"- soniqo's model card names the NVIDIA Open Model License, the licence NVIDIA's model had when soniqo "
        f"exported it; NVIDIA has since published the model under OpenMDW-1.1 (its model card, revision "
        f"ea30d66d), which this app follows.\n"
        f"- Ready in: {', '.join(ready) or 'no language yet'}.\n"
        "\nThe runtime that runs it (onnxruntime-web, MIT) is vendored with its own licence, notice and ONNX "
        "Runtime's third-party notices (`ThirdPartyNotices.txt`) under "
        "`vendor/onnxruntime-web-<version>-<revision>/`.\n"
    )
    return "".join(out)


def write_tree(out, files):
    if out.exists() and any(out.iterdir()) and not (out / "index.json").exists():
        raise BuildError(f"refusing to prune non-empty {out}: it has no index.json from an earlier build")
    out.mkdir(parents=True, exist_ok=True)

    # Files and licences first, the catalogue after them, the pointer last: a
    # run that stops part way never leaves a pointer to files not there yet.
    def order(rel):
        if rel == "index.json":
            return (3, rel)
        if rel == "LICENSES.md":
            return (2, rel)
        return (1 if re.fullmatch(r"index\.[0-9a-f]{12}\.json", rel) else 0, rel)

    for rel in sorted(files, key=order):
        kind, spec = files[rel]
        dest = out / rel
        dest.parent.mkdir(parents=True, exist_ok=True)
        if kind == "copy":
            size = spec.stat().st_size
            fresh = dest.exists() and dest.stat().st_size == size and sha256_file(dest) == sha256_file(spec)
            if not fresh:
                tmp = dest.with_name(dest.name + ".tmp")
                with open(spec, "rb") as f, open(tmp, "wb") as o:
                    for chunk in iter(lambda: f.read(1 << 22), b""):
                        o.write(chunk)
                os.replace(tmp, dest)
        elif not (dest.exists() and dest.read_bytes() == spec):
            tmp = dest.with_name(dest.name + ".tmp")
            tmp.write_bytes(spec)
            os.replace(tmp, dest)
        os.chmod(dest, 0o644)
        os.utime(dest, (EPOCH, EPOCH))
    keep = set(files)
    for p in sorted(out.rglob("*"), reverse=True):
        rel = p.relative_to(out).as_posix()
        if p.is_file() and rel not in keep:
            p.unlink()
            log(f"prune: {rel}")
        elif p.is_dir() and not any(p.iterdir()):
            p.rmdir()


def main():
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--source", type=Path, default=SOURCE_DEFAULT,
                    help="the edited graphs, weights and vocabulary (asr-proto's build/model)")
    ap.add_argument("--out", type=Path, default=OUT_DEFAULT)
    ap.add_argument("--no-commit", action="store_true", help="leave data/asr alone (a trial build)")
    args = ap.parse_args()
    try:
        if args.no_commit and args.out == OUT_DEFAULT:
            raise BuildError("a trial build (--no-commit) goes elsewhere: give --out DIR")
        found = check_sources(args.source)
        files, cat_bytes = plan(found)
        write_tree(args.out, files)
        if not args.no_commit:
            COMMITTED.mkdir(parents=True, exist_ok=True)
            (COMMITTED / "index.json").write_bytes(cat_bytes)
            (COMMITTED / "LICENSES.md").write_bytes(files["LICENSES.md"][1])
    except BuildError as e:
        print(f"✗ {e}", file=sys.stderr)
        return 1
    n = sum(1 for kind, _ in files.values() if kind == "copy")
    log(f"✓ {args.out}: {MODEL_ID}, {n} files, catalogue {json.loads(files['index.json'][1])['current']}")
    return 0


if __name__ == "__main__":
    sys.exit(main())