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