Moonshine Streaming Model(s) Recipe

usefulSensors-moonshine-streaming/cpu/README.md

@@ -0,0 +1,127 @@
+# Moonshine Streaming (CPU EP)
+
+This recipe exports **usefulsensors/moonshine-streaming-tiny** (default) or
+**usefulsensors/moonshine-streaming-small** to ONNX and produces CPU-ready
+ONNX Runtime GenAI artifacts for streaming ASR. Pick the variant with
+`--model-name` (see [Run](#run)).
+
+The streaming model is exported as **five FP32 ONNX components**, each handled
+through Olive's declarative `OnnxConversion` pass (dynamo exporter, opset 20)
+with the exact input/output names the streaming runner expects:
+
+- **Frontend** – stateful convolutional feature extractor. Consumes an audio
+  chunk plus rolling sample/conv buffers and emits log-mel-like features while
+  updating its buffers.
+- **Encoder** – sliding-window transformer encoder (features → hidden states).
+- **Adapter** – positional projection (`encoded + pos_emb(arange(T) + pos_offset)`)
+  producing the decoder memory.
+- **Cross-KV** – precomputes per-layer cross-attention K/V from the memory once
+  per segment.
+- **Decoder-KV** – autoregressive decoder with cached self-attention KV and the
+  precomputed cross-attention KV → logits.
+
+## Files
+- `cpu/moonshine_frontend_fp32_cpu.json` – Olive frontend config (convert only)
+- `cpu/moonshine_encoder_fp32_cpu.json` – Olive encoder config (convert only)
+- `cpu/moonshine_adapter_fp32_cpu.json` – Olive adapter config (convert only)
+- `cpu/moonshine_cross_kv_fp32_cpu.json` – Olive cross-KV config (convert only)
+- `cpu/moonshine_decoder_kv_fp32_cpu.json` – Olive decoder-KV config (convert only)
+- `cpu/moonshine_encoder_kquant8_cpu.json` – Olive encoder config (convert → INT8 k-quant)
+- `cpu/moonshine_decoder_kv_kquant8_cpu.json` – Olive decoder-KV config (convert → INT8 k-quant)
+- `cpu/moonshine_model_load.py` – model loaders + wrapper modules + dummy inputs
+- `cpu/optimize.py` – full pipeline script (Olive × 5 + tokenizer + configs + VAD)
+- `cpu/export_moonshine_streaming.py` – standalone exporter (no Olive; for debugging)
+
+## Setup
+From repo root:
+
+```bash
+python -m venv .venv
+source .venv/bin/activate
+pip install -r usefulSensors-moonshine-streaming/cpu/requirements.txt
+```
+
+## Run
+
+From the `usefulSensors-moonshine-streaming` directory:
+
+```bash
+cd usefulSensors-moonshine-streaming
+
+python cpu/optimize.py --output-dir build/moonshine-tiny
+```
+
+This runs the full pipeline:
+1. **Frontend / Encoder / Adapter / Cross-KV / Decoder-KV** — Olive: `OnnxConversion` (FP32) for each of the five components
+2. **Configs** — generates `genai_config.json` + `streaming_config.json`
+3. **Tokenizer** — exports `tokenizer.json` + `tokenizer_config.json`
+4. **VAD** — downloads Silero VAD ONNX model
+
+`--output-dir` is resolved relative to the `cpu/` directory unless an absolute
+path is given.
+
+Export the **small** variant instead:
+
+```bash
+python cpu/optimize.py --model-name usefulsensors/moonshine-streaming-small \
+    --output-dir build/moonshine-small
+```
+
+### Quantization (`--quantize`)
+
+Add `--quantize` to quantize the **encoder** and **decoder_kv** MatMuls
+(frontend / adapter / cross_kv stay FP32). IO names and the runtime configs
+are unchanged. `--quant-method` picks the algorithm:
+
+- `--quant-method kquant8` (default) — INT8 **weight-only k-quant**. Swaps in
+  `moonshine_encoder_kquant8_cpu.json` / `moonshine_decoder_kv_kquant8_cpu.json`,
+  which chain `OnnxConversion → OnnxKQuantQuantization` (`bits=8`, `block_size=32`;
+  weight matmuls become `MatMulNBits`). Weight-only, so activation×activation
+  attention matmuls stay FP32.
+- `--quant-method kquant8-enc` — same INT8 k-quant, **encoder only**;
+  `decoder_kv` stays FP32. Use when decoder quantization degrades transcription
+  quality and you can afford the larger decoder.
+
+```bash
+# INT8 k-quant on encoder + decoder_kv (default)
+python cpu/optimize.py --quantize --output-dir build/moonshine-tiny-kquant8
+
+# INT8 k-quant on encoder only, FP32 decoder
+python cpu/optimize.py --quantize --quant-method kquant8-enc \
+    --output-dir build/moonshine-tiny-kquant8-enc
+```
+
+On a 40s clip (CPU EP), INT8 k-quant preserves transcription quality:
+
+| variant | build | encoder | decoder_kv | total | RTF |
+|---|---|---|---|---|---|
+| small | FP32 | 168 MB | 309 MB | ~541 MB | ~7.0× |
+| small | INT8 k-quant (`--quantize`) | ~40 MB | ~120 MB | ~230 MB | ~8–9× |
+| small | official `.ort` | 42 MB | 174 MB | — | ~9.4× |
+| tiny  | INT8 k-quant (`--quantize`) | ~8 MB | ~64 MB | ~94 MB | ~17× |
+
+Transcription is essentially identical to FP32 (only quant-noise wording
+drift).
+
+Or run individual components directly with the Olive CLI:
+
+```bash
+python -m olive run --config cpu/moonshine_frontend_fp32_cpu.json
+python -m olive run --config cpu/moonshine_encoder_fp32_cpu.json
+python -m olive run --config cpu/moonshine_adapter_fp32_cpu.json
+python -m olive run --config cpu/moonshine_cross_kv_fp32_cpu.json
+python -m olive run --config cpu/moonshine_decoder_kv_fp32_cpu.json
+```
+
+## Output
+Expected artifacts in `cpu/build/moonshine-tiny/`:
+- `frontend.onnx` (+ `frontend.onnx.data`)
+- `encoder.onnx` (+ `encoder.onnx.data`)
+- `adapter.onnx` (+ `adapter.onnx.data`)
+- `cross_kv.onnx` (+ `cross_kv.onnx.data`)
+- `decoder_kv.onnx` (+ `decoder_kv.onnx.data`)
+- `genai_config.json`
+- `streaming_config.json`
+- `tokenizer.json`
+- `tokenizer_config.json`
+- `silero_vad.onnx`

usefulSensors-moonshine-streaming/cpu/export_moonshine_streaming.py

@@ -0,0 +1,173 @@
+"""Standalone exporter: HuggingFace MoonshineStreaming -> five stateful ONNX
+graphs for onnxruntime-genai (frontend / encoder / adapter / cross_kv /
+decoder_kv).
+
+Run inside a Python env with the moonshine transformers integration installed:
+
+    conda activate moonshine   # or: source .venv/bin/activate
+    python export_moonshine_streaming.py \
+        --model usefulsensors/moonshine-streaming-tiny \
+        --output-dir build/moonshine-streaming-tiny
+
+The exporter uses the TorchDynamo ONNX path (``dynamo=True``) so data-dependent
+shapes in the stateful frontend export cleanly, then rewrites every graph's
+input/output names to the exact contract genai expects.
+"""
+
+from __future__ import annotations
+
+import argparse
+import os
+
+import onnx
+import torch
+from torch.export import Dim
+
+import moonshine_model_load as mml
+
+AUTO = Dim.AUTO
+
+
+# --------------------------------------------------------------------------- #
+# Per-component export specification                                          #
+# --------------------------------------------------------------------------- #
+def build_specs():
+    """Return the ordered list of component export specs.  ``dynamic`` is a
+    tuple aligned with the positional dummy inputs: each entry is either a
+    ``{axis: Dim.AUTO}`` dict or ``None`` for a fully static input."""
+    return [
+        {
+            "name": "frontend",
+            "loader": mml.frontend_model_loader,
+            "dummy": mml.frontend_dummy_inputs,
+            "input_names": [
+                "audio_chunk", "sample_buffer", "sample_len",
+                "conv1_buffer", "conv2_buffer", "frame_count",
+            ],
+            "output_names": [
+                "features", "sample_buffer_out", "sample_len_out",
+                "conv1_buffer_out", "conv2_buffer_out", "frame_count_out",
+            ],
+            "dynamic": ({1: AUTO}, None, None, None, None, None),
+        },
+        {
+            "name": "encoder",
+            "loader": mml.encoder_model_loader,
+            "dummy": mml.encoder_dummy_inputs,
+            "input_names": ["features"],
+            "output_names": ["encoded"],
+            "dynamic": ({1: AUTO},),
+        },
+        {
+            "name": "adapter",
+            "loader": mml.adapter_model_loader,
+            "dummy": mml.adapter_dummy_inputs,
+            "input_names": ["encoded", "pos_offset"],
+            "output_names": ["memory"],
+            "dynamic": ({1: AUTO}, None),
+        },
+        {
+            "name": "cross_kv",
+            "loader": mml.cross_kv_model_loader,
+            "dummy": mml.cross_kv_dummy_inputs,
+            "input_names": ["memory"],
+            "output_names": ["k_cross", "v_cross"],
+            "dynamic": ({1: AUTO},),
+        },
+        {
+            "name": "decoder_kv",
+            "loader": mml.decoder_kv_model_loader,
+            "dummy": mml.decoder_kv_dummy_inputs,
+            "input_names": ["token", "k_self", "v_self", "out_k_cross", "out_v_cross"],
+            "output_names": [
+                "logits", "out_k_self", "out_v_self", "out_k_cross", "out_v_cross",
+            ],
+            "dynamic": ({1: AUTO}, {3: AUTO}, {3: AUTO}, {3: AUTO}, {3: AUTO}),
+        },
+    ]
+
+
+# --------------------------------------------------------------------------- #
+# ONNX I/O renaming (make names match the genai contract exactly)             #
+# --------------------------------------------------------------------------- #
+def rename_io(path, input_names, output_names):
+    model = onnx.load(path)
+    graph = model.graph
+
+    def remap(value_infos, desired):
+        mapping = {}
+        for vi, new in zip(value_infos, desired):
+            if vi.name != new:
+                mapping[vi.name] = new
+                vi.name = new
+        return mapping
+
+    in_map = remap(graph.input, input_names)
+    out_map = remap(graph.output, output_names)
+    rename = {**in_map, **out_map}
+    if rename:
+        for node in graph.node:
+            node.input[:] = [rename.get(n, n) for n in node.input]
+            node.output[:] = [rename.get(n, n) for n in node.output]
+    onnx.save(model, path)
+    return [i.name for i in graph.input], [o.name for o in graph.output]
+
+
+# --------------------------------------------------------------------------- #
+# Main                                                                        #
+# --------------------------------------------------------------------------- #
+def export_component(spec, model_name, output_dir, opset):
+    name = spec["name"]
+    print(f"\n=== exporting {name} ===")
+    module = spec["loader"](model_name).eval()
+    dummy = spec["dummy"](module)
+    # dummy_inputs_func returns a dict keyed by forward-parameter name (Olive's
+    # kwargs export path); the standalone exporter uses positional args, so
+    # flatten to a value tuple in the declared input order.
+    dummy_args = tuple(dummy.values()) if isinstance(dummy, dict) else tuple(dummy)
+    onnx_path = os.path.join(output_dir, f"{name}.onnx")
+
+    torch.onnx.export(
+        module,
+        dummy_args,
+        onnx_path,
+        input_names=spec["input_names"],
+        output_names=spec["output_names"],
+        dynamic_shapes=spec["dynamic"],
+        opset_version=opset,
+        dynamo=True,
+    )
+    ins, outs = rename_io(onnx_path, spec["input_names"], spec["output_names"])
+    onnx.checker.check_model(onnx_path)
+    print(f"  inputs : {ins}")
+    print(f"  outputs: {outs}")
+    return onnx_path
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model", default="usefulsensors/moonshine-streaming-tiny")
+    parser.add_argument("--output-dir", required=True)
+    parser.add_argument("--opset", type=int, default=20)
+    parser.add_argument("--chunk-samples", type=int, default=8000)
+    parser.add_argument(
+        "--only", nargs="*", default=None,
+        help="Optional subset of component names to export.",
+    )
+    args = parser.parse_args()
+
+    os.makedirs(args.output_dir, exist_ok=True)
+    mml.set_chunk_samples(args.chunk_samples)
+
+    specs = build_specs()
+    if args.only:
+        specs = [s for s in specs if s["name"] in args.only]
+
+    for spec in specs:
+        export_component(spec, args.model, args.output_dir, args.opset)
+
+    print(f"\nDone. Graphs written to {args.output_dir}")
+
+
+if __name__ == "__main__":
+    main()

usefulSensors-moonshine-streaming/cpu/info.yaml

@@ -0,0 +1,28 @@
+name: usefulSensors-moonshine-streaming
+provider: usefulsensors
+model_id: usefulsensors/moonshine-streaming-tiny
+# The recipe also supports usefulsensors/moonshine-streaming-small via
+# `--model-name usefulsensors/moonshine-streaming-small`.
+task: automatic-speech-recognition
+framework: ONNX Runtime
+execution_provider: CPUExecutionProvider
+summary: >
+  CPU recipe for exporting the Moonshine streaming ASR model to ONNX (FP32 or
+  INT8 k-quant) as five ONNX Runtime GenAI components (frontend, encoder,
+  adapter, cross-KV, decoder-KV) with the exact IO names the streaming runner
+  expects, plus the generated genai/streaming configs, tokenizer, and Silero
+  VAD.
+
+artifacts:
+  - cpu/moonshine_frontend_fp32_cpu.json
+  - cpu/moonshine_encoder_fp32_cpu.json
+  - cpu/moonshine_adapter_fp32_cpu.json
+  - cpu/moonshine_cross_kv_fp32_cpu.json
+  - cpu/moonshine_decoder_kv_fp32_cpu.json
+  - cpu/moonshine_encoder_kquant8_cpu.json
+  - cpu/moonshine_decoder_kv_kquant8_cpu.json
+
+scripts:
+  - cpu/optimize.py
+  - cpu/moonshine_model_load.py
+  - cpu/export_moonshine_streaming.py

usefulSensors-moonshine-streaming/cpu/moonshine_adapter_fp32_cpu.json

@@ -0,0 +1,41 @@
+{
+    "input_model": {
+        "type": "PyTorchModel",
+        "model_path": "usefulsensors/moonshine-streaming-small",
+        "model_loader": "adapter_model_loader",
+        "model_script": "cpu/moonshine_model_load.py",
+        "io_config": {
+            "input_names": ["encoded", "pos_offset"],
+            "output_names": ["memory"],
+            "dynamic_shapes": {
+                "encoded": {"1": "adp_seq"},
+                "pos_offset": {}
+            }
+        },
+        "dummy_inputs_func": "adapter_dummy_inputs"
+    },
+    "systems": {
+        "local_system": {
+            "type": "LocalSystem",
+            "accelerators": [
+                {
+                    "device": "cpu",
+                    "execution_providers": ["CPUExecutionProvider"]
+                }
+            ]
+        }
+    },
+    "passes": {
+        "convert": {
+            "type": "OnnxConversion",
+            "target_opset": 20,
+            "dynamic": true,
+            "use_dynamo_exporter": true,
+            "save_as_external_data": true,
+            "external_data_name": "adapter.onnx.data"
+        }
+    },
+    "target": "local_system",
+    "output_dir": "build/onnx/adapter.onnx",
+    "no_artifacts": true
+}