Add files using upload-large-folder tool

.gitattributes

@@ -33,3 +33,7 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+encoder_int8.onnx_data filter=lfs diff=lfs merge=lfs -text
+encoder.onnx_data filter=lfs diff=lfs merge=lfs -text
+editor.onnx_data filter=lfs diff=lfs merge=lfs -text
+editor_int8.onnx_data filter=lfs diff=lfs merge=lfs -text

LICENSE

@@ -0,0 +1,202 @@
+
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

editor.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:524863dc275c529cd198d57bf7c309c897ec93db27c86a5d51ea8dffede8b0cf
+size 1838936

editor.onnx_data

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6882668bc9eef72941824bfbabadb0d42d693398a942774e53f2553e38ad83f1
+size 6527000576

editor_int8.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3a1c6cc4f5fb9f0f926bfeb7b0a8f26272b5c5c781bf34648bcfdcfecff79d6a
+size 6418060

editor_int8.onnx_data

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c84ab94b8b79b834f55877d7c794dfb3f4a710e31770e517fb628b1d36b43d03
+size 1632247808

encoder.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bb87c12e45340b614d619cfa3549c10ead32d1f56b69bf18dbbda144d7704d32
+size 1068038

encoder.onnx_data

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:67fe37eda473f984e6a1c574f0265765c196e1cdf39f3750d026608ffc0e6d4a
+size 2491061620

encoder_int8.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5fe3e39656c436f7c5ff2f4f2e11fbc6b55aed007cba37de1fb3b81c0d946671
+size 3310811

encoder_int8.onnx_data

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:64764e9911a47cf724df22f7857e701624e3f81279973e96a20acce245d1086d
+size 935864692

export_nar_editor.py

@@ -0,0 +1,639 @@
+# Copyright 2026 Sam McLeod
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Export the NAR editor (Granite 4.0 1B LLM run bidirectionally) as a single
+ONNX graph and verify parity against the captured PyTorch baseline.
+
+Wraps `model.llm.model` + `model.llm.lm_head` in an nn.Module whose forward
+takes (inputs_embeds, position_ids, attention_mask) where attention_mask is a
+4-D additive mask (zeros = attention allowed everywhere). Exports with
+torch.onnx.export at opset 20, IR 10, single sidecar.
+
+End-to-end parity test:
+  1. Run the already-exported encoder.onnx on the reference clip.
+  2. Run CTC greedy decode + slot insertion + flat embedding assembly via the
+     upstream model's bound methods (matches what Rust glue will do).
+  3. Run the resulting `inputs_embeds` through both the patched PyTorch editor
+     and the exported editor.onnx, then compare logits.
+  4. Decode the ONNX logits at the text positions and verify the transcript
+     matches the upstream NAR transcript exactly.
+
+Usage:
+    HF_HOME=$TMPDIR/hf_home HF_MODULES_CACHE=$TMPDIR/hf_modules \
+        uv run python src/export_nar_editor.py
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import time
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import soundfile as sf
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+# Resolve roots so the script works whether it lives at <repo>/src/<name>.py
+# (project layout) or <bundle>/<name>.py (HF bundle layout). Defaults exist for
+# the project layout; bundle users should pass explicit --audio / --baseline /
+# --model-dir / --out-dir.
+SCRIPT_DIR = Path(__file__).resolve().parent
+REPO_ROOT = SCRIPT_DIR.parent if SCRIPT_DIR.name == "src" else SCRIPT_DIR
+DEFAULT_AUDIO = REPO_ROOT / "test_data" / "10226_10111_000000.wav"
+DEFAULT_BASELINE = REPO_ROOT / "test_data" / "baselines" / "nar.json"
+DEFAULT_MODEL_DIR = REPO_ROOT / "models" / "granite-speech-4.1-2b-nar"
+DEFAULT_OUT_DIR = REPO_ROOT / "exports" / "granite-speech-4.1-2b-nar"
+
+
+def load_audio(path: Path) -> np.ndarray:
+    waveform, sr = sf.read(str(path), dtype="float32")
+    if waveform.ndim > 1:
+        waveform = waveform.mean(axis=1)
+    assert sr == 16000, f"expected 16 kHz, got {sr}"
+    return waveform
+
+
+def tensor_stats(t: torch.Tensor | np.ndarray | None) -> dict[str, Any] | None:
+    if t is None:
+        return None
+    if isinstance(t, torch.Tensor):
+        x = t.detach().float().cpu().numpy()
+        dtype_str = str(t.dtype).replace("torch.", "")
+    else:
+        x = np.asarray(t).astype(np.float32, copy=False)
+        dtype_str = str(t.dtype)
+    flat = x.flatten()
+    return {
+        "shape": list(x.shape),
+        "dtype": dtype_str,
+        "mean": float(flat.mean()) if flat.size else None,
+        "std": float(flat.std()) if flat.size else None,
+        "min": float(flat.min()) if flat.size else None,
+        "max": float(flat.max()) if flat.size else None,
+        "first10": [float(v) for v in flat[:10]],
+    }
+
+
+# ---------------------------------------------------------------------------
+# Wrapper module: LLM backbone + lm_head exposed as a single graph.
+# ---------------------------------------------------------------------------
+
+
+class NAREditor(nn.Module):
+    """Wrap llm.model + llm.lm_head with a 4-D additive attention-mask input.
+
+    Inputs:
+        inputs_embeds:  float32 [1, N_total, D_llm] pre-built flat sequence
+        position_ids:   int64   [1, N_total] cumulative per-sample positions
+        attention_mask: float32 [1, 1, N_total, N_total] additive mask
+                        Zeros = attention allowed everywhere (bidirectional).
+
+    Output:
+        logits:         float32 [1, N_total, V_llm]
+
+    The Granite model in transformers 5.8 calls `create_causal_mask`, which
+    in turn calls `_preprocess_mask_arguments`. That helper short-circuits and
+    returns any 4-D attention mask as-is. So feeding zeros disables causality.
+    """
+
+    def __init__(self, llm_model: nn.Module, lm_head: nn.Module) -> None:
+        super().__init__()
+        self.llm_model = llm_model
+        self.lm_head = lm_head
+
+    def forward(
+        self,
+        inputs_embeds: torch.Tensor,
+        position_ids: torch.Tensor,
+        attention_mask: torch.Tensor,
+    ) -> torch.Tensor:
+        out = self.llm_model(
+            inputs_embeds=inputs_embeds,
+            position_ids=position_ids,
+            attention_mask=attention_mask,
+            use_cache=False,
+        )
+        return self.lm_head(out.last_hidden_state)
+
+
+# ---------------------------------------------------------------------------
+# Model loading (mirrors capture_baselines.py::capture_nar).
+# ---------------------------------------------------------------------------
+
+
+def load_nar_model(model_dir: Path) -> tuple[nn.Module, Any]:
+    """Load NAR model with the same patches capture_baselines.py uses."""
+    from transformers import AutoConfig, AutoFeatureExtractor, AutoModel
+
+    granite_local = REPO_ROOT / "models" / "granite-4.0-1b-base"
+    if not granite_local.exists():
+        raise FileNotFoundError(
+            f"Expected local Granite 4.0 base at {granite_local}; "
+            "run `hf download ibm-granite/granite-4.0-1b-base "
+            "--include '*.json' --include 'tokenizer*' --include '*.txt' "
+            f"--local-dir {granite_local}` first."
+        )
+
+    config = AutoConfig.from_pretrained(str(model_dir), trust_remote_code=True)
+    config.llm_name = str(granite_local)
+    config.attn_implementation = "eager"
+    config._attn_implementation = "eager"
+    for sub_attr in ("llm_config", "encoder_config", "projector_config"):
+        sub = getattr(config, sub_attr, None)
+        if sub is not None:
+            for attr in ("attn_implementation", "_attn_implementation"):
+                try:
+                    setattr(sub, attr, "eager")
+                except Exception:
+                    pass
+
+    print(f"  loading model from {model_dir} (eager)")
+    t0 = time.time()
+    model = AutoModel.from_pretrained(
+        str(model_dir),
+        trust_remote_code=True,
+        torch_dtype=torch.float32,
+        attn_implementation="eager",
+        config=config,
+    )
+    model.eval()
+    # The NAR config nests `llm_config.dtype = "bfloat16"`, which overrides
+    # the top-level `torch_dtype=float32` request for the LLM submodule.
+    # Force the whole model (including LLM and lm_head) to fp32.
+    model = model.to(torch.float32)
+    print(f"  loaded in {time.time() - t0:.1f}s")
+
+    fe = AutoFeatureExtractor.from_pretrained(str(model_dir), trust_remote_code=True)
+    return model, fe
+
+
+# ---------------------------------------------------------------------------
+# Build inputs_embeds for parity from the already-exported encoder.onnx.
+# ---------------------------------------------------------------------------
+
+
+def build_editor_inputs(
+    model: nn.Module,
+    fe: Any,
+    waveform: np.ndarray,
+    encoder_onnx_path: Path,
+) -> tuple[torch.Tensor, torch.Tensor, list[int], list[int], list[str]]:
+    """Mirror what Rust glue does: encoder.onnx -> CTC greedy + slot inserts ->
+    flat (audio, text_with_slots) sequence.
+
+    Returns:
+        flat_embeds:        [1, N_total, 2048]
+        flat_position_ids:  [1, N_total]
+        projected_lengths:  per-sample audio-length list
+        text_lengths:       per-sample text-length list
+        text_ctc_preds:     per-sample CTC strings (for debug)
+    """
+    import onnxruntime as ort
+
+    waveform_t = torch.from_numpy(waveform.copy())
+    inputs = fe([waveform_t], device="cpu")
+    input_features = inputs["input_features"].to(torch.float32)
+    attention_mask_int = inputs["attention_mask"].to(torch.int64)
+
+    print(f"  running encoder.onnx for inputs_embeds construction")
+    sess = ort.InferenceSession(str(encoder_onnx_path), providers=["CPUExecutionProvider"])
+    out_names = [o.name for o in sess.get_outputs()]
+    ort_outputs = sess.run(
+        out_names,
+        {
+            "input_features": input_features.numpy().astype(np.float32),
+            "attention_mask": attention_mask_int.numpy().astype(np.int64),
+        },
+    )
+    out_map = dict(zip(out_names, ort_outputs))
+    bpe_dense = torch.from_numpy(out_map["bpe_logits_dense"])  # [B, T_bpe, V_bpe]
+    bpe_mask = torch.from_numpy(out_map["bpe_mask"])  # [B, T_bpe] bool
+    audio_embeds = torch.from_numpy(out_map["audio_embeds"])  # [B, T_audio, 2048]
+    audio_lengths = torch.from_numpy(out_map["audio_lengths"])  # [B] int64
+
+    # Reconstruct sparse BPE logits ([N_valid, V_bpe]) and per-sample lengths.
+    bpe_lengths = bpe_mask.sum(dim=1).tolist()
+    bpe_logits_flat = bpe_dense[bpe_mask]  # [N_valid, V_bpe]
+
+    # CTC greedy decode -> List[str] (one per sample).
+    text_ctc_preds = model._decode_bpe_ctc_greedy(bpe_logits_flat, bpe_lengths)
+    print(f"  text_ctc_preds: {text_ctc_preds!r}")
+
+    # Build flat LLM inputs. The upstream `_build_flat_llm_inputs` calls
+    # `self.projector(encoder_embs)` to produce audio embeddings; we already
+    # have those from the encoder.onnx. To avoid re-running the projector,
+    # we replicate the slot-insertion + concat steps inline here.
+    # This must match the upstream method byte-for-byte modulo the projector
+    # source (encoder.onnx vs PyTorch projector).
+    if model.config.scale_projected_embeddings and hasattr(model.llm.config, "embedding_multiplier"):
+        audio_embeds_scaled = audio_embeds / model.llm.config.embedding_multiplier
+    else:
+        audio_embeds_scaled = audio_embeds
+    audio_embeds_scaled = audio_embeds_scaled.to(model.llm.model.embed_tokens.weight.dtype)
+
+    pred_text_llm_tokens = model.llm_tokenizer(text_ctc_preds)
+    text_ids_with_slots = [
+        model.add_insertion_slots(torch.tensor(x))
+        for x in pred_text_llm_tokens.input_ids
+    ]
+
+    embed_tokens = model.llm.model.embed_tokens
+    embeds_list = []
+    position_ids_list = []
+    text_lengths = []
+    projected_lengths = audio_lengths.tolist()
+
+    for i, audio_len in enumerate(projected_lengths):
+        audio = audio_embeds_scaled[i, :audio_len]
+        text_emb = embed_tokens(text_ids_with_slots[i])
+        sample = torch.cat([audio, text_emb], dim=0)
+        embeds_list.append(sample)
+        position_ids_list.append(torch.arange(sample.shape[0]))
+        text_lengths.append(text_ids_with_slots[i].shape[0])
+
+    flat_embeds = torch.cat(embeds_list, dim=0).unsqueeze(0).to(torch.float32)
+    flat_position_ids = torch.cat(position_ids_list, dim=0).unsqueeze(0).to(torch.int64)
+
+    print(
+        f"  flat_embeds={tuple(flat_embeds.shape)}  "
+        f"flat_position_ids={tuple(flat_position_ids.shape)}  "
+        f"projected_lengths={projected_lengths}  text_lengths={text_lengths}"
+    )
+    return flat_embeds, flat_position_ids, projected_lengths, text_lengths, text_ctc_preds
+
+
+# ---------------------------------------------------------------------------
+# Export.
+# ---------------------------------------------------------------------------
+
+
+def export_onnx(
+    wrapper: NAREditor,
+    sample_inputs_embeds: torch.Tensor,
+    sample_position_ids: torch.Tensor,
+    sample_attention_mask: torch.Tensor,
+    out_path: Path,
+    opset: int = 20,
+    ir_version: int = 10,
+) -> None:
+    import tempfile
+
+    import onnx
+
+    out_path.parent.mkdir(parents=True, exist_ok=True)
+    print(f"  exporting to {out_path} (opset={opset}, ir_version={ir_version})")
+
+    dynamic_axes = {
+        "inputs_embeds": {1: "N"},
+        "position_ids": {1: "N"},
+        "attention_mask": {2: "N", 3: "N"},
+        "logits": {1: "N"},
+    }
+
+    with tempfile.TemporaryDirectory(prefix="nar_editor_onnx_") as scratch_dir:
+        scratch_path = Path(scratch_dir) / "editor.onnx"
+        t0 = time.time()
+        torch.onnx.export(
+            wrapper,
+            (sample_inputs_embeds, sample_position_ids, sample_attention_mask),
+            str(scratch_path),
+            input_names=["inputs_embeds", "position_ids", "attention_mask"],
+            output_names=["logits"],
+            dynamic_axes=dynamic_axes,
+            opset_version=opset,
+            do_constant_folding=True,
+            export_params=True,
+            dynamo=False,
+        )
+        print(f"  stage-1 torch.onnx.export done in {time.time() - t0:.1f}s")
+
+        print("  stage-2: re-saving with single .onnx_data sidecar + ir bump")
+        model_proto = onnx.load(str(scratch_path), load_external_data=True)
+        if model_proto.ir_version < ir_version:
+            model_proto.ir_version = ir_version
+
+        for tensor in model_proto.graph.initializer:
+            tensor.ClearField("data_location")
+            tensor.ClearField("external_data")
+
+        sidecar_name = out_path.name + "_data"
+        if (out_path.parent / sidecar_name).exists():
+            (out_path.parent / sidecar_name).unlink()
+        if out_path.exists():
+            out_path.unlink()
+
+        onnx.save_model(
+            model_proto,
+            str(out_path),
+            save_as_external_data=True,
+            all_tensors_to_one_file=True,
+            location=sidecar_name,
+            size_threshold=1024,
+            convert_attribute=False,
+        )
+        onnx.checker.check_model(str(out_path), full_check=False)
+
+        # Quick op-domain audit: the doc target says no `com.microsoft` ops.
+        domains = sorted({n.domain for n in model_proto.graph.node})
+        print(f"  saved {out_path} (+ {sidecar_name})  node-domains={domains}")
+
+
+# ---------------------------------------------------------------------------
+# Parity test.
+# ---------------------------------------------------------------------------
+
+
+def run_parity(
+    model: nn.Module,
+    flat_embeds: torch.Tensor,
+    flat_position_ids: torch.Tensor,
+    projected_lengths: list[int],
+    text_lengths: list[int],
+    text_ctc_preds: list[str],
+    onnx_path: Path,
+    parity_json: Path,
+    baseline_json: Path,
+    abs_tol: float = 1e-3,
+    argmax_only: bool = False,
+) -> bool:
+    import onnxruntime as ort
+
+    print("\n=== parity check ===")
+    N = flat_embeds.shape[1]
+    attention_mask = torch.zeros(1, 1, N, N, dtype=flat_embeds.dtype)
+
+    print("  running PyTorch editor (llm.model + lm_head with 4-D zeros mask)")
+    t0 = time.time()
+    with torch.inference_mode():
+        out = model.llm.model(
+            inputs_embeds=flat_embeds,
+            position_ids=flat_position_ids,
+            attention_mask=attention_mask,
+            use_cache=False,
+        )
+        logits_pt = model.llm.lm_head(out.last_hidden_state)
+    print(f"    pytorch forward: {time.time() - t0:.2f}s")
+    print(f"    logits_pt shape={tuple(logits_pt.shape)}")
+
+    print(f"  running ONNX inference: {onnx_path}")
+    sess = ort.InferenceSession(str(onnx_path), providers=["CPUExecutionProvider"])
+    t0 = time.time()
+    ort_outputs = sess.run(
+        ["logits"],
+        {
+            "inputs_embeds": flat_embeds.detach().numpy().astype(np.float32),
+            "position_ids": flat_position_ids.detach().numpy().astype(np.int64),
+            "attention_mask": attention_mask.detach().numpy().astype(np.float32),
+        },
+    )
+    print(f"    onnx forward: {time.time() - t0:.2f}s")
+    logits_ort = ort_outputs[0]
+    print(f"    logits_ort shape={logits_ort.shape}")
+
+    pt_np = logits_pt.detach().float().cpu().numpy()
+    abs_err = np.abs(pt_np - logits_ort)
+    max_err = float(abs_err.max())
+    mean_err = float(abs_err.mean())
+    p99 = float(np.percentile(abs_err, 99))
+
+    am_pt = pt_np.argmax(-1)
+    am_ort = logits_ort.argmax(-1)
+    argmax_mismatches = int((am_pt != am_ort).sum())
+    argmax_total = int(am_pt.size)
+
+    # Per-segment argmax: only the text positions feed the transcript decode
+    # (audio positions are read-only inputs to attention). Track separately so
+    # INT8 mode can ship on transcript correctness even when audio-position
+    # logits drift under weight quantisation.
+    text_argmax_mismatches = 0
+    text_argmax_total = 0
+    audio_argmax_mismatches = 0
+    audio_argmax_total = 0
+    seg_offset = 0
+    for i in range(len(projected_lengths)):
+        a_lo, a_hi = seg_offset, seg_offset + projected_lengths[i]
+        t_lo, t_hi = a_hi, a_hi + text_lengths[i]
+        seg_offset = t_hi
+        audio_argmax_mismatches += int((am_pt[0, a_lo:a_hi] != am_ort[0, a_lo:a_hi]).sum())
+        audio_argmax_total += a_hi - a_lo
+        text_argmax_mismatches += int((am_pt[0, t_lo:t_hi] != am_ort[0, t_lo:t_hi]).sum())
+        text_argmax_total += t_hi - t_lo
+
+    # Top-5 stability check.
+    topk_pt = np.argsort(-pt_np, axis=-1)[..., :5]
+    topk_ort = np.argsort(-logits_ort, axis=-1)[..., :5]
+    top1_match = int((topk_pt[..., 0] == topk_ort[..., 0]).sum())
+    top5_set_match = int(
+        (
+            np.sort(topk_pt, axis=-1) == np.sort(topk_ort, axis=-1)
+        ).all(axis=-1).sum()
+    )
+
+    print("\n--- logits diff ---")
+    print(f"  shape pt={pt_np.shape}  ort={logits_ort.shape}")
+    print(f"  max_abs_err={max_err:.3e}  mean_abs_err={mean_err:.3e}  p99={p99:.3e}")
+    print(f"  argmax mismatches: {argmax_mismatches}/{argmax_total}")
+    print(
+        f"  text-segment argmax mismatches: {text_argmax_mismatches}/{text_argmax_total}; "
+        f"audio-segment argmax mismatches: {audio_argmax_mismatches}/{audio_argmax_total}"
+    )
+    print(f"  top1 match: {top1_match}/{argmax_total}  top5-set match: {top5_set_match}/{argmax_total}")
+
+    # End-to-end transcript check: slice text positions from ONNX logits and run
+    # the upstream argmax + unique_consecutive + EOS removal.
+    eos_id = int(model.llm.config.eos_token_id)
+    offset = 0
+    decoded_segments = []
+    for i in range(len(projected_lengths)):
+        offset += projected_lengths[i]
+        seg = logits_ort[0, offset:offset + text_lengths[i]]
+        offset += text_lengths[i]
+        pred = seg.argmax(-1)
+        # Use torch.unique_consecutive to mirror upstream exactly.
+        collapsed = torch.unique_consecutive(torch.from_numpy(pred)).tolist()
+        collapsed = [t for t in collapsed if t != eos_id]
+        text = model.llm_tokenizer.decode(collapsed, skip_special_tokens=True)
+        decoded_segments.append(text)
+
+    onnx_transcript = decoded_segments[0] if decoded_segments else ""
+    print(f"\n  ONNX transcript: {onnx_transcript!r}")
+
+    baseline_transcript = None
+    if baseline_json.exists():
+        baseline = json.loads(baseline_json.read_text())
+        baseline_transcript = baseline.get("transcript")
+        print(f"  baseline transcript: {baseline_transcript!r}")
+    transcript_match = bool(
+        baseline_transcript is not None and onnx_transcript == baseline_transcript
+    )
+
+    # Pass criteria: zero argmax mismatches AND transcript matches baseline.
+    # The max-abs threshold is informational; spec mandates argmax stability.
+    # In INT8 mode, only text-position argmax matters - audio positions feed
+    # attention but never get sliced for the transcript decode, so weight-quant
+    # drift there is harmless.
+    max_err_ok = max_err <= abs_tol
+    argmax_ok = argmax_mismatches == 0
+    text_argmax_ok = text_argmax_mismatches == 0
+    if argmax_only:
+        overall_ok = text_argmax_ok and transcript_match
+    else:
+        overall_ok = argmax_ok and transcript_match
+
+    sidecar = onnx_path.with_name(onnx_path.name + "_data")
+    int8_size = int(sidecar.stat().st_size) if sidecar.exists() else None
+    payload = {
+        "ok": overall_ok,
+        "abs_tol": abs_tol,
+        "argmax_only": argmax_only,
+        "graph_path": str(onnx_path),
+        "graph_size_bytes": int(onnx_path.stat().st_size),
+        "int8_size_bytes": int8_size,
+        "shape_pt": list(pt_np.shape),
+        "shape_ort": list(logits_ort.shape),
+        "text_argmax_mismatches": text_argmax_mismatches,
+        "text_argmax_total": text_argmax_total,
+        "audio_argmax_mismatches": audio_argmax_mismatches,
+        "audio_argmax_total": audio_argmax_total,
+        "max_abs_err": max_err,
+        "mean_abs_err": mean_err,
+        "p99_abs_err": p99,
+        "max_abs_err_ok": max_err_ok,
+        "argmax_mismatches": argmax_mismatches,
+        "argmax_total": argmax_total,
+        "argmax_ok": argmax_ok,
+        "top1_match": top1_match,
+        "top5_set_match": top5_set_match,
+        "logits_stats_pt": tensor_stats(logits_pt),
+        "logits_stats_ort": tensor_stats(logits_ort),
+        "projected_lengths": projected_lengths,
+        "text_lengths": text_lengths,
+        "text_ctc_preds": text_ctc_preds,
+        "onnx_transcript": onnx_transcript,
+        "baseline_transcript": baseline_transcript,
+        "transcript_match": transcript_match,
+    }
+    parity_json.parent.mkdir(parents=True, exist_ok=True)
+    parity_json.write_text(json.dumps(payload, indent=2))
+    print(f"\n  wrote parity report -> {parity_json}")
+
+    print("\n--- parity summary ---")
+    print(f"  max_abs_err <= {abs_tol}: {'PASS' if max_err_ok else 'FAIL'}  ({max_err:.3e})")
+    print(f"  argmax mismatches == 0: {'PASS' if argmax_ok else 'FAIL'}  ({argmax_mismatches}/{argmax_total})")
+    print(f"  transcript matches baseline: {'PASS' if transcript_match else 'FAIL'}")
+    print(f"\n{'PASS' if overall_ok else 'FAIL'}")
+    return overall_ok
+
+
+# ---------------------------------------------------------------------------
+# Main.
+# ---------------------------------------------------------------------------
+
+
+def main() -> None:
+    p = argparse.ArgumentParser()
+    p.add_argument("--audio", default=str(DEFAULT_AUDIO))
+    p.add_argument("--baseline", default=str(DEFAULT_BASELINE))
+    p.add_argument("--model-dir", default=str(DEFAULT_MODEL_DIR))
+    p.add_argument("--out-dir", default=str(DEFAULT_OUT_DIR))
+    p.add_argument("--abs-tol", type=float, default=1e-3)
+    p.add_argument("--skip-export", action="store_true", help="skip the export step (re-run parity only)")
+    p.add_argument(
+        "--graph-suffix",
+        default="",
+        help="suffix appended to the editor graph stem (e.g. '_int8') so parity runs "
+        "against editor<suffix>.onnx. Parity output goes to editor_parity<suffix>.json. "
+        "When set, --skip-export is implied.",
+    )
+    p.add_argument(
+        "--encoder-suffix",
+        default=None,
+        help="suffix for the encoder graph used to build editor inputs. Defaults to "
+        "--graph-suffix; pass '' to force the FP32 encoder.",
+    )
+    args = p.parse_args()
+
+    out_dir = Path(args.out_dir)
+    suffix = args.graph_suffix
+    if suffix and not args.skip_export:
+        print(f"  --graph-suffix={suffix!r} set; implying --skip-export")
+        args.skip_export = True
+    encoder_suffix = args.encoder_suffix if args.encoder_suffix is not None else suffix
+    onnx_path = out_dir / f"editor{suffix}.onnx"
+    parity_json = out_dir / f"editor_parity{suffix}.json"
+    encoder_onnx = out_dir / f"encoder{encoder_suffix}.onnx"
+
+    if not encoder_onnx.exists():
+        raise FileNotFoundError(
+            f"Expected exported encoder at {encoder_onnx}; "
+            "run src/export_nar_encoder.py first."
+        )
+
+    model_dir = Path(args.model_dir)
+
+    print(f"audio:    {args.audio}")
+    print(f"out_dir:  {out_dir}")
+    waveform = load_audio(Path(args.audio))
+    print(f"  duration={waveform.shape[0] / 16000:.2f}s")
+
+    print("loading model...")
+    model, fe = load_nar_model(model_dir)
+
+    # Build editor inputs from the encoder.onnx output and the model's bound
+    # methods. This is what the Rust glue will do in production.
+    flat_embeds, flat_position_ids, projected_lengths, text_lengths, text_ctc_preds = (
+        build_editor_inputs(model, fe, waveform, encoder_onnx)
+    )
+
+    wrapper = NAREditor(llm_model=model.llm.model, lm_head=model.llm.lm_head)
+    wrapper.eval()
+
+    if not args.skip_export:
+        N = flat_embeds.shape[1]
+        sample_attn = torch.zeros(1, 1, N, N, dtype=flat_embeds.dtype)
+        with torch.inference_mode():
+            export_onnx(
+                wrapper=wrapper,
+                sample_inputs_embeds=flat_embeds,
+                sample_position_ids=flat_position_ids,
+                sample_attention_mask=sample_attn,
+                out_path=onnx_path,
+                opset=20,
+                ir_version=10,
+            )
+
+    ok = run_parity(
+        model=model,
+        flat_embeds=flat_embeds,
+        flat_position_ids=flat_position_ids,
+        projected_lengths=projected_lengths,
+        text_lengths=text_lengths,
+        text_ctc_preds=text_ctc_preds,
+        onnx_path=onnx_path,
+        parity_json=parity_json,
+        baseline_json=Path(args.baseline),
+        abs_tol=args.abs_tol,
+        argmax_only=bool(suffix),
+    )
+    if not ok:
+        raise SystemExit(1)
+
+
+if __name__ == "__main__":
+    main()

export_nar_encoder.py

@@ -0,0 +1,743 @@
+# Copyright 2026 Sam McLeod
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Export the NAR encoder + projector as a single ONNX graph and verify parity.
+
+Wraps `model.encoder` (NLECTCEncoder) and `model.projector` (EncoderProjectorQFormer)
+into one nn.Module whose forward takes (input_features, attention_mask) and returns
+char_logits, dense BPE logits, BPE-mask, audio embeddings, and audio_lengths. Exports
+with torch.onnx.export (TorchScript-style; opset 20, IR 10) using external-data
+storage. Then runs the ONNX graph via onnxruntime CPU on the reference clip and
+diffs against the live PyTorch forward.
+
+Usage:
+    HF_HOME=$TMPDIR/hf_home HF_MODULES_CACHE=$TMPDIR/hf_modules \
+        uv run python src/export_nar_encoder.py
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import time
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import soundfile as sf
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+# Resolve roots so the script works whether it lives at <repo>/src/<name>.py
+# (project layout) or <bundle>/<name>.py (HF bundle layout). Defaults exist for
+# the project layout; bundle users should pass explicit --audio / --baseline /
+# --model-dir / --out-dir.
+SCRIPT_DIR = Path(__file__).resolve().parent
+REPO_ROOT = SCRIPT_DIR.parent if SCRIPT_DIR.name == "src" else SCRIPT_DIR
+DEFAULT_AUDIO = REPO_ROOT / "test_data" / "10226_10111_000000.wav"
+DEFAULT_BASELINE = REPO_ROOT / "test_data" / "baselines" / "nar.json"
+DEFAULT_MODEL_DIR = REPO_ROOT / "models" / "granite-speech-4.1-2b-nar"
+DEFAULT_OUT_DIR = REPO_ROOT / "exports" / "granite-speech-4.1-2b-nar"
+
+
+def load_audio(path: Path) -> np.ndarray:
+    waveform, sr = sf.read(str(path), dtype="float32")
+    if waveform.ndim > 1:
+        waveform = waveform.mean(axis=1)
+    assert sr == 16000, f"expected 16 kHz, got {sr}"
+    return waveform
+
+
+def tensor_stats(t: torch.Tensor | np.ndarray | None) -> dict[str, Any] | None:
+    if t is None:
+        return None
+    if isinstance(t, torch.Tensor):
+        x = t.detach().float().cpu().numpy()
+        dtype_str = str(t.dtype).replace("torch.", "")
+    else:
+        x = np.asarray(t).astype(np.float32, copy=False)
+        dtype_str = str(t.dtype)
+    flat = x.flatten()
+    return {
+        "shape": list(x.shape),
+        "dtype": dtype_str,
+        "mean": float(flat.mean()) if flat.size else None,
+        "std": float(flat.std()) if flat.size else None,
+        "min": float(flat.min()) if flat.size else None,
+        "max": float(flat.max()) if flat.size else None,
+        "first10": [float(v) for v in flat[:10]],
+    }
+
+
+# ---------------------------------------------------------------------------
+# Wrapper module: encoder + projector exposed as a single graph.
+# ---------------------------------------------------------------------------
+
+
+class NAREncoderProjector(nn.Module):
+    """Combined NAR encoder + projector wrapper for ONNX export.
+
+    Inputs:
+        input_features: float32 [B, T, 160]
+        attention_mask: int64   [B, T] (1 = valid, 0 = pad)
+
+    Outputs (all dense; BPE mask carries the validity layout):
+        char_logits:       float32 [B, T_enc, 348]
+        bpe_logits_dense:  float32 [B, T_bpe, 100353] where T_bpe = ceil(T_enc / 4)
+        bpe_mask:          bool    [B, T_bpe]
+        audio_embeds:      float32 [B, T_audio, 2048] where T_audio = nblocks * (block_size/downsample_rate)
+        audio_lengths:     int64   [B] = attention_mask.sum(-1) // downsample_rate
+
+    Notes:
+        - The encoder produces a *sparse* BPE tensor [N_valid, V_bpe] in the upstream
+          forward. We re-densify it to [B, T_bpe, V_bpe] and emit the corresponding
+          [B, T_bpe] bool mask. T_bpe is the encoder time dim downsampled by the
+          encoder's bpe_pooling_window (= 4 for this checkpoint), NOT T_enc.
+        - attention_mask must be int64; the wrapper casts it to bool internally.
+    """
+
+    def __init__(self, encoder: nn.Module, projector: nn.Module, encoder_layer_indices: list[int]) -> None:
+        super().__init__()
+        self.encoder = encoder
+        self.projector = projector
+        self.encoder_layer_indices = list(encoder_layer_indices)
+        self.bpe_pool = int(encoder.config.bpe_pooling_window)
+        self.downsample_rate = int(projector.config.downsample_rate)
+
+    def forward(
+        self,
+        input_features: torch.Tensor,
+        attention_mask: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        mask_bool = attention_mask.to(torch.bool)
+
+        enc_out = self.encoder(
+            input_features=input_features,
+            attention_mask=mask_bool,
+            output_hidden_states=True,
+        )
+        char_logits = enc_out.logits  # [B, T_enc, 348]
+
+        # Concatenate the four selected encoder hidden layers along the feature dim.
+        all_h = enc_out.all_hidden_states
+        selected = [all_h[idx] for idx in self.encoder_layer_indices]
+        encoder_embs = torch.cat(selected, dim=-1)  # [B, T_enc, 4 * 1024]
+
+        # Projector: produces [B, T_audio, llm_dim].
+        audio_embeds = self.projector(encoder_embs)
+
+        # Densify BPE logits: encoder gives flat [N_valid, V_bpe] selected by bpe_mask.
+        # We rebuild a [B, T_bpe, V_bpe] tensor using scatter / index_put.
+        # bpe_mask matches the encoder's pooled mask: pad attention_mask to a multiple
+        # of bpe_pool, then take stride bpe_pool, like the upstream code.
+        T = mask_bool.shape[1]
+        pad_T = (-T) % self.bpe_pool  # = (bpe_pool - T % bpe_pool) % bpe_pool
+        bpe_mask = F.pad(mask_bool, (0, pad_T), value=False)[:, :: self.bpe_pool]
+
+        bpe_logits_sparse = enc_out.logits_bpe  # [N_valid, V_bpe]
+        B = mask_bool.shape[0]
+        T_bpe = bpe_mask.shape[1]
+        V_bpe = bpe_logits_sparse.shape[-1]
+        bpe_dense = torch.zeros(
+            B, T_bpe, V_bpe, dtype=bpe_logits_sparse.dtype, device=bpe_logits_sparse.device
+        )
+        # Pure tensor scatter via masked_scatter so the indexing traces cleanly.
+        bpe_dense = bpe_dense.masked_scatter(bpe_mask.unsqueeze(-1), bpe_logits_sparse)
+
+        # audio_lengths in projector resolution.
+        audio_lengths = (attention_mask.to(torch.int64).sum(dim=1) // self.downsample_rate)
+
+        return char_logits, bpe_dense, bpe_mask, audio_embeds, audio_lengths
+
+
+# ---------------------------------------------------------------------------
+# Model loading (mirrors capture_baselines.py).
+# ---------------------------------------------------------------------------
+
+
+def load_nar_model(model_dir: Path) -> tuple[nn.Module, Any]:
+    """Load NAR model with the same patches capture_baselines.py uses."""
+    from transformers import AutoConfig, AutoFeatureExtractor, AutoModel
+
+    granite_local = REPO_ROOT / "models" / "granite-4.0-1b-base"
+    if not granite_local.exists():
+        raise FileNotFoundError(
+            f"Expected local Granite 4.0 base at {granite_local}; "
+            "run `hf download ibm-granite/granite-4.0-1b-base "
+            "--include '*.json' --include 'tokenizer*' --include '*.txt' "
+            f"--local-dir {granite_local}` first."
+        )
+
+    config = AutoConfig.from_pretrained(str(model_dir), trust_remote_code=True)
+    config.llm_name = str(granite_local)
+    config.attn_implementation = "eager"
+    config._attn_implementation = "eager"
+    for sub_attr in ("llm_config", "encoder_config", "projector_config"):
+        sub = getattr(config, sub_attr, None)
+        if sub is not None:
+            for attr in ("attn_implementation", "_attn_implementation"):
+                try:
+                    setattr(sub, attr, "eager")
+                except Exception:
+                    pass
+
+    print(f"  loading model from {model_dir} (eager)")
+    t0 = time.time()
+    model = AutoModel.from_pretrained(
+        str(model_dir),
+        trust_remote_code=True,
+        torch_dtype=torch.float32,
+        attn_implementation="eager",
+        config=config,
+    )
+    model.eval()
+    print(f"  loaded in {time.time() - t0:.1f}s")
+
+    fe = AutoFeatureExtractor.from_pretrained(str(model_dir), trust_remote_code=True)
+    return model, fe
+
+
+# ---------------------------------------------------------------------------
+# Trace-friendly monkey-patches.
+# ---------------------------------------------------------------------------
+
+
+def patch_for_tracing(model: nn.Module) -> None:
+    """Replace forward implementations that rely on Python control flow over tensor
+    shapes with versions that always execute the same op path. Lets torch.onnx.export
+    produce a single graph valid for any T.
+
+    Affected modules:
+      - NLEConformerAttention.forward: replaces SDPA with a plain matmul-based
+        attention so the ONNX exporter doesn't need an SDPA decomposition. Also
+        keeps the conditional pad path on a tensor-shaped pad amount (always
+        executes pad with `(-num_features) % context_size`).
+      - QFormerCrossAttention.forward: same SDPA -> matmul replacement.
+      - EncoderProjectorQFormer.forward: replaces the data-dependent `if rest > 0`
+        branch with an unconditional pad whose length is `(-seq_len) % block_size`
+        (zero when already a multiple) and uses ceil-div for nblocks.
+    """
+    encoder = model.encoder
+    projector = model.projector
+
+    # ---- patch NLEConformerAttention.forward (every layer shares the same class) ----
+    attn0 = encoder.layers[0].attn
+    attn_cls = type(attn0)
+
+    def attn_forward(self, hidden_states: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor:
+        hidden_states = self.pre_norm(hidden_states)
+        bsz, num_features, _ = hidden_states.shape
+        # Always-pad: pad amount may be zero. Use modulo so the graph is valid
+        # for any T at runtime.
+        pad_amount = (-num_features) % self.context_size
+        num_blocks = (num_features + self.context_size - 1) // self.context_size
+        if self.config.old_encoder_mask:
+            attention_mask = torch.ones_like(attention_mask)
+
+        hidden_states = F.pad(hidden_states, (0, 0, 0, pad_amount))
+        attention_mask = F.pad(attention_mask, (0, pad_amount))
+
+        query_states = self.to_q(hidden_states)
+        key_states, value_states = self.to_kv(hidden_states).chunk(2, dim=-1)
+
+        query_states = query_states.reshape(
+            bsz, num_blocks, self.context_size, self.num_heads, -1
+        ).transpose(2, 3)
+        key_states = key_states.reshape(
+            bsz, num_blocks, self.context_size, self.num_heads, -1
+        ).transpose(2, 3)
+        value_states = value_states.reshape(
+            bsz, num_blocks, self.context_size, self.num_heads, -1
+        ).transpose(2, 3)
+
+        seq = torch.arange(self.config.context_size, device=hidden_states.device)
+        dist = seq.view(-1, 1) - seq.view(1, -1) + self.config.max_pos_emb
+        rel_pos_emb = self.rel_pos_emb(dist).to(query_states.dtype)
+        # query_states: [B, M, H, C, D]; rel_pos_emb: [C, R, D]
+        # Output: [B, M, H, C, R]
+        pos_attn = torch.einsum("b m h c d, c r d -> b m h c r", query_states, rel_pos_emb) * self.scale
+        mask_value = -torch.finfo(pos_attn.dtype).max
+        expanded_attention_mask = attention_mask.reshape(bsz, num_blocks, 1, 1, -1)
+        # Avoid in-place masked_fill_ which can confuse the exporter.
+        pos_attn = pos_attn.masked_fill(~expanded_attention_mask, mask_value)
+
+        # Plain matmul attention (matches MATH SDPA backend numerically).
+        # query_states: [B, M, H, C, D]; key_states: [B, M, H, C, D]
+        attn_logits = torch.matmul(query_states, key_states.transpose(-1, -2)) * self.scale
+        attn_logits = attn_logits + pos_attn
+        attn_weights = torch.softmax(attn_logits, dim=-1)
+        out = torch.matmul(attn_weights, value_states)  # [B, M, H, C, D]
+
+        out = out.transpose(2, 3).reshape(bsz, hidden_states.shape[1], -1)
+        out = self.to_out(out[:, :num_features, :])
+        return self.dropout(out)
+
+    attn_cls.forward = attn_forward
+
+    # ---- patch QFormerCrossAttention.forward (replace SDPA with matmul) ----
+    qformer_attn = projector.qformer.layers[0].cross_attention
+    qformer_attn_cls = type(qformer_attn)
+
+    def qformer_attn_forward(self, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor) -> torch.Tensor:
+        batch_size, query_len, _ = hidden_states.shape
+        encoder_len = encoder_hidden_states.shape[1]
+
+        q = (
+            self.q_proj(hidden_states)
+            .view(batch_size, query_len, self.num_heads, self.head_dim)
+            .transpose(1, 2)
+        )
+        k = (
+            self.k_proj(encoder_hidden_states)
+            .view(batch_size, encoder_len, self.num_heads, self.head_dim)
+            .transpose(1, 2)
+        )
+        v = (
+            self.v_proj(encoder_hidden_states)
+            .view(batch_size, encoder_len, self.num_heads, self.head_dim)
+            .transpose(1, 2)
+        )
+        scale = self.head_dim ** -0.5
+        attn_logits = torch.matmul(q, k.transpose(-1, -2)) * scale
+        attn_weights = torch.softmax(attn_logits, dim=-1)
+        attn_output = torch.matmul(attn_weights, v)
+        attn_output = attn_output.transpose(1, 2).contiguous().view(batch_size, query_len, self.hidden_size)
+        return self.o_proj(attn_output)
+
+    qformer_attn_cls.forward = qformer_attn_forward
+
+    # ---- patch EncoderProjectorQFormer.forward (always-pad; no `if rest > 0`) ----
+    projector_cls = type(projector)
+
+    def projector_forward(self, x: torch.Tensor) -> torch.Tensor:
+        batch_size, seq_len, dim = x.size()
+
+        x = x.view(batch_size, seq_len, self.config.num_encoder_layers, self.config.encoder_dim)
+        normalized_layers = []
+        for i, layer_norm in enumerate(self.layer_norms):
+            normalized_layers.append(layer_norm(x[:, :, i]))
+        x = torch.cat(normalized_layers, dim=-1)
+
+        x = self.projector_act(self.layer_projector(x))
+
+        block_size = self.config.block_size
+        # Always pad to next multiple of block_size; pad may be zero.
+        pad_len = (-seq_len) % block_size
+        x = F.pad(x, (0, 0, 0, pad_len), "constant", 0)
+        nblocks = (seq_len + block_size - 1) // block_size
+
+        x = x.view(batch_size * nblocks, block_size, self.config.hidden_size)
+        query_length = self.query.shape[1]
+        mean_pool = x.view(
+            batch_size * nblocks,
+            query_length,
+            self.config.downsample_rate,
+            self.config.hidden_size,
+        ).mean(dim=-2)
+
+        query_output = self.qformer(
+            query_embeds=self.dropout(self.query + mean_pool),
+            encoder_hidden_states=self.dropout(x + self.window_positions),
+        )
+
+        query_output = query_output.view(batch_size, nblocks * query_length, -1)
+        query_output = self.dropout(self.out_norm(query_output))
+        return self.out_linear(query_output)
+
+    projector_cls.forward = projector_forward
+
+
+# ---------------------------------------------------------------------------
+# Export.
+# ---------------------------------------------------------------------------
+
+
+def export_onnx(
+    wrapper: NAREncoderProjector,
+    sample_input_features: torch.Tensor,
+    sample_attention_mask: torch.Tensor,
+    out_path: Path,
+    opset: int = 20,
+    ir_version: int = 10,
+) -> None:
+    import tempfile
+
+    import onnx
+
+    out_path.parent.mkdir(parents=True, exist_ok=True)
+    print(f"  exporting to {out_path} (opset={opset}, ir_version={ir_version})")
+
+    dynamic_axes = {
+        "input_features": {0: "B", 1: "T"},
+        "attention_mask": {0: "B", 1: "T"},
+        "char_logits": {0: "B", 1: "T_enc"},
+        "bpe_logits_dense": {0: "B", 1: "T_bpe"},
+        "bpe_mask": {0: "B", 1: "T_bpe"},
+        "audio_embeds": {0: "B", 1: "T_audio"},
+        "audio_lengths": {0: "B"},
+    }
+
+    # Stage 1: torch.onnx.export to a scratch directory. The legacy TorchScript
+    # exporter spills weights as individual sidecar files; we move them out of
+    # the final target dir so we end up with exactly two artefacts on disk.
+    with tempfile.TemporaryDirectory(prefix="nar_onnx_") as scratch_dir:
+        scratch_path = Path(scratch_dir) / "encoder.onnx"
+        t0 = time.time()
+        torch.onnx.export(
+            wrapper,
+            (sample_input_features, sample_attention_mask),
+            str(scratch_path),
+            input_names=["input_features", "attention_mask"],
+            output_names=[
+                "char_logits",
+                "bpe_logits_dense",
+                "bpe_mask",
+                "audio_embeds",
+                "audio_lengths",
+            ],
+            dynamic_axes=dynamic_axes,
+            opset_version=opset,
+            do_constant_folding=True,
+            export_params=True,
+            dynamo=False,
+        )
+        print(f"  stage-1 torch.onnx.export done in {time.time() - t0:.1f}s")
+
+        # Stage 2: load with external data resolved, bump IR version, then
+        # rewrite all weights into a single sidecar at the final location.
+        print("  stage-2: re-saving with single .onnx_data sidecar + ir bump")
+        model_proto = onnx.load(str(scratch_path), load_external_data=True)
+        if model_proto.ir_version < ir_version:
+            model_proto.ir_version = ir_version
+
+        # Strip any pre-existing external-data references so save_model rewrites
+        # them cleanly into the new sidecar.
+        for tensor in model_proto.graph.initializer:
+            tensor.ClearField("data_location")
+            tensor.ClearField("external_data")
+
+        sidecar_name = out_path.name + "_data"
+        # If a previous run left a sidecar / loose tensor files, remove them.
+        if (out_path.parent / sidecar_name).exists():
+            (out_path.parent / sidecar_name).unlink()
+        if out_path.exists():
+            out_path.unlink()
+
+        onnx.save_model(
+            model_proto,
+            str(out_path),
+            save_as_external_data=True,
+            all_tensors_to_one_file=True,
+            location=sidecar_name,
+            size_threshold=1024,
+            convert_attribute=False,
+        )
+        onnx.checker.check_model(str(out_path), full_check=False)
+        print(f"  saved {out_path} (+ {sidecar_name})")
+
+
+# ---------------------------------------------------------------------------
+# Parity test.
+# ---------------------------------------------------------------------------
+
+
+def run_parity(
+    wrapper: NAREncoderProjector,
+    fe: Any,
+    waveform: np.ndarray,
+    onnx_path: Path,
+    parity_json: Path,
+    baseline_json: Path,
+    abs_tol: float = 1e-3,
+    argmax_only: bool = False,
+) -> bool:
+    import onnxruntime as ort
+
+    print("\n=== parity check ===")
+    waveform_t = torch.from_numpy(waveform.copy())
+    inputs = fe([waveform_t], device="cpu")
+    input_features = inputs["input_features"].to(torch.float32)
+    attention_mask = inputs["attention_mask"].to(torch.int64)
+    print(f"  input_features: {tuple(input_features.shape)}  attention_mask: {tuple(attention_mask.shape)}")
+
+    # PyTorch reference (full-tensor diff target).
+    print("  running PyTorch wrapper forward")
+    t0 = time.time()
+    with torch.inference_mode():
+        char_pt, bpe_pt, bpe_mask_pt, audio_pt, alen_pt = wrapper(input_features, attention_mask)
+    print(f"    pytorch forward: {time.time() - t0:.2f}s")
+
+    # ORT.
+    print(f"  running ONNX inference: {onnx_path}")
+    sess = ort.InferenceSession(str(onnx_path), providers=["CPUExecutionProvider"])
+    ort_inputs = {
+        "input_features": input_features.numpy().astype(np.float32),
+        "attention_mask": attention_mask.numpy().astype(np.int64),
+    }
+    t0 = time.time()
+    out_names = [o.name for o in sess.get_outputs()]
+    ort_outputs = sess.run(out_names, ort_inputs)
+    print(f"    onnx forward: {time.time() - t0:.2f}s")
+    out_map = dict(zip(out_names, ort_outputs))
+
+    char_ort = out_map["char_logits"]
+    bpe_ort = out_map["bpe_logits_dense"]
+    bpe_mask_ort = out_map["bpe_mask"]
+    audio_ort = out_map["audio_embeds"]
+    alen_ort = out_map["audio_lengths"]
+
+    def diff(
+        name: str,
+        pt: torch.Tensor,
+        ort_arr: np.ndarray,
+        mask: np.ndarray | None = None,
+        tol: float | None = None,
+        check_argmax: bool = False,
+    ) -> dict[str, Any]:
+        local_tol = abs_tol if tol is None else tol
+        pt_np = pt.detach().float().cpu().numpy()
+        if pt_np.shape != ort_arr.shape:
+            return {
+                "name": name,
+                "shape_pt": list(pt_np.shape),
+                "shape_ort": list(ort_arr.shape),
+                "max_abs_err": None,
+                "ok": False,
+                "reason": "shape mismatch",
+                "tol": local_tol,
+            }
+        if mask is not None:
+            # restrict diff to masked-True positions for sparsity-bearing tensors
+            sel_pt = pt_np[mask]
+            sel_ort = ort_arr[mask]
+        else:
+            sel_pt = pt_np
+            sel_ort = ort_arr
+        if sel_pt.size == 0:
+            err = 0.0
+            mean_err = 0.0
+            p99 = 0.0
+        else:
+            abs_err = np.abs(sel_pt - sel_ort)
+            err = float(abs_err.max())
+            mean_err = float(abs_err.mean())
+            p99 = float(np.percentile(abs_err, 99))
+        out: dict[str, Any] = {
+            "name": name,
+            "shape": list(pt_np.shape),
+            "max_abs_err": err,
+            "mean_abs_err": mean_err,
+            "p99_abs_err": p99,
+            "tol": local_tol,
+            "mean_pt": float(pt_np.mean()),
+            "std_pt": float(pt_np.std()),
+            "mean_ort": float(ort_arr.mean()),
+            "std_ort": float(ort_arr.std()),
+            "first10_pt": [float(v) for v in pt_np.flatten()[:10]],
+            "first10_ort": [float(v) for v in ort_arr.flatten()[:10]],
+            "ok": err <= local_tol,
+        }
+        if check_argmax and sel_pt.ndim >= 2:
+            am_pt = sel_pt.reshape(-1, sel_pt.shape[-1]).argmax(-1)
+            am_ort = sel_ort.reshape(-1, sel_ort.shape[-1]).argmax(-1)
+            out["argmax_mismatches"] = int((am_pt != am_ort).sum())
+            out["argmax_total"] = int(am_pt.size)
+            # If max-abs fails but every argmax decision matches, treat this as OK.
+            # The Linear over a 100k-vocab head accumulates fp32 rounding error
+            # uniformly across logits; argmax stability is the actual semantic test.
+            if not out["ok"] and out["argmax_mismatches"] == 0 and err <= 1e-2:
+                out["ok"] = True
+                out["ok_reason"] = "argmax-stable; max_abs slightly above tol due to fp32 cascade"
+            # In INT8 mode the weight quantisation introduces larger logit deltas
+            # but argmax stability is the actual ship gate. Char logits are an
+            # unused intermediate (only BPE feeds the CTC draft) so even argmax
+            # drift there is harmless to the transcript.
+            if argmax_only:
+                if out["argmax_mismatches"] == 0:
+                    out["ok"] = True
+                    out["ok_reason"] = "argmax-only int8 mode; max_abs delta tolerated"
+                elif name == "char_logits":
+                    out["ok"] = True
+                    out["ok_reason"] = (
+                        "argmax-only int8 mode; char_logits drift tolerated "
+                        "(unused downstream)"
+                    )
+        elif argmax_only and not out["ok"]:
+            # Non-logit tensors (audio_embeds, audio_lengths, bpe_mask): under
+            # INT8 the audio_embeds drift but they're not directly compared
+            # downstream - the LLM consumes them and we test argmax there. Soften
+            # the gate so the report still flags the FP32-relative drift.
+            out["ok"] = True
+            out["ok_reason"] = "argmax-only int8 mode; audio_embeds drift tolerated"
+        return out
+
+    bpe_mask_np = bpe_mask_pt.detach().cpu().numpy()
+    diffs = [
+        diff("char_logits", char_pt, char_ort, check_argmax=True),
+        diff(
+            "bpe_logits_dense",
+            bpe_pt,
+            bpe_ort,
+            mask=bpe_mask_np,
+            check_argmax=True,
+        ),
+        diff("bpe_mask", bpe_mask_pt.to(torch.int64), bpe_mask_ort.astype(np.int64)),
+        diff("audio_embeds", audio_pt, audio_ort),
+        diff("audio_lengths", alen_pt, alen_ort.astype(np.int64)),
+    ]
+
+    # Compare projector output against captured baseline (informational).
+    baseline_proj_stats = None
+    if baseline_json.exists():
+        baseline = json.loads(baseline_json.read_text())
+        baseline_proj_stats = baseline.get("projector_output")
+
+    # Stats for each tensor (ONNX side, used for archive).
+    onnx_stats = {
+        "char_logits": tensor_stats(char_ort),
+        "bpe_logits_dense": tensor_stats(bpe_ort),
+        "audio_embeds": tensor_stats(audio_ort),
+        "audio_lengths": tensor_stats(alen_ort),
+    }
+    pt_stats = {
+        "char_logits": tensor_stats(char_pt),
+        "bpe_logits_dense": tensor_stats(bpe_pt),
+        "audio_embeds": tensor_stats(audio_pt),
+        "audio_lengths": tensor_stats(alen_pt),
+    }
+
+    all_ok = all(d["ok"] for d in diffs)
+
+    payload = {
+        "ok": all_ok,
+        "abs_tol": abs_tol,
+        "argmax_only": argmax_only,
+        "input_features": tensor_stats(input_features),
+        "attention_mask_sum": int(attention_mask.sum().item()),
+        "diffs": diffs,
+        "onnx_stats": onnx_stats,
+        "pytorch_stats": pt_stats,
+        "baseline_projector_output": baseline_proj_stats,
+    }
+
+    # Sidecar size info for both the graph under test and (if it exists) the
+    # FP32 sibling that the int8 graph was derived from.
+    sidecar = onnx_path.with_name(onnx_path.name + "_data")
+    payload["graph_path"] = str(onnx_path)
+    payload["graph_size_bytes"] = int(onnx_path.stat().st_size)
+    payload["int8_size_bytes"] = int(sidecar.stat().st_size) if sidecar.exists() else None
+
+    parity_json.parent.mkdir(parents=True, exist_ok=True)
+    parity_json.write_text(json.dumps(payload, indent=2))
+    print(f"  wrote parity report -> {parity_json}")
+
+    print("\n--- parity summary ---")
+    for d in diffs:
+        status = "PASS" if d["ok"] else "FAIL"
+        print(f"  {status}  {d['name']:<20s}  shape={d.get('shape')}  max_abs_err={d.get('max_abs_err')}")
+    print(f"\n{'PASS' if all_ok else 'FAIL'} (abs_tol={abs_tol})")
+
+    return all_ok
+
+
+# ---------------------------------------------------------------------------
+# Main.
+# ---------------------------------------------------------------------------
+
+
+def main() -> None:
+    p = argparse.ArgumentParser()
+    p.add_argument("--audio", default=str(DEFAULT_AUDIO))
+    p.add_argument("--baseline", default=str(DEFAULT_BASELINE))
+    p.add_argument("--model-dir", default=str(DEFAULT_MODEL_DIR))
+    p.add_argument("--out-dir", default=str(DEFAULT_OUT_DIR))
+    p.add_argument("--abs-tol", type=float, default=1e-3)
+    p.add_argument("--skip-export", action="store_true", help="skip the export step (re-run parity only)")
+    p.add_argument(
+        "--graph-suffix",
+        default="",
+        help="suffix appended to graph stem (e.g. '_int8') so parity runs against "
+        "encoder<suffix>.onnx. Parity output goes to encoder_parity<suffix>.json. "
+        "When set, --skip-export is implied.",
+    )
+    args = p.parse_args()
+
+    out_dir = Path(args.out_dir)
+    suffix = args.graph_suffix
+    if suffix and not args.skip_export:
+        print(f"  --graph-suffix={suffix!r} set; implying --skip-export")
+        args.skip_export = True
+    onnx_path = out_dir / f"encoder{suffix}.onnx"
+    parity_json = out_dir / f"encoder_parity{suffix}.json"
+
+    model_dir = Path(args.model_dir)
+
+    print(f"audio:    {args.audio}")
+    print(f"out_dir:  {out_dir}")
+    waveform = load_audio(Path(args.audio))
+    print(f"  duration={waveform.shape[0] / 16000:.2f}s")
+
+    print("loading model...")
+    model, fe = load_nar_model(model_dir)
+
+    print("patching modules for tracing...")
+    patch_for_tracing(model)
+
+    encoder_layer_indices = list(model.config.encoder_layer_indices)
+    print(f"  encoder_layer_indices={encoder_layer_indices}")
+
+    wrapper = NAREncoderProjector(
+        encoder=model.encoder,
+        projector=model.projector,
+        encoder_layer_indices=encoder_layer_indices,
+    )
+    wrapper.eval()
+
+    # Build trace inputs from the actual reference clip so the trace matches the
+    # baseline regime exactly (T=843, B=1).
+    waveform_t = torch.from_numpy(waveform.copy())
+    inputs = fe([waveform_t], device="cpu")
+    sample_features = inputs["input_features"].to(torch.float32)
+    sample_mask = inputs["attention_mask"].to(torch.int64)
+
+    if not args.skip_export:
+        with torch.inference_mode():
+            export_onnx(
+                wrapper=wrapper,
+                sample_input_features=sample_features,
+                sample_attention_mask=sample_mask,
+                out_path=onnx_path,
+                opset=20,
+                ir_version=10,
+            )
+
+    ok = run_parity(
+        wrapper=wrapper,
+        fe=fe,
+        waveform=waveform,
+        onnx_path=onnx_path,
+        parity_json=parity_json,
+        baseline_json=Path(args.baseline),
+        abs_tol=args.abs_tol,
+        argmax_only=bool(suffix),
+    )
+
+    if not ok:
+        raise SystemExit(1)
+
+
+if __name__ == "__main__":
+    main()

granite_export_metadata.json

@@ -0,0 +1,373 @@
+{
+  "variant": "nar",
+  "upstream": {
+    "repo": "ibm-granite/granite-speech-4.1-2b-nar",
+    "url": "https://huggingface.co/ibm-granite/granite-speech-4.1-2b-nar",
+    "license": "Apache-2.0"
+  },
+  "topology": "encoder + editor (NAR bidirectional)",
+  "graphs": [
+    {
+      "name": "encoder.onnx",
+      "sidecar": "encoder.onnx_data",
+      "precision": "fp32",
+      "size_bytes": 1068038,
+      "sidecar_size_bytes": 2491061620,
+      "opset": 20,
+      "ir_version": 10,
+      "ai_onnx_only": true,
+      "inputs": [
+        {
+          "name": "input_features",
+          "shape": [
+            "B",
+            "T",
+            160
+          ],
+          "dtype": "float32"
+        },
+        {
+          "name": "attention_mask",
+          "shape": [
+            "B",
+            "T"
+          ],
+          "dtype": "int64"
+        }
+      ],
+      "outputs": [
+        {
+          "name": "char_logits",
+          "shape": [
+            "B",
+            "T_enc",
+            348
+          ],
+          "dtype": "float32"
+        },
+        {
+          "name": "bpe_logits_dense",
+          "shape": [
+            "B",
+            "T_bpe",
+            "V_bpe"
+          ],
+          "dtype": "float32"
+        },
+        {
+          "name": "bpe_mask",
+          "shape": [
+            "B",
+            "T_bpe"
+          ],
+          "dtype": "float32"
+        },
+        {
+          "name": "audio_embeds",
+          "shape": [
+            "B",
+            "T_audio",
+            2048
+          ],
+          "dtype": "float32"
+        },
+        {
+          "name": "audio_lengths",
+          "shape": [
+            "B"
+          ],
+          "dtype": "int64"
+        }
+      ]
+    },
+    {
+      "name": "encoder_int8.onnx",
+      "sidecar": "encoder_int8.onnx_data",
+      "precision": "int8-weights-only",
+      "size_bytes": 3310811,
+      "sidecar_size_bytes": 935864692,
+      "opset": 20,
+      "ir_version": 10,
+      "ai_onnx_only": true,
+      "inputs": [
+        {
+          "name": "input_features",
+          "shape": [
+            "B",
+            "T",
+            160
+          ],
+          "dtype": "float32"
+        },
+        {
+          "name": "attention_mask",
+          "shape": [
+            "B",
+            "T"
+          ],
+          "dtype": "int64"
+        }
+      ],
+      "outputs": [
+        {
+          "name": "char_logits",
+          "shape": [
+            "B",
+            "T_enc",
+            348
+          ],
+          "dtype": "float32"
+        },
+        {
+          "name": "bpe_logits_dense",
+          "shape": [
+            "B",
+            "T_bpe",
+            "V_bpe"
+          ],
+          "dtype": "float32"
+        },
+        {
+          "name": "bpe_mask",
+          "shape": [
+            "B",
+            "T_bpe"
+          ],
+          "dtype": "float32"
+        },
+        {
+          "name": "audio_embeds",
+          "shape": [
+            "B",
+            "T_audio",
+            2048
+          ],
+          "dtype": "float32"
+        },
+        {
+          "name": "audio_lengths",
+          "shape": [
+            "B"
+          ],
+          "dtype": "int64"
+        }
+      ]
+    },
+    {
+      "name": "editor.onnx",
+      "sidecar": "editor.onnx_data",
+      "precision": "fp32",
+      "size_bytes": 1838936,
+      "sidecar_size_bytes": 6527000576,
+      "opset": 20,
+      "ir_version": 10,
+      "ai_onnx_only": true,
+      "inputs": [
+        {
+          "name": "inputs_embeds",
+          "shape": [
+            1,
+            "N",
+            2048
+          ],
+          "dtype": "float32"
+        },
+        {
+          "name": "position_ids",
+          "shape": [
+            1,
+            "N"
+          ],
+          "dtype": "int64"
+        },
+        {
+          "name": "attention_mask",
+          "shape": [
+            1,
+            1,
+            "N",
+            "N"
+          ],
+          "dtype": "float32"
+        }
+      ],
+      "outputs": [
+        {
+          "name": "logits",
+          "shape": [
+            "B_out",
+            "N",
+            100352
+          ],
+          "dtype": "float32"
+        }
+      ]
+    },
+    {
+      "name": "editor_int8.onnx",
+      "sidecar": "editor_int8.onnx_data",
+      "precision": "int8-weights-only",
+      "size_bytes": 6418060,
+      "sidecar_size_bytes": 1632247808,
+      "opset": 20,
+      "ir_version": 10,
+      "ai_onnx_only": true,
+      "inputs": [
+        {
+          "name": "inputs_embeds",
+          "shape": [
+            1,
+            "N",
+            2048
+          ],
+          "dtype": "float32"
+        },
+        {
+          "name": "position_ids",
+          "shape": [
+            1,
+            "N"
+          ],
+          "dtype": "int64"
+        },
+        {
+          "name": "attention_mask",
+          "shape": [
+            1,
+            1,
+            "N",
+            "N"
+          ],
+          "dtype": "float32"
+        }
+      ],
+      "outputs": [
+        {
+          "name": "logits",
+          "shape": [
+            "B_out",
+            "N",
+            100352
+          ],
+          "dtype": "float32"
+        }
+      ]
+    }
+  ],
+  "parity": {
+    "fp32": {
+      "encoder": {
+        "argmax_only": false,
+        "ok": true,
+        "char_logits": {
+          "max_abs_err": 0.0007429122924804688,
+          "mean_abs_err": 2.1003936126362532e-05,
+          "argmax_mismatches": 0,
+          "argmax_total": 843
+        },
+        "bpe_logits_dense": {
+          "max_abs_err": 0.002044081687927246,
+          "mean_abs_err": 3.075435233768076e-05,
+          "argmax_mismatches": 0,
+          "argmax_total": 211
+        },
+        "audio_embeds": {
+          "max_abs_err": 2.384185791015625e-05,
+          "mean_abs_err": 2.855549382729805e-06
+        },
+        "attention_mask_sum": 843
+      },
+      "editor": {
+        "argmax_only": false,
+        "ok": true,
+        "max_abs_err": 0.00147247314453125,
+        "mean_abs_err": 4.470655039767735e-05,
+        "argmax_mismatches": 0,
+        "argmax_total": 257,
+        "transcript_match": true,
+        "top1_match": 257,
+        "top5_set_match": 257
+      }
+    },
+    "int8": {
+      "encoder": {
+        "argmax_only": true,
+        "ok": true,
+        "char_logits": {
+          "max_abs_err": 7.266112327575684,
+          "mean_abs_err": 0.5243923664093018,
+          "argmax_mismatches": 76,
+          "argmax_total": 843
+        },
+        "bpe_logits_dense": {
+          "max_abs_err": 1.8408704996109009,
+          "mean_abs_err": 0.20938457548618317,
+          "argmax_mismatches": 0,
+          "argmax_total": 211
+        },
+        "audio_embeds": {
+          "max_abs_err": 0.7828707695007324,
+          "mean_abs_err": 0.06720388680696487
+        },
+        "attention_mask_sum": 843
+      },
+      "editor": {
+        "argmax_only": true,
+        "ok": true,
+        "max_abs_err": 94.53038024902344,
+        "mean_abs_err": 8.335189819335938,
+        "argmax_mismatches": 15,
+        "argmax_total": 257,
+        "transcript_match": true,
+        "top1_match": 242,
+        "top5_set_match": 24
+      }
+    }
+  },
+  "multi_clip_parity": {
+    "rows": [
+      {
+        "name": "is-it-more-wood",
+        "duration_s": 46.9,
+        "fp32_byte_exact_vs_pt": true,
+        "int8_byte_exact_vs_pt": false,
+        "int8_wer_vs_pt": 0.036,
+        "int8_vs_fp32_lev": 13,
+        "fp32_transcript": "well, hello, sam. guess who? yeah, it's robert clotworthy, the narrator of your favorite television show, the curse of oak island. yes, i'm the is it possible? could it be? and what else do we say on oak island? a couple of words. they're not coming to me. oh, yeah, more wood. but let's not forget it is an island named after a tree. well, here's the question. why am i reaching out to you? is it possible that i'm reaching out to you because it's your birthday? could it be that emma let the cat out of the bag? well, the answer to those questions is yes. and she said, well, she contacted me. she said, robert, you know, sam is an amazing boyfriend. in fact, she used the word great. she said he is a great boyfriend.",
+        "int8_transcript": "well, hello, sam, guess who? yeah, it's robert clotworthy, the narrator of your favorite television show, the curse of oak island. yes, i'm the... is it possible? could it be? and what else do we say on oak island? couple of words. they're not coming to me. oh, yeah, more wood. let's not forget it is an island named after a tree. well, here's the question. why am i reaching out to you? is it possible that i'm reaching out to you because it's your birthday? could it be that emma let the cat out of the bag? well, the answer to those questions is yes. and she said... well, she contacted me. she said, robert, you know, sam is an amazing boyfriend. in fact, she used the word great. she said he is a great boyfriend.",
+        "pt_transcript": "well, hello, sam. guess who? yeah, it's robert clotworthy, the narrator of your favorite television show, the curse of oak island. yes, i'm the is it possible? could it be? and what else do we say on oak island? a couple of words. they're not coming to me. oh, yeah, more wood. but let's not forget it is an island named after a tree. well, here's the question. why am i reaching out to you? is it possible that i'm reaching out to you because it's your birthday? could it be that emma let the cat out of the bag? well, the answer to those questions is yes. and she said, well, she contacted me. she said, robert, you know, sam is an amazing boyfriend. in fact, she used the word great. she said he is a great boyfriend."
+      },
+      {
+        "name": "two-speakers-1",
+        "duration_s": 93.8,
+        "fp32_byte_exact_vs_pt": false,
+        "int8_byte_exact_vs_pt": false,
+        "int8_wer_vs_pt": 0.0351,
+        "int8_vs_fp32_lev": 26,
+        "fp32_transcript": "today, it is a true honor to speak with demisavas, who is the ceo of deepmind. demis, welcome to the podcast. thanks for having me. first question, given your neuroscience background, how do you think about intelligence? specifically, do you think it's like one higher level general reasoning circuit, or do you think it's thousands of independent subkills and heuristics? well, it's interesting because intelligence is so broad and what we use it for is so sort of generally applicable. i think that suggests that you know there must be some sort of high level common things, you know common kind of algorithmic themes, i think, around how the brain processes the world around us. so of course, then there are specialized parts of the brain that do specific things, but i think there are probably some underlying principles that underpin all of that. yeah. how do you make sense of the fact that in these llms, though, when you give them a lot of data in any specific domain, they tend to get asymmetrically better in that domain? wouldn't we expect a sort of like general improvement across all the different areas? well, i think you, first of all, i think you do actually sometimes get surprising improvement in other domains when you improve in a specific domain. so, for example, when these large models sort of improve at coding, that can actually improve their general reasoning. so there is some evidence of some transfer, although i think we would like a lot more evidence of that. but also, you know, that's how the human brain learns too, is if we experience and practice a lot of things like chess or, you know, writing.",
+        "int8_transcript": "today, it is a true honor to speak with disesisavas, who is the ceo of deepmind. demis, welcome to the podcast. thanks for having me. first question, given your neuroscience background, how do you think about intelligence? specifically, do you think it's like one higher level general reasoning circuit, or do you think it's thousands of independent subkills and heuristics? well, it's interesting because intelligence is so broad and, you, what we use it for is so sort of generally applicable. i think that suggests that, know there must be some sort of high level-level common things know common kind of algorithmic themes, i think, around how the brain processes the world around us. so of course, then there are specialized parts of the brain that do specific things, but i think there are probably some underlying principles that underpin all of that. yeah. how do you make sense of the fact that in these llms though, when you give them a lot of data in any specific domain, they tend to get asymmetrically better in that domain? wouldn't we expect a sort of like general improvement across all the different areas?? well, i think you, first of all, i think you do actually sometimes get surprising improvement in other domains when you improve in a specific domain. so, for example, when these large models sort of improve at coding, that can actually improve their general reasoning. so there is some evidence of some transfer, although i think we would like a lot more evidence of that. but also, you know, that's how the human brain learns too, is if we experience and practice a lot of things like chess or, you know, writing.",
+        "pt_transcript": "today, it is a true honor to speak with demisavas, who is the ceo of deepmind. demis, welcome to the podcast. thanks for having me. first question, given your neuroscience background, how do you think about intelligence? specifically, do you think it's like one higher level general reasoning circuit, or do you think it's thousands of independent subkills and heuristics? well, it's interesting because intelligence is so broad and what we use it for is so sort of generally applicable. i think that suggests that you know there must be some sort of high level common things, you know common kind of algorithmic themes, i think, around how the brain processes the world around us. so of course, then there are specialized parts of the brain that do specific things, but i think there are probably some underlying principles that underpin all of that. yeah. how do you make sense of the fact that in these lms, though, when you give them a lot of data in any specific domain, they tend to get asymmetrically better in that domain? wouldn't we expect a sort of like general improvement across all the different areas? well, i think you, first of all, i think you do actually sometimes get surprising improvement in other domains when you improve in a specific domain. so, for example, when these large models sort of improve at coding, that can actually improve their general reasoning. so there is some evidence of some transfer, although i think we would like a lot more evidence of that. but also, you know, that's how the human brain learns too, is if we experience and practice a lot of things like chess or, you know, writing."
+      },
+      {
+        "name": "two-speakers-2",
+        "duration_s": 38.8,
+        "fp32_byte_exact_vs_pt": true,
+        "int8_byte_exact_vs_pt": false,
+        "int8_wer_vs_pt": 0.0202,
+        "int8_vs_fp32_lev": 10,
+        "fp32_transcript": "for the first time ever, we may have things more intelligent than us. you believe they can understand? yes. you believe they are intelligent? yes. you believe these systems have experiences of their own and can make decisions based on those experiences? in the same sense as people do, yes. are they conscious? i think they probably don't have much self--awareness at present. so in that sense, i don't think they're conscious. will they have self-awareness? yes.. oh, yes, i think they will in time. and so human beings will be the second most intelligent beings on the planet.",
+        "int8_transcript": "for the first time ever, we may have things more intelligent than us. you believe they can understand? yes. you believe they are intelligent? yes. you believe these systems have experiences of their own and can make decisions based on those experiences? in the same sense as people do, yes. are they conscious? i think they probably don't have much self--awareness at present. so in that sense, i don't think they're conscious. will they have self-aware?. i oh, yes, i think they will in time. and so human beings will be the second most intelligent beings on the planet.",
+        "pt_transcript": "for the first time ever, we may have things more intelligent than us. you believe they can understand? yes. you believe they are intelligent? yes. you believe these systems have experiences of their own and can make decisions based on those experiences? in the same sense as people do, yes. are they conscious? i think they probably don't have much self--awareness at present. so in that sense, i don't think they're conscious. will they have self-awareness? yes.. oh, yes, i think they will in time. and so human beings will be the second most intelligent beings on the planet."
+      }
+    ]
+  },
+  "toolchain": {
+    "transformers": "5.8.0",
+    "torch": "2.11.0",
+    "onnx": "1.21.0",
+    "onnxruntime": "1.25.1",
+    "exporter": "torch.onnx.export TorchScript path (dynamo=False)"
+  },
+  "ort_compatibility": "ort 2.0-rc.x (Rust crate); validated against onnxruntime 1.17 - 1.25",
+  "audio_token_id": 100352
+}

preprocessor_config.json

@@ -0,0 +1,11 @@
+{
+  "auto_map": {
+    "AutoFeatureExtractor": "feature_extraction_nle.NLEFeatureExtractor"
+  },
+  "feature_extractor_type": "NLEFeatureExtractor",
+  "hop_length": 160,
+  "n_fft": 512,
+  "n_mels": 80,
+  "sampling_rate": 16000,
+  "win_length": 400
+}

quantise.py

@@ -0,0 +1,299 @@
+# Copyright 2026 Sam McLeod
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Dynamic INT8 (weights-only) quantiser for the Granite Speech 4.1 ONNX
+exports.
+
+Wraps `onnxruntime.quantization.quantize_dynamic` with the conventions used
+by the Granite Speech ONNX bundles:
+
+  - Single external-data sidecar per graph (mirrors the FP32 export layout).
+  - Pure `ai.onnx` opset 20 / IR 10. The default operator set is restricted
+    to `MatMul` so the dynamic quantiser emits `MatMulInteger` (standard
+    `ai.onnx`) rather than the `com.microsoft.Attention` /
+    `com.microsoft.EmbedLayerNormalization` quantised variants. Override at
+    your own risk - those domain ops are forbidden by the parakeet-rs
+    consumer contract.
+  - `per_channel=True` and `weight_type=QInt8` by default (better accuracy
+    on the LLM weight tensors with no measurable speed cost on
+    arm64 / x86 CPU EP).
+
+The script is self-contained (no project-internal imports) so it ships
+inside each Hugging Face bundle alongside the export script.
+
+Usage:
+    python quantise.py --input PATH --output PATH \\
+        [--per-channel | --no-per-channel] \\
+        [--reduce-range] \\
+        [--weight-type qint8|quint8] \\
+        [--op-types MatMul,Gemm] \\
+        [--exclude-pattern REGEX] \\
+        [--exclude-nodes NODE1,NODE2]
+
+Examples:
+    # Quantise the NAR editor with defaults.
+    python quantise.py \\
+        --input  exports/granite-speech-4.1-2b-nar/editor.onnx \\
+        --output exports/granite-speech-4.1-2b-nar/editor_int8.onnx
+
+    # Skip the lm_head MatMul if it hurts parity.
+    python quantise.py \\
+        --input  exports/granite-speech-4.1-2b-nar/editor.onnx \\
+        --output exports/granite-speech-4.1-2b-nar/editor_int8.onnx \\
+        --exclude-nodes /lm_head/MatMul
+"""
+
+from __future__ import annotations
+
+import argparse
+import re
+import sys
+import tempfile
+import time
+from pathlib import Path
+
+import onnx
+from onnxruntime.quantization import QuantType, quantize_dynamic
+
+
+WEIGHT_TYPE_MAP = {
+    "qint8": QuantType.QInt8,
+    "quint8": QuantType.QUInt8,
+}
+
+
+def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
+    p = argparse.ArgumentParser(
+        description="Dynamic INT8 (weights-only) ONNX quantiser for Granite Speech 4.1 graphs.",
+    )
+    p.add_argument(
+        "--input",
+        required=True,
+        type=Path,
+        help="Path to the FP32 .onnx graph (external sidecar must sit alongside it).",
+    )
+    p.add_argument(
+        "--output",
+        required=True,
+        type=Path,
+        help="Destination .onnx path. A single sidecar named <output>_data is written next to it.",
+    )
+    p.add_argument(
+        "--per-channel",
+        dest="per_channel",
+        action="store_true",
+        default=True,
+        help="Quantise weights per output channel (default: on).",
+    )
+    p.add_argument(
+        "--no-per-channel",
+        dest="per_channel",
+        action="store_false",
+        help="Disable per-channel quantisation.",
+    )
+    p.add_argument(
+        "--reduce-range",
+        action="store_true",
+        default=False,
+        help="Quantise to 7 bits instead of 8. Improves accuracy on non-VNNI hardware "
+        "but reduces the quantisation gain. Off by default.",
+    )
+    p.add_argument(
+        "--weight-type",
+        choices=sorted(WEIGHT_TYPE_MAP.keys()),
+        default="qint8",
+        help="Weight quantisation dtype (default: qint8).",
+    )
+    p.add_argument(
+        "--op-types",
+        default="MatMul",
+        help=(
+            "Comma-separated op types to quantise. Default: 'MatMul' (emits "
+            "MatMulInteger only, all ai.onnx). Adding 'Conv' enables ConvInteger "
+            "for the Conformer encoder's depthwise convolutions; this shrinks the "
+            "encoder INT8 sidecar by ~40 percent but on this model family feeds "
+            "enough weight-quant noise into the LLM head that it flips "
+            "capitalisation and drops sentence-final punctuation on short clips - "
+            "see task 17 in dev-plan.md. MatMul-only is the validated default. "
+            "Adding 'Attention' or 'EmbedLayerNormalization' would introduce "
+            "com.microsoft domain ops, which are forbidden by the parakeet-rs "
+            "contract."
+        ),
+    )
+    p.add_argument(
+        "--exclude-pattern",
+        default=None,
+        help="Regex applied to ONNX node names. Matching nodes are excluded from "
+        "quantisation. Useful for skipping e.g. lm_head if its quantisation "
+        "breaks parity.",
+    )
+    p.add_argument(
+        "--exclude-nodes",
+        default="",
+        help="Explicit comma-separated list of node names to exclude from quantisation.",
+    )
+    p.add_argument(
+        "--ir-version",
+        type=int,
+        default=10,
+        help="ONNX IR version to write (default: 10, matches the FP32 exports).",
+    )
+    return p.parse_args(argv)
+
+
+def collect_excluded_nodes(
+    input_path: Path,
+    exclude_pattern: str | None,
+    exclude_nodes: list[str],
+) -> list[str]:
+    """Resolve --exclude-pattern against the FP32 graph's node names and merge
+    with the explicit --exclude-nodes list. Loaded without external data so we
+    only touch the small graph proto.
+    """
+    excluded = set(n for n in exclude_nodes if n)
+    if exclude_pattern:
+        rx = re.compile(exclude_pattern)
+        proto = onnx.load(str(input_path), load_external_data=False)
+        for node in proto.graph.node:
+            if node.name and rx.search(node.name):
+                excluded.add(node.name)
+    return sorted(excluded)
+
+
+def assert_pure_ai_onnx(model_path: Path) -> list[str]:
+    """Reload the produced graph and verify only `ai.onnx` nodes are present.
+    Returns the sorted list of domains for reporting.
+    """
+    proto = onnx.load(str(model_path), load_external_data=False)
+    domains = sorted({(n.domain or "ai.onnx") for n in proto.graph.node})
+    forbidden = [d for d in domains if d not in ("ai.onnx", "")]
+    if forbidden:
+        raise RuntimeError(
+            f"Quantised graph contains forbidden op domains {forbidden}. "
+            "Re-run with a narrower --op-types list."
+        )
+    return domains
+
+
+def consolidate_single_sidecar(
+    quantised_in: Path,
+    final_out: Path,
+    ir_version: int,
+) -> None:
+    """The dynamic quantiser may scatter weights across multiple external-data
+    files. Reload + resave through a tempdir to land on the single-sidecar
+    layout that matches the FP32 exports.
+    """
+    print("  consolidating to single .onnx_data sidecar")
+    proto = onnx.load(str(quantised_in), load_external_data=True)
+    if proto.ir_version < ir_version:
+        proto.ir_version = ir_version
+
+    for tensor in proto.graph.initializer:
+        tensor.ClearField("data_location")
+        tensor.ClearField("external_data")
+
+    sidecar_name = final_out.name + "_data"
+    if (final_out.parent / sidecar_name).exists():
+        (final_out.parent / sidecar_name).unlink()
+    if final_out.exists():
+        final_out.unlink()
+    final_out.parent.mkdir(parents=True, exist_ok=True)
+
+    onnx.save_model(
+        proto,
+        str(final_out),
+        save_as_external_data=True,
+        all_tensors_to_one_file=True,
+        location=sidecar_name,
+        size_threshold=1024,
+        convert_attribute=False,
+    )
+    onnx.checker.check_model(str(final_out), full_check=False)
+
+
+def quantise_graph(args: argparse.Namespace) -> None:
+    input_path: Path = args.input.resolve()
+    output_path: Path = args.output.resolve()
+    if not input_path.exists():
+        raise SystemExit(f"input not found: {input_path}")
+    op_types = [s.strip() for s in args.op_types.split(",") if s.strip()]
+    explicit_excludes = [s.strip() for s in args.exclude_nodes.split(",") if s.strip()]
+
+    excluded = collect_excluded_nodes(input_path, args.exclude_pattern, explicit_excludes)
+    weight_type = WEIGHT_TYPE_MAP[args.weight_type]
+
+    print(f"input:        {input_path}")
+    print(f"output:       {output_path}")
+    print(f"op_types:     {op_types}")
+    print(f"per_channel:  {args.per_channel}")
+    print(f"reduce_range: {args.reduce_range}")
+    print(f"weight_type:  {args.weight_type}")
+    if excluded:
+        print(f"excluded nodes ({len(excluded)}): {excluded}")
+    else:
+        print("excluded nodes: (none)")
+
+    fp32_size = input_path.stat().st_size
+    sidecar = input_path.with_name(input_path.name + "_data")
+    fp32_data_size = sidecar.stat().st_size if sidecar.exists() else 0
+    print(
+        f"  fp32 graph={fp32_size / 1e6:.2f} MB  "
+        f"sidecar={fp32_data_size / 1e9:.2f} GB"
+    )
+
+    with tempfile.TemporaryDirectory(prefix="quantise_int8_") as scratch_dir:
+        scratch_path = Path(scratch_dir) / output_path.name
+        t0 = time.time()
+        quantize_dynamic(
+            model_input=input_path,
+            model_output=scratch_path,
+            op_types_to_quantize=op_types,
+            per_channel=args.per_channel,
+            reduce_range=args.reduce_range,
+            weight_type=weight_type,
+            nodes_to_exclude=excluded or None,
+            use_external_data_format=True,
+        )
+        print(f"  quantize_dynamic done in {time.time() - t0:.1f}s")
+
+        # Stage 2: consolidate any scattered external-data files into a single
+        # sidecar at the final destination.
+        consolidate_single_sidecar(scratch_path, output_path, args.ir_version)
+
+    # Verify pure ai.onnx after the move.
+    domains = assert_pure_ai_onnx(output_path)
+    int8_size = output_path.stat().st_size
+    int8_data = output_path.with_name(output_path.name + "_data")
+    int8_data_size = int8_data.stat().st_size if int8_data.exists() else 0
+    print(
+        f"  saved {output_path} (+ {int8_data.name})  "
+        f"graph={int8_size / 1e6:.2f} MB  sidecar={int8_data_size / 1e9:.2f} GB"
+    )
+    print(f"  node-domains={domains}")
+    if fp32_data_size > 0:
+        ratio = int8_data_size / fp32_data_size
+        print(f"  sidecar size ratio (int8 / fp32) = {ratio:.3f}")
+
+
+def main(argv: list[str] | None = None) -> None:
+    args = parse_args(argv)
+    try:
+        quantise_graph(args)
+    except RuntimeError as exc:
+        print(f"FAIL: {exc}", file=sys.stderr)
+        raise SystemExit(2) from exc
+
+
+if __name__ == "__main__":
+    main()

special_tokens_map.json

@@ -0,0 +1,30 @@
+{
+  "bos_token": {
+    "content": "<|end_of_text|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "eos_token": {
+    "content": "<|end_of_text|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": {
+    "content": "<|pad|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "unk_token": {
+    "content": "<|unk|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

tokenizer_config.json

@@ -0,0 +1,783 @@
+{
+  "add_bos_token": false,
+  "add_prefix_space": false,
+  "added_tokens_decoder": {
+    "100256": {
+      "content": "<|pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "100257": {
+      "content": "<|end_of_text|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "100258": {
+      "content": "<|fim_prefix|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "100259": {
+      "content": "<|fim_middle|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "100260": {
+      "content": "<|fim_suffix|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "100261": {
+      "content": "<|fim_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "100262": {
+      "content": "<|filename|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "100263": {
+      "content": "<|reponame|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "100264": {
+      "content": "<|start_of_role|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "100265": {
+      "content": "<|end_of_role|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "100266": {
+      "content": "<|unused_1|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "100267": {
+      "content": "<|start_of_plugin|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "100268": {
+      "content": "<|end_of_plugin|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "100269": {
+      "content": "<|unk|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "100270": {
+      "content": "<tool_call>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "100271": {
+      "content": "</tool_call>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "100272": {
+      "content": "<tool_response>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "100273": {
+      "content": "</tool_response>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "100274": {
+      "content": "<think>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "100275": {
+      "content": "</think>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "100276": {
+      "content": "<think_on>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "100277": {
+      "content": "<think_off>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "100278": {
+      "content": "<schema>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "100279": {
+      "content": "</schema>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "100280": {
+      "content": "<tools>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "100281": {
+      "content": "</tools>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "100282": {
+      "content": "<documents>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "100283": {
+      "content": "</documents>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "100284": {
+      "content": "<|unused_15|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "100285": {
+      "content": "<|unused_16|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "100286": {
+      "content": "<|unused_17|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "100287": {
+      "content": "<|unused_18|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "100288": {
+      "content": "<|unused_19|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "100289": {
+      "content": "<|unused_20|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "100290": {
+      "content": "<|unused_21|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "100291": {
+      "content": "<|unused_22|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "100292": {
+      "content": "<|unused_23|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "100293": {
+      "content": "<|unused_24|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "100294": {
+      "content": "<|unused_25|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "100295": {
+      "content": "<|unused_26|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "100296": {
+      "content": "<|unused_27|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "100297": {
+      "content": "<|unused_28|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "100298": {
+      "content": "<|unused_29|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "100299": {
+      "content": "<|unused_30|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "100300": {
+      "content": "<|unused_31|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "100301": {
+      "content": "<|unused_32|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "100302": {
+      "content": "<|unused_33|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "100303": {
+      "content": "<|unused_34|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "100304": {
+      "content": "<|unused_35|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "100305": {
+      "content": "<|unused_36|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "100306": {
+      "content": "<|unused_37|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "100307": {
+      "content": "<|unused_38|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "100308": {
+      "content": "<|unused_39|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "100309": {
+      "content": "<|unused_40|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "100310": {
+      "content": "<|unused_41|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "100311": {
+      "content": "<|unused_42|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "100312": {
+      "content": "<|unused_43|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "100313": {
+      "content": "<|unused_44|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "100314": {
+      "content": "<|unused_45|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "100315": {
+      "content": "<|unused_46|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "100316": {
+      "content": "<|unused_47|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "100317": {
+      "content": "<|unused_48|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "100318": {
+      "content": "<|unused_49|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "100319": {
+      "content": "<|unused_50|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "100320": {
+      "content": "<|unused_51|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "100321": {
+      "content": "<|unused_52|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "100322": {
+      "content": "<|unused_53|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "100323": {
+      "content": "<|unused_54|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "100324": {
+      "content": "<|unused_55|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "100325": {
+      "content": "<|unused_56|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "100326": {
+      "content": "<|unused_57|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "100327": {
+      "content": "<|unused_58|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "100328": {
+      "content": "<|unused_59|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "100329": {
+      "content": "<|unused_60|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "100330": {
+      "content": "<|unused_61|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "100331": {
+      "content": "<|unused_62|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "100332": {
+      "content": "<|unused_63|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "100333": {
+      "content": "<|unused_64|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "100334": {
+      "content": "<|unused_65|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "100335": {
+      "content": "<|unused_66|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "100336": {
+      "content": "<|unused_67|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "100337": {
+      "content": "<|unused_68|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "100338": {
+      "content": "<|unused_69|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "100339": {
+      "content": "<|unused_70|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "100340": {
+      "content": "<|unused_71|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "100341": {
+      "content": "<|unused_72|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "100342": {
+      "content": "<|unused_73|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "100343": {
+      "content": "<|unused_74|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "100344": {
+      "content": "<|unused_75|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "100345": {
+      "content": "<|unused_76|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "100346": {
+      "content": "<|unused_77|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "100347": {
+      "content": "<|unused_78|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "100348": {
+      "content": "<|unused_79|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "100349": {
+      "content": "<|unused_80|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "100350": {
+      "content": "<|unused_81|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "100351": {
+      "content": "<|unused_82|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "bos_token": "<|end_of_text|>",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|end_of_text|>",
+  "extra_special_tokens": {},
+  "model_max_length": 1000000000000000019884624838656,
+  "pad_token": "<|pad|>",
+  "padding_side": "left",
+  "tokenizer_class": "GPT2Tokenizer",
+  "unk_token": "<|unk|>"
+}