add config.json, remove non-CTC scripts

- Restore config.json (model metadata: vocab_size, token IDs)
- Remove convert/parakeet-tdt-v2-0.6b/ (TDT 0.6b, not CTC)
- Remove scripts/run_benchmarks.py (general FluidAudio, not CTC)
- Keep convert/parakeet-tdt-ctc-110m/ (CTC conversion scripts)

config.json

@@ -0,0 +1,7 @@
+{
+  "bos_token_id": 1,
+  "eos_token_id": 2,
+  "nemo_model_type": "parakeet",
+  "pad_token_id": 0,
+  "vocab_size": 1024,
+}

convert/parakeet-tdt-v2-0.6b/coreml/compare-components.py

@@ -1,958 +0,0 @@
-#!/usr/bin/env python3
-"""Compare Parakeet TDT v2 Torch vs CoreML components on a fixed 15s window.
-
-Writes numeric diffs to the specified output directory (metadata.json) and
-saves plots under a repo-tracked directory: plots/<script-name>/.
-"""
-from __future__ import annotations
-
-import json
-import time
-from dataclasses import dataclass
-from pathlib import Path
-from typing import Dict, Optional, Tuple
-
-import coremltools as ct
-import numpy as np
-import soundfile as sf
-import torch
-import typer
-
-import nemo.collections.asr as nemo_asr
-
-# Optional plotting
-try:
-    import matplotlib
-    matplotlib.use("Agg")
-    import matplotlib.pyplot as plt
-    HAS_MPL = True
-except Exception:
-    HAS_MPL = False
-
-
-@dataclass
-class ValidationSettings:
-    audio_path: Optional[Path]
-    seconds: float
-    seed: Optional[int]
-    rtol: float
-    atol: float
-
-
-def _compute_length(seconds: float, sample_rate: int) -> int:
-    return int(round(seconds * sample_rate))
-
-
-def _prepare_audio(
-    validation_audio: Optional[Path],
-    sample_rate: int,
-    max_samples: int,
-    seed: Optional[int],
-) -> torch.Tensor:
-    if validation_audio is None:
-        if seed is not None:
-            torch.manual_seed(seed)
-        return torch.randn(1, max_samples, dtype=torch.float32)
-
-    data, sr = sf.read(str(validation_audio), dtype="float32")
-    if sr != sample_rate:
-        raise typer.BadParameter(
-            f"Validation audio sample rate {sr} does not match model rate {sample_rate}"
-        )
-    if data.ndim > 1:
-        data = data[:, 0]
-    if data.size == 0:
-        raise typer.BadParameter("Validation audio is empty")
-    if data.size < max_samples:
-        data = np.pad(data, (0, max_samples - data.size))
-    elif data.size > max_samples:
-        data = data[:max_samples]
-    return torch.from_numpy(data).unsqueeze(0).to(dtype=torch.float32)
-
-
-def _np(x: torch.Tensor, dtype=None) -> np.ndarray:
-    arr = x.detach().cpu().numpy()
-    if dtype is not None:
-        return arr.astype(dtype, copy=False)
-    return arr
-
-
-def _to_t(x) -> torch.Tensor:
-    if isinstance(x, torch.Tensor):
-        return x.detach().cpu()
-    elif isinstance(x, np.ndarray):
-        # Ensure a separate tensor (avoid shared memory weirdness)
-        return torch.from_numpy(np.array(x, copy=True))
-    else:
-        return torch.tensor(x)
-
-
-def _max_diffs(a, b, rtol: float, atol: float) -> Tuple[float, float, bool]:
-    # Use NumPy for comparisons to avoid invoking the PyTorch C-API in contexts
-    # where the GIL may not be held (which can trigger PyEval_SaveThread errors).
-    na = np.array(a, dtype=np.float32, copy=True)
-    nb = np.array(b, dtype=np.float32, copy=True)
-    if na.size == 0:
-        return 0.0, 0.0, True
-    diff = np.abs(na - nb)
-    max_abs = float(diff.max())
-    denom = np.maximum(np.abs(na), np.abs(nb))
-    with np.errstate(divide="ignore", invalid="ignore"):
-        rel = np.where(denom == 0.0, 0.0, diff / denom)
-    max_rel = float(rel.max())
-    ok = bool(np.allclose(na, nb, rtol=rtol, atol=atol))
-    return max_abs, max_rel, ok
-
-
-def _plot_line(x_ref: np.ndarray, x_ml: np.ndarray, title: str, path: Path, also_delta: bool = False):
-    if not HAS_MPL:
-        return None
-    if also_delta:
-        fig, axes = plt.subplots(2, 1, figsize=(8, 5), sharex=True)
-        axes[0].plot(x_ref, label="torch", linewidth=1)
-        axes[0].plot(x_ml, label="coreml", linewidth=1, alpha=0.8)
-        axes[0].set_title(title)
-        axes[0].legend()
-        delta = np.asarray(x_ref) - np.asarray(x_ml)
-        axes[1].plot(delta, color="C3", linewidth=1)
-        axes[1].set_title("Delta (torch - coreml)")
-        axes[1].set_xlabel("time/step")
-        plt.tight_layout()
-        plt.savefig(path)
-        plt.close(fig)
-    else:
-        plt.figure(figsize=(8, 3))
-        plt.plot(x_ref, label="torch", linewidth=1)
-        plt.plot(x_ml, label="coreml", linewidth=1, alpha=0.8)
-        plt.title(title)
-        plt.legend()
-        plt.tight_layout()
-        plt.savefig(path)
-        plt.close()
-    return str(path.name)
-
-
-def _plot_image(img: np.ndarray, title: str, path: Path, vmin=None, vmax=None):
-    if not HAS_MPL:
-        return None
-    plt.figure(figsize=(6, 4))
-    plt.imshow(img, aspect='auto', origin='lower', interpolation='nearest', vmin=vmin, vmax=vmax)
-    plt.title(title)
-    plt.colorbar(shrink=0.8)
-    plt.tight_layout()
-    plt.savefig(path)
-    plt.close()
-    return str(path.name)
-
-
-def _plot_mel_composite(
-    mel_torch: np.ndarray,
-    mel_coreml: np.ndarray,
-    path: Path,
-    vmin=None,
-    vmax=None,
-):
-    """Create a single PNG with mel torch, mel coreml, abs diff heatmap, and mean-over-time curves with delta."""
-    if not HAS_MPL:
-        return None
-    mel_torch = np.asarray(mel_torch)
-    mel_coreml = np.asarray(mel_coreml)
-    absdiff = np.abs(mel_torch - mel_coreml)
-    mean_t = mel_torch.mean(axis=0)
-    mean_c = mel_coreml.mean(axis=0)
-    delta = mean_t - mean_c
-
-    fig = plt.figure(figsize=(12, 8))
-    gs = fig.add_gridspec(2, 2, height_ratios=[1, 1])
-    ax1 = fig.add_subplot(gs[0, 0])
-    im1 = ax1.imshow(mel_torch, aspect='auto', origin='lower', interpolation='nearest', vmin=vmin, vmax=vmax)
-    ax1.set_title("Mel (Torch)")
-    fig.colorbar(im1, ax=ax1, shrink=0.8)
-
-    ax2 = fig.add_subplot(gs[0, 1])
-    im2 = ax2.imshow(mel_coreml, aspect='auto', origin='lower', interpolation='nearest', vmin=vmin, vmax=vmax)
-    ax2.set_title("Mel (CoreML)")
-    fig.colorbar(im2, ax=ax2, shrink=0.8)
-
-    ax3 = fig.add_subplot(gs[1, 0])
-    im3 = ax3.imshow(absdiff, aspect='auto', origin='lower', interpolation='nearest')
-    ax3.set_title("Mel |diff|")
-    fig.colorbar(im3, ax=ax3, shrink=0.8)
-
-    ax4 = fig.add_subplot(gs[1, 1])
-    ax4.plot(mean_t, label="torch", linewidth=1)
-    ax4.plot(mean_c, label="coreml", linewidth=1, alpha=0.8)
-    ax4.plot(delta, label="delta", linewidth=1, color="C3")
-    ax4.set_title("Mel mean over time + delta")
-    ax4.legend()
-
-    plt.tight_layout()
-    plt.savefig(path)
-    plt.close(fig)
-    return str(path.name)
-
-
-def _plot_latency_bars(
-    labels,
-    torch_means,
-    torch_stds,
-    coreml_means,
-    coreml_stds,
-    path: Path,
-):
-    if not HAS_MPL:
-        return None
-    x = np.arange(len(labels))
-    width = 0.35
-    fig, ax = plt.subplots(figsize=(8, 4))
-    b1 = ax.bar(x - width/2, torch_means, width, yerr=torch_stds, label="torch", color="C0", alpha=0.9)
-    b2 = ax.bar(x + width/2, coreml_means, width, yerr=coreml_stds, label="coreml", color="C1", alpha=0.9)
-    ax.set_xticks(x, labels, rotation=15)
-    ax.set_ylabel("latency (ms)")
-    ax.set_title("Component latency (15s window inputs)")
-    ax.legend()
-    # Add value labels on bars
-    def _annotate(bars):
-        for bar in bars:
-            h = bar.get_height()
-            if np.isnan(h):
-                continue
-            ax.annotate(f"{h:.0f}",
-                        xy=(bar.get_x() + bar.get_width()/2, h),
-                        xytext=(0, 3), textcoords="offset points",
-                        ha='center', va='bottom', fontsize=8)
-    _annotate(b1)
-    _annotate(b2)
-    plt.tight_layout()
-    plt.savefig(path)
-    plt.close(fig)
-    return str(path.name)
-
-
-def _plot_speedup_bars(labels, torch_means, coreml_means, path: Path):
-    if not HAS_MPL:
-        return None
-    speedup = []
-    for t, c in zip(torch_means, coreml_means):
-        if c and c > 0:
-            speedup.append(float(t) / float(c))
-        else:
-            speedup.append(np.nan)
-    x = np.arange(len(labels))
-    fig, ax = plt.subplots(figsize=(8, 4))
-    bars = ax.bar(x, speedup, color="C2")
-    ax.set_xticks(x, labels, rotation=15)
-    ax.set_ylabel("torch/coreml speedup")
-    ax.set_title("CoreML speedup vs Torch (higher is better)")
-    ax.axhline(1.0, color="gray", linestyle="--", linewidth=1)
-    # Add value labels
-    for bar in bars:
-        h = bar.get_height()
-        if np.isnan(h):
-            continue
-        ax.annotate(f"{h:.2f}",
-                    xy=(bar.get_x() + bar.get_width()/2, h),
-                    xytext=(0, 3), textcoords="offset points",
-                    ha='center', va='bottom', fontsize=8)
-    plt.tight_layout()
-    plt.savefig(path)
-    plt.close(fig)
-    return str(path.name)
-
-
-
-app = typer.Typer(add_completion=False, pretty_exceptions_show_locals=False)
-
-
-@app.command()
-def compare(
-    output_dir: Path = typer.Option(Path("parakeet_coreml"), help="Directory containing mlpackages + metadata.json"),
-    nemo_path: Optional[Path] = typer.Option(None, "--nemo-path", exists=True, resolve_path=True, help="Path to .nemo checkpoint"),
-    model_id: str = typer.Option("nvidia/parakeet-tdt-0.6b-v2", "--model-id", help="HF model id if --nemo-path omitted"),
-    validation_audio: Optional[Path] = typer.Option(None, exists=True, resolve_path=True, help="15s, 16kHz wav for validation (defaults to audio/yc_first_minute_16k_15s.wav if present)"),
-    seed: Optional[int] = typer.Option(None, help="Random seed for synthetic input when audio is not provided"),
-    rtol: float = typer.Option(1e-3, help="Relative tolerance for comparisons"),
-    atol: float = typer.Option(1e-4, help="Absolute tolerance for comparisons"),
-    runs: int = typer.Option(10, help="Timed runs per model for latency measurement"),
-    warmup: int = typer.Option(3, help="Warmup runs before timing (compilation, caches)"),
-    symbol_steps: int = typer.Option(
-        32,
-        help="Number of sequential decoder steps to validate with streaming U=1 inputs",
-    ),
-) -> None:
-    """Run Torch vs CoreML comparisons and update metadata.json with plots and diffs."""
-    output_dir.mkdir(parents=True, exist_ok=True)
-    if symbol_steps < 1:
-        raise typer.BadParameter("symbol_steps must be >= 1")
-
-    meta_path = output_dir / "metadata.json"
-    exported_meta: Dict[str, object] = {}
-    if meta_path.exists():
-        try:
-            exported_meta = json.loads(meta_path.read_text())
-        except Exception:
-            exported_meta = {}
-    exported_max_u = int(exported_meta.get("max_symbol_steps", 1))
-    if exported_max_u != 1:
-        typer.echo(
-            f"Note: CoreML export reports max_symbol_steps={exported_max_u}; "
-            "comparison still drives decoder step-wise with U=1 inputs."
-        )
-    if nemo_path is not None:
-        typer.echo(f"Loading NeMo model from {nemo_path}…")
-        asr_model = nemo_asr.models.EncDecRNNTBPEModel.restore_from(str(nemo_path), map_location="cpu")
-    else:
-        typer.echo(f"Downloading NeMo model via {model_id}…")
-        asr_model = nemo_asr.models.EncDecRNNTBPEModel.from_pretrained(model_id, map_location="cpu")
-    asr_model.eval()
-
-    sample_rate = int(asr_model.cfg.preprocessor.sample_rate)
-    max_samples = _compute_length(15.0, sample_rate)
-    default_audio = (Path(__file__).parent / "audio" / "yc_first_minute_16k_15s.wav").resolve()
-    chosen_audio = validation_audio if validation_audio is not None else (default_audio if default_audio.exists() else None)
-    if chosen_audio is not None and validation_audio is None:
-        typer.echo(f"Using default validation audio: {chosen_audio}")
-        
-    audio_tensor = _prepare_audio(chosen_audio, sample_rate, max_samples, seed)
-    audio_length = torch.tensor([max_samples], dtype=torch.int32)
-
-    asr_model.decoder._rnnt_export = True
-    # Disable fused loss/WER computation for simpler joint inference
-    asr_model.joint.set_fuse_loss_wer(False)
-    # Important: ensure the joint returns raw logits (not log-softmax)
-    # RNNTJoint applies log_softmax on CPU by default when `log_softmax is None`.
-    # Our exported CoreML joint emits pre-softmax logits, so make the Torch
-    # reference do the same to avoid systematic offsets in comparisons/plots.
-    try:
-        # Some versions expose this as a plain attribute
-        asr_model.joint.log_softmax = False
-    except Exception:
-        pass
-
-    # Generate reference outputs directly from NeMo model components
-    with torch.inference_mode():
-        # Preprocessor - direct NeMo call
-        mel_ref, mel_length_ref = asr_model.preprocessor(
-            input_signal=audio_tensor,
-            length=audio_length.to(dtype=torch.long)
-        )
-        mel_length_ref = mel_length_ref.to(dtype=torch.int32)
-
-        # Encoder - direct NeMo call
-        encoder_ref, encoder_length_ref = asr_model.encoder(
-            audio_signal=mel_ref,
-            length=mel_length_ref.to(dtype=torch.long)
-        )
-        encoder_length_ref = encoder_length_ref.to(dtype=torch.int32)
-
-    vocab_size = int(asr_model.tokenizer.vocab_size)
-    num_extra = int(asr_model.joint.num_extra_outputs)
-    decoder_hidden = int(asr_model.decoder.pred_hidden)
-    decoder_layers = int(asr_model.decoder.pred_rnn_layers)
-    blank_id = int(asr_model.decoder.blank_idx)
-
-    blank_targets = torch.tensor([[blank_id]], dtype=torch.int32)
-    blank_target_lengths = torch.tensor([1], dtype=torch.int32)
-    blank_targets_long = blank_targets.to(dtype=torch.long)
-    blank_target_lengths_long = blank_target_lengths.to(dtype=torch.long)
-
-    def _decoder_rollout_torch(num_steps: int) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-        outputs = []
-        h_state = torch.zeros(decoder_layers, 1, decoder_hidden, dtype=torch.float32)
-        c_state = torch.zeros(decoder_layers, 1, decoder_hidden, dtype=torch.float32)
-        state = [h_state, c_state]
-        with torch.inference_mode():
-            for _ in range(num_steps):
-                y, _, new_state = asr_model.decoder(
-                    targets=blank_targets_long,
-                    target_length=blank_target_lengths_long,
-                    states=state,
-                )
-                outputs.append(y.detach())
-                state = [new_state[0].detach(), new_state[1].detach()]
-        if outputs:
-            decoder_seq = torch.cat(outputs, dim=-1)
-        else:
-            decoder_seq = torch.zeros(1, decoder_hidden, 0, dtype=torch.float32)
-        return decoder_seq, state[0], state[1]
-
-    decoder_ref, h_ref, c_ref = _decoder_rollout_torch(symbol_steps)
-
-    with torch.inference_mode():
-        logits_ref = asr_model.joint(
-            encoder_outputs=encoder_ref,
-            decoder_outputs=decoder_ref,
-        )
-
-    # Convert tensors to numpy for CoreML
-    def _np32(x):
-        return np.array(x.detach().cpu().numpy(), dtype=np.float32, copy=True)
-
-    # Prepare plot dir (write to repo-tracked plots/<script-name>/)
-    plots_root = Path(__file__).parent / "plots"
-    plots_dir = plots_root / Path(__file__).stem
-    plots_dir.mkdir(parents=True, exist_ok=True)
-
-    encoder_np = _np32(encoder_ref)
-    decoder_ref_np = _np32(decoder_ref)
-
-    summary: Dict[str, object] = {
-        "requested": True,
-        "status": "ok",
-        "atol": atol,
-        "rtol": rtol,
-        "symbol_steps": int(symbol_steps),
-        "audio_path": None if validation_audio is None else str(validation_audio),
-        "components": {},
-    }
-
-    # Preprocessor
-    pre = ct.models.MLModel(str(output_dir / "parakeet_preprocessor.mlpackage"), compute_units=ct.ComputeUnit.CPU_AND_NE)
-    t0 = time.perf_counter()
-    pre_out = pre.predict({"audio_signal": _np32(audio_tensor), "audio_length": _np32(audio_length).astype(np.int32)})
-    t1 = time.perf_counter()
-    pre_first_ms = (t1 - t0) * 1000.0
-    mel_ml = np.array(pre_out["mel"], dtype=np.float32, copy=True)
-    mel_len_ml = np.array(pre_out["mel_length"], dtype=np.int32, copy=True)
-    pre_atol, pre_rtol = max(atol, 1.0), max(rtol, 1e-2)
-    a_mel, r_mel, ok_mel = _max_diffs(_np32(mel_ref), mel_ml, pre_rtol, pre_atol)
-    ok_len = int(_np32(mel_length_ref).astype(np.int32)[0]) == int(np.array(mel_len_ml).astype(np.int32)[0])
-    mel_t = _np32(mel_ref)[0]
-    mel_c = mel_ml[0]
-    vmin = float(min(mel_t.min(), mel_c.min()))
-    vmax = float(max(mel_t.max(), mel_c.max()))
-    pre_plots = {
-        "mel_composite.png": _plot_mel_composite(mel_t, mel_c, plots_dir / "mel_composite.png", vmin=vmin, vmax=vmax),
-    }
-    # Latency measurements: Torch and CoreML
-    def _time_coreml(model: ct.models.MLModel, inputs: Dict[str, np.ndarray]) -> Tuple[float, float]:
-        # Warmup
-        for _ in range(max(0, warmup)):
-            _ = model.predict(inputs)
-        times = []
-        for _ in range(max(1, runs)):
-            t0 = time.perf_counter()
-            _ = model.predict(inputs)
-            t1 = time.perf_counter()
-            times.append((t1 - t0) * 1000.0)
-        arr = np.array(times, dtype=np.float64)
-        return float(arr.mean()), float(arr.std(ddof=1) if arr.size > 1 else 0.0)
-
-    def _time_torch(fn, *args, **kwargs) -> Tuple[float, float]:
-        with torch.inference_mode():
-            for _ in range(max(0, warmup)):
-                _ = fn(*args, **kwargs)
-            times = []
-            for _ in range(max(1, runs)):
-                t0 = time.perf_counter()
-                _ = fn(*args, **kwargs)
-                t1 = time.perf_counter()
-                times.append((t1 - t0) * 1000.0)
-        arr = np.array(times, dtype=np.float64)
-        return float(arr.mean()), float(arr.std(ddof=1) if arr.size > 1 else 0.0)
-
-    pre_torch_ms_mean, pre_torch_ms_std = _time_torch(
-        asr_model.preprocessor, input_signal=audio_tensor, length=audio_length.to(dtype=torch.long)
-    )
-    pre_coreml_ms_mean, pre_coreml_ms_std = _time_coreml(
-        pre,
-        {"audio_signal": _np32(audio_tensor), "audio_length": _np32(audio_length).astype(np.int32)},
-    )
-    seconds = 15.0
-    pre_coreml_rtf = float(pre_coreml_ms_mean / (seconds * 1000.0)) if pre_coreml_ms_mean > 0 else None
-    pre_torch_rtf = float(pre_torch_ms_mean / (seconds * 1000.0)) if pre_torch_ms_mean > 0 else None
-
-    summary["components"]["preprocessor"] = {
-        "mel": {"max_abs": a_mel, "max_rel": r_mel, "match": bool(ok_mel)},
-        "length_match": bool(ok_len),
-        "latency": {
-            "runs": int(runs),
-            "warmup": int(warmup),
-            "coreml_first_ms": pre_first_ms,
-            "torch_ms": {"mean": pre_torch_ms_mean, "std": pre_torch_ms_std},
-            "coreml_ms": {"mean": pre_coreml_ms_mean, "std": pre_coreml_ms_std},
-            "rtf": {"torch": pre_torch_rtf, "coreml": pre_coreml_rtf},
-        },
-        "plots": {k: v for k, v in pre_plots.items() if v},
-    }
-
-    # Encoder
-    enc = ct.models.MLModel(str(output_dir / "parakeet_encoder.mlpackage"), compute_units=ct.ComputeUnit.CPU_AND_NE)
-    t0 = time.perf_counter()
-    enc_out = enc.predict({"mel": _np32(mel_ref), "mel_length": _np32(mel_length_ref).astype(np.int32)})
-    t1 = time.perf_counter()
-    enc_first_ms = (t1 - t0) * 1000.0
-    enc_ml = np.array(enc_out["encoder"], dtype=np.float32, copy=True)
-    enc_len_ml = np.array(enc_out["encoder_length"], dtype=np.int32, copy=True)
-    a_enc, r_enc, ok_enc = _max_diffs(_np32(encoder_ref), enc_ml, max(rtol, 5e-3), max(atol, 5e-2))
-    ok_enc_len = int(_np32(encoder_length_ref).astype(np.int32)[0]) == int(np.array(enc_len_ml).astype(np.int32)[0])
-    enc_t = _np32(encoder_ref)[0]
-    enc_c = enc_ml[0]
-    enc_plots = {
-        "encoder_time_l2.png": _plot_line(
-            np.linalg.norm(enc_t, axis=0),  # L2 norm over features (D) for each time step
-            np.linalg.norm(enc_c, axis=0),  # enc_t shape is (D, T), so axis=0 is features
-            "Encoder L2 over time",
-            plots_dir / "encoder_time_l2.png",
-            also_delta=True,
-        ),
-    }
-    enc_torch_ms_mean, enc_torch_ms_std = _time_torch(
-        asr_model.encoder, audio_signal=mel_ref, length=mel_length_ref.to(dtype=torch.long)
-    )
-    enc_coreml_ms_mean, enc_coreml_ms_std = _time_coreml(
-        enc, {"mel": _np32(mel_ref), "mel_length": _np32(mel_length_ref).astype(np.int32)}
-    )
-    enc_coreml_rtf = float(enc_coreml_ms_mean / (seconds * 1000.0)) if enc_coreml_ms_mean > 0 else None
-    enc_torch_rtf = float(enc_torch_ms_mean / (seconds * 1000.0)) if enc_torch_ms_mean > 0 else None
-
-    summary["components"]["encoder"] = {
-        "encoder": {"max_abs": a_enc, "max_rel": r_enc, "match": bool(ok_enc)},
-        "length_match": bool(ok_enc_len),
-        "latency": {
-            "runs": int(runs),
-            "warmup": int(warmup),
-            "coreml_first_ms": enc_first_ms,
-            "torch_ms": {"mean": enc_torch_ms_mean, "std": enc_torch_ms_std},
-            "coreml_ms": {"mean": enc_coreml_ms_mean, "std": enc_coreml_ms_std},
-            "rtf": {"torch": enc_torch_rtf, "coreml": enc_coreml_rtf},
-        },
-        "plots": {k: v for k, v in enc_plots.items() if v},
-    }
-
-    # Decoder (sequential U=1 rollout)
-    dec = ct.models.MLModel(str(output_dir / "parakeet_decoder.mlpackage"), compute_units=ct.ComputeUnit.CPU_AND_NE)
-
-    zero_state_np = np.zeros((decoder_layers, 1, decoder_hidden), dtype=np.float32)
-    blank_targets_np = np.array(blank_targets.detach().cpu().numpy(), dtype=np.int32, copy=True)
-    blank_target_lengths_np = np.array(blank_target_lengths.detach().cpu().numpy(), dtype=np.int32, copy=True)
-
-    def _decoder_rollout_coreml(num_steps: int) -> Tuple[np.ndarray, np.ndarray, np.ndarray, float]:
-        outputs = []
-        h_np = zero_state_np.copy()
-        c_np = zero_state_np.copy()
-        first_ms: Optional[float] = None
-        for i in range(num_steps):
-            t0_i = time.perf_counter() if i == 0 else None
-            res = dec.predict(
-                {
-                    "targets": blank_targets_np,
-                    "target_length": blank_target_lengths_np,
-                    "h_in": h_np,
-                    "c_in": c_np,
-                }
-            )
-            if t0_i is not None:
-                t1_i = time.perf_counter()
-                first_ms = (t1_i - t0_i) * 1000.0
-            outputs.append(np.array(res["decoder"], dtype=np.float32, copy=True))
-            h_np = np.array(res["h_out"], dtype=np.float32, copy=True)
-            c_np = np.array(res["c_out"], dtype=np.float32, copy=True)
-        if outputs:
-            decoder_seq = np.concatenate(outputs, axis=-1)
-        else:
-            decoder_seq = np.zeros((1, decoder_hidden, 0), dtype=np.float32)
-        return decoder_seq, h_np, c_np, (0.0 if first_ms is None else float(first_ms))
-
-    dec_ml, h_ml, c_ml, dec_first_ms = _decoder_rollout_coreml(symbol_steps)
-    h_ref_np = _np32(h_ref)
-    c_ref_np = _np32(c_ref)
-
-    a_dec, r_dec, ok_dec = _max_diffs(decoder_ref_np, dec_ml, max(rtol, 1e-2), max(atol, 1e-1))
-    a_h, r_h, ok_h = _max_diffs(h_ref_np, h_ml, max(rtol, 1e-2), max(atol, 2.5e-1))
-    a_c, r_c, ok_c = _max_diffs(c_ref_np, c_ml, max(rtol, 5e-2), max(atol, 1.5e0))
-
-    dec_t = decoder_ref_np[0]
-    dec_c = dec_ml[0]
-    dec_plots = {
-        "decoder_steps_l2.png": _plot_line(
-            np.linalg.norm(dec_t, axis=0),
-            np.linalg.norm(dec_c, axis=0),
-            "Decoder L2 over steps",
-            plots_dir / "decoder_steps_l2.png",
-            also_delta=True,
-        ),
-    }
-
-    def _time_decoder_coreml() -> Tuple[float, float]:
-        for _ in range(max(0, warmup)):
-            _decoder_rollout_coreml(symbol_steps)
-        times = []
-        for _ in range(max(1, runs)):
-            t0 = time.perf_counter()
-            _decoder_rollout_coreml(symbol_steps)
-            t1 = time.perf_counter()
-            times.append((t1 - t0) * 1000.0)
-        arr = np.array(times, dtype=np.float64)
-        return float(arr.mean()), float(arr.std(ddof=1) if arr.size > 1 else 0.0)
-
-    dec_torch_ms_mean, dec_torch_ms_std = _time_torch(lambda: _decoder_rollout_torch(symbol_steps))
-    dec_coreml_ms_mean, dec_coreml_ms_std = _time_decoder_coreml()
-    dec_coreml_rtf = float(dec_coreml_ms_mean / (seconds * 1000.0)) if dec_coreml_ms_mean > 0 else None
-    dec_torch_rtf = float(dec_torch_ms_mean / (seconds * 1000.0)) if dec_torch_ms_mean > 0 else None
-
-    summary["components"]["decoder"] = {
-        "decoder": {"max_abs": a_dec, "max_rel": r_dec, "match": bool(ok_dec)},
-        "h_out": {"max_abs": a_h, "max_rel": r_h, "match": bool(ok_h)},
-        "c_out": {"max_abs": a_c, "max_rel": r_c, "match": bool(ok_c)},
-        "latency": {
-            "runs": int(runs),
-            "warmup": int(warmup),
-            "coreml_first_ms": dec_first_ms,
-            "torch_ms": {"mean": dec_torch_ms_mean, "std": dec_torch_ms_std},
-            "coreml_ms": {"mean": dec_coreml_ms_mean, "std": dec_coreml_ms_std},
-            "rtf": {"torch": dec_torch_rtf, "coreml": dec_coreml_rtf},
-        },
-        "plots": {k: v for k, v in dec_plots.items() if v},
-    }
-
-    # Joint (sequential U=1 rollouts)
-    j = ct.models.MLModel(str(output_dir / "parakeet_joint.mlpackage"), compute_units=ct.ComputeUnit.CPU_AND_NE)
-
-    def _joint_rollout_coreml(decoder_seq_np: np.ndarray) -> Tuple[np.ndarray, float]:
-        logits_steps = []
-        first_ms: Optional[float] = None
-        for u in range(decoder_seq_np.shape[2]):
-            dec_slice = decoder_seq_np[:, :, u : u + 1]
-            t0_u = time.perf_counter() if u == 0 else None
-            res = j.predict({"encoder": encoder_np, "decoder": dec_slice})
-            if t0_u is not None:
-                t1_u = time.perf_counter()
-                first_ms = (t1_u - t0_u) * 1000.0
-            logits_steps.append(np.array(res["logits"], dtype=np.float32, copy=True))
-        if not logits_steps:
-            raise RuntimeError("No decoder steps provided for joint rollout")
-        return np.concatenate(logits_steps, axis=2), (0.0 if first_ms is None else float(first_ms))
-
-    logits_ml, joint_first_ms = _joint_rollout_coreml(decoder_ref_np)
-    logits_ref_np = _np32(logits_ref)
-    a_j, r_j, ok_j = _max_diffs(logits_ref_np, logits_ml, max(rtol, 1e-2), max(atol, 1e-1))
-    joint_plots = {}
-    if HAS_MPL:
-        lt = logits_ref_np[0, 0, 0, :]
-        lc = logits_ml[0, 0, 0, :]
-        top_idx = np.argsort(-np.abs(lt))[:50]
-        path = plots_dir / "joint_top50.png"
-        plt.figure(figsize=(8, 3))
-        plt.plot(lt[top_idx], label="torch")
-        plt.plot(lc[top_idx], label="coreml", alpha=0.8)
-        plt.title("Joint logits (t=0,u=0) top-50 |torch|")
-        plt.legend(); plt.tight_layout(); plt.savefig(path); plt.close()
-        joint_plots["joint_top50.png"] = str(path.name)
-
-        # Delta-over-time visualization (fix u=0; summarize over vocab)
-        jt = logits_ref_np[0, :, 0, :]
-        jc = logits_ml[0, :, 0, :]
-        l2_t = np.linalg.norm(jt, axis=1)
-        l2_c = np.linalg.norm(jc, axis=1)
-        path2 = plots_dir / "joint_time_l2.png"
-        _plot_line(l2_t, l2_c, "Joint L2 over time (u=0)", path2, also_delta=True)
-        joint_plots["joint_time_l2.png"] = str(path2.name)
-
-    def _time_joint_coreml() -> Tuple[float, float]:
-        for _ in range(max(0, warmup)):
-            _joint_rollout_coreml(decoder_ref_np)
-        times = []
-        for _ in range(max(1, runs)):
-            t0 = time.perf_counter()
-            _joint_rollout_coreml(decoder_ref_np)
-            t1 = time.perf_counter()
-            times.append((t1 - t0) * 1000.0)
-        arr = np.array(times, dtype=np.float64)
-        return float(arr.mean()), float(arr.std(ddof=1) if arr.size > 1 else 0.0)
-
-    joint_torch_ms_mean, joint_torch_ms_std = _time_torch(
-        asr_model.joint, encoder_outputs=encoder_ref, decoder_outputs=decoder_ref
-    )
-    joint_coreml_ms_mean, joint_coreml_ms_std = _time_joint_coreml()
-    joint_coreml_rtf = float(joint_coreml_ms_mean / (seconds * 1000.0)) if joint_coreml_ms_mean > 0 else None
-    joint_torch_rtf = float(joint_torch_ms_mean / (seconds * 1000.0)) if joint_torch_ms_mean > 0 else None
-
-    summary["components"]["joint"] = {
-        "logits": {"max_abs": a_j, "max_rel": r_j, "match": bool(ok_j)},
-        "latency": {
-            "runs": int(runs),
-            "warmup": int(warmup),
-            "coreml_first_ms": joint_first_ms,
-            "torch_ms": {"mean": joint_torch_ms_mean, "std": joint_torch_ms_std},
-            "coreml_ms": {"mean": joint_coreml_ms_mean, "std": joint_coreml_ms_std},
-            "rtf": {"torch": joint_torch_rtf, "coreml": joint_coreml_rtf},
-        },
-        "plots": joint_plots,
-    }
-
-    # Fused components
-    # 1) Mel+Encoder fused vs separate
-    mel_enc_plots = {}
-    try:
-        mel_enc = ct.models.MLModel(str(output_dir / "parakeet_mel_encoder.mlpackage"), compute_units=ct.ComputeUnit.CPU_AND_NE)
-        t0 = time.perf_counter()
-        mel_enc_out = mel_enc.predict({
-            "audio_signal": _np32(audio_tensor),
-            "audio_length": _np32(audio_length).astype(np.int32),
-        })
-        t1 = time.perf_counter()
-        mel_enc_first_ms = (t1 - t0) * 1000.0
-        mel_enc_ml = np.array(mel_enc_out["encoder"], dtype=np.float32, copy=True)
-        mel_enc_len_ml = np.array(mel_enc_out["encoder_length"], dtype=np.int32, copy=True)
-        # Compare fused output vs Torch reference encoder
-        a_melenc, r_melenc, ok_melenc = _max_diffs(_np32(encoder_ref), mel_enc_ml, max(rtol, 5e-3), max(atol, 5e-2))
-        ok_melenc_len = int(_np32(encoder_length_ref).astype(np.int32)[0]) == int(mel_enc_len_ml.astype(np.int32)[0])
-        # Also compare fused vs separate CoreML pipeline (pre -> enc)
-        a_melenc_vs_sep, r_melenc_vs_sep, ok_melenc_vs_sep = _max_diffs(enc_ml, mel_enc_ml, max(rtol, 5e-3), max(atol, 5e-2))
-
-        # Plots: L2 over time (fused vs torch)
-        enc_t_ref = _np32(encoder_ref)[0]
-        enc_c_fused = mel_enc_ml[0]
-        mel_enc_plots["mel_encoder_time_l2.png"] = _plot_line(
-            np.linalg.norm(enc_t_ref, axis=0),
-            np.linalg.norm(enc_c_fused, axis=0),
-            "Mel+Encoder (fused) L2 over time",
-            plots_dir / "mel_encoder_time_l2.png",
-            also_delta=True,
-        )
-
-        # Latency: fused CoreML vs separate (CoreML pre + CoreML enc)
-        mel_enc_coreml_ms_mean, mel_enc_coreml_ms_std = _time_coreml(
-            mel_enc,
-            {"audio_signal": _np32(audio_tensor), "audio_length": _np32(audio_length).astype(np.int32)},
-        )
-        sep_coreml_ms_mean = float(pre_coreml_ms_mean + enc_coreml_ms_mean)
-        sep_coreml_ms_std = float((pre_coreml_ms_std ** 2 + enc_coreml_ms_std ** 2) ** 0.5)
-        # Torch baseline (separate torch pre + enc)
-        sep_torch_ms_mean = float(pre_torch_ms_mean + enc_torch_ms_mean)
-        sep_torch_ms_std = float((pre_torch_ms_std ** 2 + enc_torch_ms_std ** 2) ** 0.5)
-
-        mel_enc_coreml_rtf = float(mel_enc_coreml_ms_mean / (seconds * 1000.0)) if mel_enc_coreml_ms_mean > 0 else None
-        sep_coreml_rtf = float(sep_coreml_ms_mean / (seconds * 1000.0)) if sep_coreml_ms_mean > 0 else None
-        sep_torch_rtf = float(sep_torch_ms_mean / (seconds * 1000.0)) if sep_torch_ms_mean > 0 else None
-
-        summary["components"]["mel_encoder"] = {
-            "encoder": {"max_abs": a_melenc, "max_rel": r_melenc, "match": bool(ok_melenc)},
-            "length_match": bool(ok_melenc_len),
-            "vs_separate_coreml": {"max_abs": a_melenc_vs_sep, "max_rel": r_melenc_vs_sep, "match": bool(ok_melenc_vs_sep)},
-            "latency": {
-                "runs": int(runs),
-                "warmup": int(warmup),
-                "fused_coreml_first_ms": mel_enc_first_ms,
-                "fused_coreml_ms": {"mean": mel_enc_coreml_ms_mean, "std": mel_enc_coreml_ms_std},
-                "separate_coreml_ms": {"mean": sep_coreml_ms_mean, "std": sep_coreml_ms_std},
-                "separate_torch_ms": {"mean": sep_torch_ms_mean, "std": sep_torch_ms_std},
-                "rtf": {"fused_coreml": mel_enc_coreml_rtf, "separate_coreml": sep_coreml_rtf, "separate_torch": sep_torch_rtf},
-            },
-            "plots": {k: v for k, v in mel_enc_plots.items() if v},
-        }
-    except Exception as e:
-        summary["components"]["mel_encoder_error"] = str(e)
-
-    # 2) JointDecision fused vs CPU PyTorch post-processing
-    jd_plots = {}
-    try:
-        # Fused CoreML joint decision
-        jd = ct.models.MLModel(str(output_dir / "parakeet_joint_decision.mlpackage"), compute_units=ct.ComputeUnit.CPU_AND_NE)
-        def _joint_decision_rollout_coreml(decoder_seq_np: np.ndarray) -> Tuple[np.ndarray, np.ndarray, np.ndarray, float]:
-            token_ids = []
-            token_probs = []
-            durations = []
-            first_ms: Optional[float] = None
-            for u in range(decoder_seq_np.shape[2]):
-                dec_slice = decoder_seq_np[:, :, u : u + 1]
-                t0_u = time.perf_counter() if u == 0 else None
-                res = jd.predict({"encoder": encoder_np, "decoder": dec_slice})
-                if t0_u is not None:
-                    t1_u = time.perf_counter()
-                    first_ms = (t1_u - t0_u) * 1000.0
-                token_ids.append(np.array(res["token_id"], dtype=np.int32, copy=True))
-                token_probs.append(np.array(res["token_prob"], dtype=np.float32, copy=True))
-                durations.append(np.array(res["duration"], dtype=np.int32, copy=True))
-            if not token_ids:
-                raise RuntimeError("No decoder steps provided for joint decision rollout")
-            return (
-                np.concatenate(token_ids, axis=2),
-                np.concatenate(token_probs, axis=2),
-                np.concatenate(durations, axis=2),
-                (0.0 if first_ms is None else float(first_ms)),
-            )
-
-        token_id_ml, token_prob_ml, duration_ml, jd_first_ms = _joint_decision_rollout_coreml(decoder_ref_np)
-
-        # CPU PyTorch decision using Torch logits
-        vocab_with_blank = int(vocab_size) + 1
-        with torch.inference_mode():
-            logits_t = logits_ref
-            token_logits_t = logits_t[..., :vocab_with_blank]
-            duration_logits_t = logits_t[..., -num_extra:] if num_extra > 0 else None
-            token_ids_t = torch.argmax(token_logits_t, dim=-1).to(dtype=torch.int32)
-            token_probs_all_t = torch.softmax(token_logits_t, dim=-1)
-            token_prob_t = torch.gather(
-                token_probs_all_t, dim=-1, index=token_ids_t.long().unsqueeze(-1)
-            ).squeeze(-1)
-            if duration_logits_t is not None and duration_logits_t.numel() > 0:
-                duration_t = torch.argmax(duration_logits_t, dim=-1).to(dtype=torch.int32)
-            else:
-                duration_t = torch.zeros_like(token_ids_t, dtype=torch.int32)
-
-        # Also derive CPU decision from CoreML joint logits for "separate" path
-        token_logits_c = _to_t(logits_ml)[..., :vocab_with_blank]
-        duration_logits_c = _to_t(logits_ml)[..., -num_extra:] if num_extra > 0 else None
-        token_ids_c = torch.argmax(token_logits_c, dim=-1).to(dtype=torch.int32)
-        token_probs_all_c = torch.softmax(token_logits_c, dim=-1)
-        token_prob_c = torch.gather(
-            token_probs_all_c, dim=-1, index=token_ids_c.long().unsqueeze(-1)
-        ).squeeze(-1)
-        if duration_logits_c is not None and duration_logits_c.numel() > 0:
-            duration_c = torch.argmax(duration_logits_c, dim=-1).to(dtype=torch.int32)
-        else:
-            duration_c = torch.zeros_like(token_ids_c, dtype=torch.int32)
-
-        # Compare fused outputs to CPU PyTorch decisions
-        a_tid_t, r_tid_t, ok_tid_t = _max_diffs(_np(token_ids_t), token_id_ml, 0.0, 0.0)
-        a_tprob_t, r_tprob_t, ok_tprob_t = _max_diffs(_np(token_prob_t), token_prob_ml, max(rtol, 1e-2), max(atol, 1e-1))
-        a_dur_t, r_dur_t, ok_dur_t = _max_diffs(_np(duration_t), duration_ml, 0.0, 0.0)
-
-        a_tid_c, r_tid_c, ok_tid_c = _max_diffs(_np(token_ids_c), token_id_ml, 0.0, 0.0)
-        a_tprob_c, r_tprob_c, ok_tprob_c = _max_diffs(_np(token_prob_c), token_prob_ml, max(rtol, 1e-2), max(atol, 1e-1))
-        a_dur_c, r_dur_c, ok_dur_c = _max_diffs(_np(duration_c), duration_ml, 0.0, 0.0)
-
-        # Plots: token_prob over time for u=0 (fused vs torch CPU)
-        if HAS_MPL:
-            prob_t = _np(token_prob_t)[0, :, 0]
-            prob_ml = token_prob_ml[0, :, 0]
-            jd_plots["joint_decision_prob_u0.png"] = _plot_line(
-                prob_t,
-                prob_ml,
-                "JointDecision token_prob (u=0)",
-                plots_dir / "joint_decision_prob_u0.png",
-                also_delta=True,
-            )
-
-            # Agreement heatmap for token_id
-            agree = (_np(token_ids_t)[0] == token_id_ml[0]).astype(np.float32)
-            jd_plots["joint_decision_token_agree.png"] = _plot_image(
-                agree,
-                "token_id agreement (torch CPU vs fused)",
-                plots_dir / "joint_decision_token_agree.png",
-                vmin=0.0,
-                vmax=1.0,
-            )
-
-        # Latency: fused CoreML vs separate (CoreML joint + CPU PyTorch decision)
-        def _time_joint_decision_coreml() -> Tuple[float, float]:
-            for _ in range(max(0, warmup)):
-                _joint_decision_rollout_coreml(decoder_ref_np)
-            times = []
-            for _ in range(max(1, runs)):
-                t0 = time.perf_counter()
-                _joint_decision_rollout_coreml(decoder_ref_np)
-                t1 = time.perf_counter()
-                times.append((t1 - t0) * 1000.0)
-            arr = np.array(times, dtype=np.float64)
-            return float(arr.mean()), float(arr.std(ddof=1) if arr.size > 1 else 0.0)
-
-        jd_coreml_ms_mean, jd_coreml_ms_std = _time_joint_decision_coreml()
-
-        # Time CPU post-processing only (Torch) on top of CoreML or Torch logits. Use Torch logits.
-        def _decision_torch_call():
-            with torch.inference_mode():
-                tl = logits_ref
-                tl_token = tl[..., :vocab_with_blank]
-                tl_ids = torch.argmax(tl_token, dim=-1)
-                tl_probs = torch.softmax(tl_token, dim=-1)
-                _ = torch.gather(tl_probs, -1, tl_ids.long().unsqueeze(-1)).squeeze(-1)
-                if num_extra > 0:
-                    _ = torch.argmax(tl[..., -num_extra:], dim=-1)
-            return None
-
-        jd_decision_torch_ms_mean, jd_decision_torch_ms_std = _time_torch(lambda: _decision_torch_call())
-        sep_joint_plus_cpu_ms_mean = float(joint_coreml_ms_mean + jd_decision_torch_ms_mean)
-        sep_joint_plus_cpu_ms_std = float((joint_coreml_ms_std ** 2 + jd_decision_torch_ms_std ** 2) ** 0.5)
-        jd_coreml_rtf = float(jd_coreml_ms_mean / (seconds * 1000.0)) if jd_coreml_ms_mean > 0 else None
-        sep_joint_cpu_rtf = float(sep_joint_plus_cpu_ms_mean / (seconds * 1000.0)) if sep_joint_plus_cpu_ms_mean > 0 else None
-
-        summary["components"]["joint_decision"] = {
-            "vs_torch_cpu": {
-                "token_id": {"max_abs": a_tid_t, "max_rel": r_tid_t, "match": bool(ok_tid_t)},
-                "token_prob": {"max_abs": a_tprob_t, "max_rel": r_tprob_t, "match": bool(ok_tprob_t)},
-                "duration": {"max_abs": a_dur_t, "max_rel": r_dur_t, "match": bool(ok_dur_t)},
-            },
-            "vs_coreml_joint_cpu": {
-                "token_id": {"max_abs": a_tid_c, "max_rel": r_tid_c, "match": bool(ok_tid_c)},
-                "token_prob": {"max_abs": a_tprob_c, "max_rel": r_tprob_c, "match": bool(ok_tprob_c)},
-                "duration": {"max_abs": a_dur_c, "max_rel": r_dur_c, "match": bool(ok_dur_c)},
-            },
-            "latency": {
-                "runs": int(runs),
-                "warmup": int(warmup),
-                "fused_coreml_first_ms": jd_first_ms,
-                "fused_coreml_ms": {"mean": jd_coreml_ms_mean, "std": jd_coreml_ms_std},
-                "separate_joint_coreml_plus_cpu_ms": {"mean": sep_joint_plus_cpu_ms_mean, "std": sep_joint_plus_cpu_ms_std},
-                "rtf": {"fused_coreml": jd_coreml_rtf, "separate_joint_coreml_plus_cpu": sep_joint_cpu_rtf},
-            },
-            "plots": {k: v for k, v in jd_plots.items() if v},
-        }
-    except Exception as e:
-        summary["components"]["joint_decision_error"] = str(e)
-
-    # Latency overview plots (saved alongside component plots)
-    latency_plots = {}
-    labels = ["preprocessor", "encoder", "decoder", "joint"]
-    torch_means = [pre_torch_ms_mean, enc_torch_ms_mean, dec_torch_ms_mean, joint_torch_ms_mean]
-    torch_stds = [pre_torch_ms_std, enc_torch_ms_std, dec_torch_ms_std, joint_torch_ms_std]
-    coreml_means = [pre_coreml_ms_mean, enc_coreml_ms_mean, dec_coreml_ms_mean, joint_coreml_ms_mean]
-    coreml_stds = [pre_coreml_ms_std, enc_coreml_ms_std, dec_coreml_ms_std, joint_coreml_ms_std]
-    lat_path = plots_dir / "latency_summary.png"
-    spd_path = plots_dir / "latency_speedup.png"
-    latency_plots["latency_summary.png"] = _plot_latency_bars(
-        labels, torch_means, torch_stds, coreml_means, coreml_stds, lat_path
-    )
-    latency_plots["latency_speedup.png"] = _plot_speedup_bars(
-        labels, torch_means, coreml_means, spd_path
-    )
-
-    # Fused vs separate latency summary
-    fused_labels = ["mel+encoder", "joint_decision"]
-    fused_baseline_means = [
-        float(pre_torch_ms_mean + enc_torch_ms_mean),
-        float(joint_coreml_ms_mean + jd_decision_torch_ms_mean if 'jd_coreml_ms_mean' in locals() else joint_coreml_ms_mean),
-    ]
-    fused_coreml_means = [
-        float(mel_enc_coreml_ms_mean if 'mel_enc_coreml_ms_mean' in locals() else np.nan),
-        float(jd_coreml_ms_mean if 'jd_coreml_ms_mean' in locals() else np.nan),
-    ]
-    fused_latency_path = plots_dir / "latency_fused_vs_separate.png"
-    fused_speedup_path = plots_dir / "latency_fused_speedup.png"
-    latency_plots["latency_fused_vs_separate.png"] = _plot_latency_bars(
-        fused_labels, fused_baseline_means, [0, 0], fused_coreml_means, [0, 0], fused_latency_path
-    )
-    latency_plots["latency_fused_speedup.png"] = _plot_speedup_bars(
-        fused_labels, fused_baseline_means, fused_coreml_means, fused_speedup_path
-    )
-
-    all_ok = (
-        summary["components"]["preprocessor"]["mel"]["match"]
-        and summary["components"]["preprocessor"]["length_match"]
-        and summary["components"]["encoder"]["encoder"]["match"]
-        and summary["components"]["encoder"]["length_match"]
-        and summary["components"]["decoder"]["decoder"]["match"]
-        and summary["components"]["decoder"]["h_out"]["match"]
-        and summary["components"]["decoder"]["c_out"]["match"]
-        and summary["components"]["joint"]["logits"]["match"]
-    )
-    summary["status"] = "ok" if all_ok else "mismatch"
-
-    # Update metadata.json
-    meta_path = output_dir / "metadata.json"
-    try:
-        meta = json.loads(meta_path.read_text())
-    except Exception:
-        meta = {}
-    meta["validation"] = summary
-    meta_path.write_text(json.dumps(meta, indent=2))
-
-    typer.echo(f"Validation {'passed' if all_ok else 'mismatched'}. Updated {meta_path}")
-    if HAS_MPL:
-        typer.echo(f"Saved plots to {plots_dir}")
-
-
-if __name__ == "__main__":
-    app()

convert/parakeet-tdt-v2-0.6b/coreml/compile_modelc.py

@@ -1,91 +0,0 @@
-#!/usr/bin/env python3
-"""Compile Core ML packages into ``.mlmodelc`` bundles via ``xcrun``.
-
-This script walks through the default Parakeet CoreML directories, finds
-all ``*.mlpackage`` bundles, and compiles each of them with
-``xcrun coremlcompiler`` into ``./compiled`` while preserving the
-relative directory structure.
-"""
-from __future__ import annotations
-
-import shutil
-import subprocess
-import sys
-from pathlib import Path
-
-BASE_DIR = Path(__file__).resolve().parent
-OUTPUT_ROOT = BASE_DIR / "compiled"
-SOURCE_DIRS = [BASE_DIR / "parakeet_coreml", BASE_DIR / "parakeet_coreml_quantized"]
-
-
-def ensure_coremlcompiler() -> None:
-    """Ensure ``xcrun coremlcompiler`` is available for the active Xcode."""
-    xcrun_path = shutil.which("xcrun")
-    if xcrun_path is None:
-        print("Error: 'xcrun' not found on PATH. Install Xcode command line tools.", file=sys.stderr)
-        sys.exit(1)
-
-    try:
-        subprocess.run([
-            xcrun_path,
-            "--find",
-            "coremlcompiler",
-        ], check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
-    except subprocess.CalledProcessError:
-        print("Error: 'coremlcompiler' not found via xcrun. Check your Xcode installation.", file=sys.stderr)
-        sys.exit(1)
-
-
-def gather_packages() -> list[Path]:
-    """Return a list of all ``*.mlpackage`` bundles under the source dirs."""
-    packages: list[Path] = []
-    for source in SOURCE_DIRS:
-        if not source.exists():
-            print(f"Warning: {source.relative_to(BASE_DIR)} does not exist; skipping", file=sys.stderr)
-            continue
-        packages.extend(source.rglob("*.mlpackage"))
-    return packages
-
-
-def compile_package(package: Path) -> None:
-    """Compile a single ``.mlpackage`` bundle using ``xcrun coremlcompiler``."""
-    relative_pkg = package.relative_to(BASE_DIR)
-    output_dir = OUTPUT_ROOT / relative_pkg.parent
-    output_dir.mkdir(parents=True, exist_ok=True)
-    output_path = output_dir / f"{package.stem}.mlmodelc"
-
-    if output_path.exists():
-        shutil.rmtree(output_path)
-
-    cmd = [
-        "xcrun",
-        "coremlcompiler",
-        "compile",
-        str(package),
-        str(output_dir),
-    ]
-
-    print(f"Compiling {relative_pkg} -> {output_path.relative_to(BASE_DIR)}")
-    subprocess.run(cmd, check=True)
-
-
-def main() -> None:
-    ensure_coremlcompiler()
-    packages = gather_packages()
-
-    if not packages:
-        print("No .mlpackage bundles found to compile.")
-        return
-
-    for package in packages:
-        try:
-            compile_package(package)
-        except subprocess.CalledProcessError as exc:
-            print(f"Failed to compile {package}: {exc}", file=sys.stderr)
-            sys.exit(exc.returncode)
-
-    print(f"Finished compiling {len(packages)} package(s) into {OUTPUT_ROOT.relative_to(BASE_DIR)}.")
-
-
-if __name__ == "__main__":
-    main()

convert/parakeet-tdt-v2-0.6b/coreml/convert-parakeet.py

@@ -1,619 +0,0 @@
-#!/usr/bin/env python3
-"""CLI for exporting Parakeet TDT v2 components to CoreML."""
-from __future__ import annotations
-
-import json
-from dataclasses import asdict
-from pathlib import Path
-from typing import Dict, Optional, Tuple
-
-import coremltools as ct
-import numpy as np
-import soundfile as sf
-import torch
-import typer
-
-import nemo.collections.asr as nemo_asr
-
-from individual_components import (
-    DecoderWrapper,
-    EncoderWrapper,
-    ExportSettings,
-    JointWrapper,
-    JointDecisionWrapper,
-    JointDecisionSingleStep,
-    PreprocessorWrapper,
-    MelEncoderWrapper,
-    _coreml_convert,
-)
-
-DEFAULT_MODEL_ID = "nvidia/parakeet-tdt-0.6b-v2"
-AUTHOR = "Fluid Inference"
-
-
-def _compute_length(seconds: float, sample_rate: int) -> int:
-    return int(round(seconds * sample_rate))
-
-
-def _prepare_audio(
-    validation_audio: Optional[Path],
-    sample_rate: int,
-    max_samples: int,
-    seed: Optional[int],
-) -> torch.Tensor:
-    if validation_audio is None:
-        if seed is not None:
-            torch.manual_seed(seed)
-        audio = torch.randn(1, max_samples, dtype=torch.float32)
-        return audio
-
-    data, sr = sf.read(str(validation_audio), dtype="float32")
-    if sr != sample_rate:
-        raise typer.BadParameter(
-            f"Validation audio sample rate {sr} does not match model rate {sample_rate}"
-        )
-
-    if data.ndim > 1:
-        data = data[:, 0]
-
-    if data.size == 0:
-        raise typer.BadParameter("Validation audio is empty")
-
-    if data.size < max_samples:
-        pad_width = max_samples - data.size
-        data = np.pad(data, (0, pad_width))
-    elif data.size > max_samples:
-        data = data[:max_samples]
-
-    audio = torch.from_numpy(data).unsqueeze(0).to(dtype=torch.float32)
-    return audio
-
-
-def _save_mlpackage(model: ct.models.MLModel, path: Path, description: str) -> None:
-    # Ensure iOS 17+ target for MLProgram ops and ANE readiness
-    try:
-        model.minimum_deployment_target = ct.target.iOS17
-    except Exception:
-        pass
-    model.short_description = description
-    model.author = AUTHOR
-    path.parent.mkdir(parents=True, exist_ok=True)
-    model.save(str(path))
-
-
-def _tensor_shape(tensor: torch.Tensor) -> Tuple[int, ...]:
-    return tuple(int(dim) for dim in tensor.shape)
-
-
-def _parse_compute_units(name: str) -> ct.ComputeUnit:
-    """Parse a human-friendly compute units string into ct.ComputeUnit.
-
-    Accepted (case-insensitive): ALL, CPU_ONLY, CPU_AND_GPU, CPU_AND_NE.
-    """
-    normalized = str(name).strip().upper()
-    mapping = {
-        "ALL": ct.ComputeUnit.ALL,
-        "CPU_ONLY": ct.ComputeUnit.CPU_ONLY,
-        "CPU_AND_GPU": ct.ComputeUnit.CPU_AND_GPU,
-        "CPU_AND_NE": ct.ComputeUnit.CPU_AND_NE,
-        "CPU_AND_NEURALENGINE": ct.ComputeUnit.CPU_AND_NE,
-    }
-    if normalized not in mapping:
-        raise typer.BadParameter(
-            f"Unknown compute units '{name}'. Choose from: " + ", ".join(mapping.keys())
-        )
-    return mapping[normalized]
-
-
-def _parse_compute_precision(name: Optional[str]) -> Optional[ct.precision]:
-    """Parse compute precision string into ct.precision or None.
-
-    Accepted (case-insensitive): FLOAT32, FLOAT16. If None/empty, returns None (tool default).
-    """
-    if name is None:
-        return None
-    normalized = str(name).strip().upper()
-    if normalized == "":
-        return None
-    mapping = {
-        "FLOAT32": ct.precision.FLOAT32,
-        "FLOAT16": ct.precision.FLOAT16,
-    }
-    if normalized not in mapping:
-        raise typer.BadParameter(
-            f"Unknown compute precision '{name}'. Choose from: " + ", ".join(mapping.keys())
-        )
-    return mapping[normalized]
-
-
-# Validation logic removed; use compare-compnents.py for comparisons.
-
-
-# Fixed export choices: CPU_ONLY + FP32, min target iOS17
-
-
-app = typer.Typer(add_completion=False, pretty_exceptions_show_locals=False)
-
-
-@app.command()
-def convert(
-    nemo_path: Optional[Path] = typer.Option(
-        None,
-        "--nemo-path",
-        exists=True,
-        resolve_path=True,
-        help="Path to parakeet-tdt-0.6b-v2 .nemo checkpoint (skip to auto-download)",
-    ),
-    model_id: str = typer.Option(
-        DEFAULT_MODEL_ID,
-        "--model-id",
-        help="Model identifier to download when --nemo-path is omitted",
-    ),
-    output_dir: Path = typer.Option(Path("parakeet_coreml"), help="Directory where mlpackages and metadata will be written"),
-    preprocessor_cu: str = typer.Option(
-        "CPU_ONLY",
-        "--preprocessor-cu",
-        help="Compute units for preprocessor (default CPU_ONLY)",
-    ),
-    mel_encoder_cu: str = typer.Option(
-        "CPU_ONLY",
-        "--mel-encoder-cu",
-        help="Compute units for fused mel+encoder (default CPU_ONLY)",
-    ),
-    compute_precision: Optional[str] = typer.Option(
-        None,
-        "--compute-precision",
-        help="Export precision: FLOAT32 (default) or FLOAT16 to shrink non-quantized weights.",
-    ),
-) -> None:
-    """Export all Parakeet sub-modules to CoreML with a fixed 15-second window."""
-    # Runtime CoreML contract keeps U=1 so the prediction net matches the streaming decoder.
-    export_settings = ExportSettings(
-        output_dir=output_dir,
-        compute_units=ct.ComputeUnit.CPU_ONLY,  # Default: CPU-only for all components
-        deployment_target=ct.target.iOS17,  # iOS 17+ features and kernels
-        compute_precision=_parse_compute_precision(compute_precision),
-        max_audio_seconds=15.0,
-        max_symbol_steps=1,
-    )
-
-    typer.echo("Export configuration:")
-    typer.echo(asdict(export_settings))
-
-    output_dir.mkdir(parents=True, exist_ok=True)
-    pre_cu = _parse_compute_units(preprocessor_cu)
-    melenc_cu = _parse_compute_units(mel_encoder_cu)
-
-    if nemo_path is not None:
-        typer.echo(f"Loading NeMo model from {nemo_path}…")
-        asr_model = nemo_asr.models.EncDecRNNTBPEModel.restore_from(
-            str(nemo_path), map_location="cpu"
-        )
-        checkpoint_meta = {
-            "type": "file",
-            "path": str(nemo_path),
-        }
-    else:
-        typer.echo(f"Downloading NeMo model via {model_id}…")
-        asr_model = nemo_asr.models.EncDecRNNTBPEModel.from_pretrained(
-            model_id, map_location="cpu"
-        )
-        checkpoint_meta = {
-            "type": "pretrained",
-            "model_id": model_id,
-        }
-    asr_model.eval()
-
-    sample_rate = int(asr_model.cfg.preprocessor.sample_rate)
-    max_samples = _compute_length(export_settings.max_audio_seconds, sample_rate)
-    # Prefer a bundled 15s 16kHz audio if available
-    default_audio = (Path(__file__).parent / "audio" / "yc_first_minute_16k_15s.wav").resolve()
-    if not default_audio.exists():
-        raise typer.BadParameter(f"Expected 15s trace audio at {default_audio}; add the file to proceed.")
-    typer.echo(f"Using trace audio: {default_audio}")
-    audio_tensor = _prepare_audio(default_audio, sample_rate, max_samples, seed=None)
-    audio_length = torch.tensor([max_samples], dtype=torch.int32)
-
-    preprocessor = PreprocessorWrapper(asr_model.preprocessor.eval())
-    encoder = EncoderWrapper(asr_model.encoder.eval())
-    decoder = DecoderWrapper(asr_model.decoder.eval())
-    joint = JointWrapper(asr_model.joint.eval())
-
-    decoder_export_flag = getattr(asr_model.decoder, "_rnnt_export", False)
-    asr_model.decoder._rnnt_export = True
-
-    try:
-        with torch.inference_mode():
-            mel_ref, mel_length_ref = preprocessor(audio_tensor, audio_length)
-            mel_length_ref = mel_length_ref.to(dtype=torch.int32)
-            encoder_ref, encoder_length_ref = encoder(mel_ref, mel_length_ref)
-            encoder_length_ref = encoder_length_ref.to(dtype=torch.int32)
-
-        # Clone Tensors to drop the inference tensor flag before tracing
-        mel_ref = mel_ref.clone()
-        mel_length_ref = mel_length_ref.clone()
-        encoder_ref = encoder_ref.clone()
-        encoder_length_ref = encoder_length_ref.clone()
-
-        vocab_size = int(asr_model.tokenizer.vocab_size)
-        num_extra = int(asr_model.joint.num_extra_outputs)
-        decoder_hidden = int(asr_model.decoder.pred_hidden)
-        decoder_layers = int(asr_model.decoder.pred_rnn_layers)
-
-        targets = torch.full(
-            (1, export_settings.max_symbol_steps),
-            fill_value=asr_model.decoder.blank_idx,
-            dtype=torch.int32,
-        )
-        target_lengths = torch.tensor(
-            [export_settings.max_symbol_steps], dtype=torch.int32
-        )
-        zero_state = torch.zeros(
-            decoder_layers,
-            1,
-            decoder_hidden,
-            dtype=torch.float32,
-        )
-
-        with torch.inference_mode():
-            decoder_ref, h_ref, c_ref = decoder(targets, target_lengths, zero_state, zero_state)
-            joint_ref = joint(encoder_ref, decoder_ref)
-
-        decoder_ref = decoder_ref.clone()
-        h_ref = h_ref.clone()
-        c_ref = c_ref.clone()
-        joint_ref = joint_ref.clone()
-
-        typer.echo("Tracing and converting preprocessor…")
-        # Ensure tracing happens on CPU explicitly
-        preprocessor = preprocessor.cpu()
-        audio_tensor = audio_tensor.cpu()
-        audio_length = audio_length.cpu()
-        traced_preprocessor = torch.jit.trace(
-            preprocessor, (audio_tensor, audio_length), strict=False
-        )
-        traced_preprocessor.eval()
-        preprocessor_inputs = [
-            # Allow variable-length audio up to the fixed 15s window using RangeDim
-            ct.TensorType(
-                name="audio_signal",
-                shape=(1, ct.RangeDim(1, max_samples)),
-                dtype=np.float32,
-            ),
-            ct.TensorType(name="audio_length", shape=(1,), dtype=np.int32),
-        ]
-        preprocessor_outputs = [
-            ct.TensorType(name="mel", dtype=np.float32),
-            ct.TensorType(name="mel_length", dtype=np.int32),
-        ]
-        # Preprocessor compute units (parametrized; default CPU_ONLY)
-        preprocessor_model = _coreml_convert(
-            traced_preprocessor,
-            preprocessor_inputs,
-            preprocessor_outputs,
-            export_settings,
-            compute_units_override=pre_cu,
-        )
-        preprocessor_path = output_dir / "parakeet_preprocessor.mlpackage"
-        _save_mlpackage(
-            preprocessor_model,
-            preprocessor_path,
-            "Parakeet preprocessor (15 s window)",
-        )
-
-        typer.echo("Tracing and converting encoder…")
-        traced_encoder = torch.jit.trace(
-            encoder, (mel_ref, mel_length_ref), strict=False
-        )
-        traced_encoder.eval()
-        encoder_inputs = [
-            ct.TensorType(name="mel", shape=_tensor_shape(mel_ref), dtype=np.float32),
-            ct.TensorType(name="mel_length", shape=(1,), dtype=np.int32),
-        ]
-        encoder_outputs = [
-            ct.TensorType(name="encoder", dtype=np.float32),
-            ct.TensorType(name="encoder_length", dtype=np.int32),
-        ]
-        # Encoder: CPU only
-        encoder_model = _coreml_convert(
-            traced_encoder,
-            encoder_inputs,
-            encoder_outputs,
-            export_settings,
-            compute_units_override=ct.ComputeUnit.CPU_ONLY,
-        )
-        encoder_path = output_dir / "parakeet_encoder.mlpackage"
-        _save_mlpackage(
-            encoder_model,
-            encoder_path,
-            "Parakeet encoder (15 s window)",
-        )
-
-        # Optional fused export: Preprocessor + Encoder
-        typer.echo("Tracing and converting fused mel+encoder…")
-        mel_encoder = MelEncoderWrapper(preprocessor, encoder)
-        traced_mel_encoder = torch.jit.trace(
-            mel_encoder, (audio_tensor, audio_length), strict=False
-        )
-        traced_mel_encoder.eval()
-        mel_encoder_inputs = [
-            # Keep fixed 15s window for fused Mel+Encoder
-            ct.TensorType(name="audio_signal", shape=(1, max_samples), dtype=np.float32),
-            ct.TensorType(name="audio_length", shape=(1,), dtype=np.int32),
-        ]
-        mel_encoder_outputs = [
-            ct.TensorType(name="encoder", dtype=np.float32),
-            ct.TensorType(name="encoder_length", dtype=np.int32),
-        ]
-        # Fused mel+encoder compute units (parametrized; default CPU_ONLY)
-        mel_encoder_model = _coreml_convert(
-            traced_mel_encoder,
-            mel_encoder_inputs,
-            mel_encoder_outputs,
-            export_settings,
-            compute_units_override=melenc_cu,
-        )
-        mel_encoder_path = output_dir / "parakeet_mel_encoder.mlpackage"
-        _save_mlpackage(
-            mel_encoder_model,
-            mel_encoder_path,
-            "Parakeet fused Mel+Encoder (15 s window)",
-        )
-
-        typer.echo("Tracing and converting decoder…")
-        traced_decoder = torch.jit.trace(
-            decoder,
-            (targets, target_lengths, zero_state, zero_state),
-            strict=False,
-        )
-        traced_decoder.eval()
-        decoder_inputs = [
-            ct.TensorType(name="targets", shape=_tensor_shape(targets), dtype=np.int32),
-            ct.TensorType(name="target_length", shape=(1,), dtype=np.int32),
-            ct.TensorType(name="h_in", shape=_tensor_shape(zero_state), dtype=np.float32),
-            ct.TensorType(name="c_in", shape=_tensor_shape(zero_state), dtype=np.float32),
-        ]
-        decoder_outputs = [
-            ct.TensorType(name="decoder", dtype=np.float32),
-            ct.TensorType(name="h_out", dtype=np.float32),
-            ct.TensorType(name="c_out", dtype=np.float32),
-        ]
-        # Decoder: CPU only
-        decoder_model = _coreml_convert(
-            traced_decoder,
-            decoder_inputs,
-            decoder_outputs,
-            export_settings,
-            compute_units_override=ct.ComputeUnit.CPU_ONLY,
-        )
-        decoder_path = output_dir / "parakeet_decoder.mlpackage"
-        _save_mlpackage(
-            decoder_model,
-            decoder_path,
-            "Parakeet decoder (RNNT prediction network)",
-        )
-
-        typer.echo("Tracing and converting joint…")
-        traced_joint = torch.jit.trace(
-            joint,
-            (encoder_ref, decoder_ref),
-            strict=False,
-        )
-        traced_joint.eval()
-        joint_inputs = [
-            ct.TensorType(name="encoder", shape=_tensor_shape(encoder_ref), dtype=np.float32),
-            ct.TensorType(name="decoder", shape=_tensor_shape(decoder_ref), dtype=np.float32),
-        ]
-        joint_outputs = [
-            ct.TensorType(name="logits", dtype=np.float32),
-        ]
-        # Joint: CPU only
-        joint_model = _coreml_convert(
-            traced_joint,
-            joint_inputs,
-            joint_outputs,
-            export_settings,
-            compute_units_override=ct.ComputeUnit.CPU_ONLY,
-        )
-        joint_path = output_dir / "parakeet_joint.mlpackage"
-        _save_mlpackage(
-            joint_model,
-            joint_path,
-            "Parakeet joint network (RNNT)",
-        )
-
-        # Joint + decision head (split logits, softmax, argmax)
-        typer.echo("Tracing and converting joint decision head…")
-        vocab_size = int(asr_model.tokenizer.vocab_size)
-        num_extra = int(asr_model.joint.num_extra_outputs)
-        joint_decision = JointDecisionWrapper(joint, vocab_size=vocab_size, num_extra=num_extra)
-        traced_joint_decision = torch.jit.trace(
-            joint_decision,
-            (encoder_ref, decoder_ref),
-            strict=False,
-        )
-        traced_joint_decision.eval()
-        joint_decision_inputs = [
-            ct.TensorType(name="encoder", shape=_tensor_shape(encoder_ref), dtype=np.float32),
-            ct.TensorType(name="decoder", shape=_tensor_shape(decoder_ref), dtype=np.float32),
-        ]
-        joint_decision_outputs = [
-            ct.TensorType(name="token_id", dtype=np.int32),
-            ct.TensorType(name="token_prob", dtype=np.float32),
-            ct.TensorType(name="duration", dtype=np.int32),
-        ]
-        # JointDecision: CPU only
-        joint_decision_model = _coreml_convert(
-            traced_joint_decision,
-            joint_decision_inputs,
-            joint_decision_outputs,
-            export_settings,
-            compute_units_override=ct.ComputeUnit.CPU_ONLY,
-        )
-        joint_decision_path = output_dir / "parakeet_joint_decision.mlpackage"
-        _save_mlpackage(
-            joint_decision_model,
-            joint_decision_path,
-            "Parakeet joint + decision head (split, softmax, argmax)",
-        )
-
-        # Single-step JointDecision for [1,1024,1] x [1,640,1] -> [1,1,1]
-        typer.echo("Tracing and converting single-step joint decision…")
-        jd_single = JointDecisionSingleStep(joint, vocab_size=vocab_size, num_extra=num_extra)
-        # Create single-step slices from refs
-        enc_step = encoder_ref[:, :, :1].contiguous()
-        dec_step = decoder_ref[:, :, :1].contiguous()
-        traced_jd_single = torch.jit.trace(
-            jd_single,
-            (enc_step, dec_step),
-            strict=False,
-        )
-        traced_jd_single.eval()
-        jd_single_inputs = [
-            ct.TensorType(name="encoder_step", shape=(1, enc_step.shape[1], 1), dtype=np.float32),
-            ct.TensorType(name="decoder_step", shape=(1, dec_step.shape[1], 1), dtype=np.float32),
-        ]
-        jd_single_outputs = [
-            ct.TensorType(name="token_id", dtype=np.int32),
-            ct.TensorType(name="token_prob", dtype=np.float32),
-            ct.TensorType(name="duration", dtype=np.int32),
-        ]
-        # Single-step JointDecision: CPU only
-        jd_single_model = _coreml_convert(
-            traced_jd_single,
-            jd_single_inputs,
-            jd_single_outputs,
-            export_settings,
-            compute_units_override=ct.ComputeUnit.CPU_ONLY,
-        )
-        jd_single_path = output_dir / "parakeet_joint_decision_single_step.mlpackage"
-        _save_mlpackage(
-            jd_single_model,
-            jd_single_path,
-            "Parakeet single-step joint decision (current frame)",
-        )
-
-        metadata: Dict[str, object] = {
-            "model_id": model_id,
-            "sample_rate": sample_rate,
-            "max_audio_seconds": export_settings.max_audio_seconds,
-            "max_audio_samples": max_samples,
-            "max_symbol_steps": export_settings.max_symbol_steps,
-            "vocab_size": vocab_size,
-            "joint_extra_outputs": num_extra,
-            "checkpoint": checkpoint_meta,
-            "coreml": {
-                "compute_units": export_settings.compute_units.name,
-                "compute_precision": (
-                    export_settings.compute_precision.name
-                    if export_settings.compute_precision is not None
-                    else "FLOAT32"
-                ),
-            },
-            "components": {
-                "preprocessor": {
-                    "inputs": {
-                        "audio_signal": list(_tensor_shape(audio_tensor)),
-                        "audio_length": [1],
-                    },
-                    "outputs": {
-                        "mel": list(_tensor_shape(mel_ref)),
-                        "mel_length": [1],
-                    },
-                    "path": preprocessor_path.name,
-                },
-                "encoder": {
-                    "inputs": {
-                        "mel": list(_tensor_shape(mel_ref)),
-                        "mel_length": [1],
-                    },
-                    "outputs": {
-                        "encoder": list(_tensor_shape(encoder_ref)),
-                        "encoder_length": [1],
-                    },
-                    "path": encoder_path.name,
-                },
-                "mel_encoder": {
-                    "inputs": {
-                        "audio_signal": [1, max_samples],
-                        "audio_length": [1],
-                    },
-                    "outputs": {
-                        "encoder": list(_tensor_shape(encoder_ref)),
-                        "encoder_length": [1],
-                    },
-                    "path": mel_encoder_path.name,
-                },
-                "decoder": {
-                    "inputs": {
-                        "targets": list(_tensor_shape(targets)),
-                        "target_length": [1],
-                        "h_in": list(_tensor_shape(zero_state)),
-                        "c_in": list(_tensor_shape(zero_state)),
-                    },
-                    "outputs": {
-                        "decoder": list(_tensor_shape(decoder_ref)),
-                        "h_out": list(_tensor_shape(h_ref)),
-                        "c_out": list(_tensor_shape(c_ref)),
-                    },
-                    "path": decoder_path.name,
-                },
-                "joint": {
-                    "inputs": {
-                        "encoder": list(_tensor_shape(encoder_ref)),
-                        "decoder": list(_tensor_shape(decoder_ref)),
-                    },
-                    "outputs": {
-                        "logits": list(_tensor_shape(joint_ref)),
-                    },
-                    "path": joint_path.name,
-                },
-                "joint_decision": {
-                    "inputs": {
-                        "encoder": list(_tensor_shape(encoder_ref)),
-                        "decoder": list(_tensor_shape(decoder_ref)),
-                    },
-                    "outputs": {
-                        "token_id": [
-                            _tensor_shape(encoder_ref)[0],
-                            _tensor_shape(encoder_ref)[1],
-                            _tensor_shape(decoder_ref)[1],
-                        ],
-                        "token_prob": [
-                            _tensor_shape(encoder_ref)[0],
-                            _tensor_shape(encoder_ref)[1],
-                            _tensor_shape(decoder_ref)[1],
-                        ],
-                        "duration": [
-                            _tensor_shape(encoder_ref)[0],
-                            _tensor_shape(encoder_ref)[1],
-                            _tensor_shape(decoder_ref)[1],
-                        ],
-                    },
-                    "path": joint_decision_path.name,
-                },
-                "joint_decision_single_step": {
-                    "inputs": {
-                        "encoder_step": [1, _tensor_shape(encoder_ref)[2-1], 1],
-                        "decoder_step": [1, _tensor_shape(decoder_ref)[2-1], 1],
-                    },
-                    "outputs": {
-                        "token_id": [1, 1, 1],
-                        "token_prob": [1, 1, 1],
-                        "duration": [1, 1, 1],
-                    },
-                    "path": jd_single_path.name,
-                },
-            },
-        }
-
-        metadata_path = output_dir / "metadata.json"
-        metadata_path.write_text(json.dumps(metadata, indent=2))
-        typer.echo(f"Export complete. Metadata written to {metadata_path}")
-
-    finally:
-        asr_model.decoder._rnnt_export = decoder_export_flag
-
-
-if __name__ == "__main__":
-    app()

convert/parakeet-tdt-v2-0.6b/coreml/individual_components.py

@@ -1,229 +0,0 @@
-#!/usr/bin/env python3
-"""Export Parakeet TDT v2 RNNT components into CoreML and validate outputs."""
-from __future__ import annotations
-
-from dataclasses import dataclass
-from pathlib import Path
-from typing import  Optional, Tuple
-
-import coremltools as ct
-import torch
-
-
-@dataclass
-class ExportSettings:
-    output_dir: Path
-    compute_units: ct.ComputeUnit
-    deployment_target: Optional[ct.target.iOS17]
-    compute_precision: Optional[ct.precision]
-    max_audio_seconds: float
-    max_symbol_steps: int
-
-
-@dataclass
-class ValidationSettings:
-    audio_path: Optional[Path]
-    seconds: float
-    seed: Optional[int]
-    rtol: float
-    atol: float
-    skip: bool
-
-
-@dataclass
-class ValidationDiff:
-    name: str
-    max_abs_diff: float
-    max_rel_diff: float
-
-
-@dataclass
-class ValidationResult:
-    source: str
-    audio_num_samples: int
-    audio_seconds: float
-    token_length: int
-    atol: float
-    rtol: float
-    diffs: Tuple[ValidationDiff, ...]
-
-
-class PreprocessorWrapper(torch.nn.Module):
-    def __init__(self, module: torch.nn.Module) -> None:
-        super().__init__()
-        self.module = module
-
-    def forward(self, audio_signal: torch.Tensor, length: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
-        mel, mel_length = self.module(input_signal=audio_signal, length=length.to(dtype=torch.long))
-        return mel, mel_length
-
-
-class EncoderWrapper(torch.nn.Module):
-    def __init__(self, module: torch.nn.Module) -> None:
-        super().__init__()
-        self.module = module
-
-    def forward(self, features: torch.Tensor, length: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
-        encoded, encoded_lengths = self.module(audio_signal=features, length=length.to(dtype=torch.long))
-        return encoded, encoded_lengths
-
-
-class DecoderWrapper(torch.nn.Module):
-    def __init__(self, module: torch.nn.Module) -> None:
-        super().__init__()
-        self.module = module
-
-    def forward(
-        self,
-        targets: torch.Tensor,
-        target_lengths: torch.Tensor,
-        h_in: torch.Tensor,
-        c_in: torch.Tensor,
-    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-        state = [h_in, c_in]
-        decoder_output, _, new_state = self.module(
-            targets=targets.to(dtype=torch.long),
-            target_length=target_lengths.to(dtype=torch.long),
-            states=state,
-        )
-        return decoder_output, new_state[0], new_state[1]
-
-
-class JointWrapper(torch.nn.Module):
-    def __init__(self, module: torch.nn.Module) -> None:
-        super().__init__()
-        self.module = module
-
-    def forward(self, encoder_outputs: torch.Tensor, decoder_outputs: torch.Tensor) -> torch.Tensor:
-        # Input: encoder_outputs [B, D, T], decoder_outputs [B, D, U]
-        # Transpose to match what projection layers expect
-        encoder_outputs = encoder_outputs.transpose(1, 2)  # [B, T, D]
-        decoder_outputs = decoder_outputs.transpose(1, 2)  # [B, U, D]
-
-        # Apply projections
-        enc_proj = self.module.enc(encoder_outputs)        # [B, T, 640]
-        dec_proj = self.module.pred(decoder_outputs)       # [B, U, 640]
-
-        # Explicit broadcasting along T and U to avoid converter ambiguity
-        x = enc_proj.unsqueeze(2) + dec_proj.unsqueeze(1)  # [B, T, U, 640]
-        x = self.module.joint_net[0](x)                   # ReLU
-        x = self.module.joint_net[1](x)                   # Dropout (no-op in eval)
-        out = self.module.joint_net[2](x)                 # Linear -> logits [B, T, U, 8198]
-        return out
-
-
-class MelEncoderWrapper(torch.nn.Module):
-    """Fused wrapper: waveform -> mel -> encoder.
-
-    Inputs:
-      - audio_signal: [B, S]
-      - audio_length: [B]
-
-    Outputs:
-      - encoder: [B, D, T_enc]
-      - encoder_length: [B]
-    """
-    def __init__(self, preprocessor: PreprocessorWrapper, encoder: EncoderWrapper) -> None:
-        super().__init__()
-        self.preprocessor = preprocessor
-        self.encoder = encoder
-
-    def forward(self, audio_signal: torch.Tensor, audio_length: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
-        mel, mel_length = self.preprocessor(audio_signal, audio_length)
-        encoded, enc_len = self.encoder(mel, mel_length.to(dtype=torch.int32))
-        return encoded, enc_len
-
-
-class JointDecisionWrapper(torch.nn.Module):
-    """Joint + decision head: outputs label id, label prob, duration frames.
-
-    Splits joint logits into token logits and duration logits, applies softmax
-    over tokens, argmax for both heads, and gathers probability of the chosen token.
-
-    Inputs:
-      - encoder_outputs: [B, D, T]
-      - decoder_outputs: [B, D, U]
-
-    Returns:
-      - token_id: [B, T, U] int32
-      - token_prob: [B, T, U] float32
-      - duration: [B, T, U] int32  (frames; duration buckets as defined by the model)
-    """
-    def __init__(self, joint: JointWrapper, vocab_size: int, num_extra: int) -> None:
-        super().__init__()
-        self.joint = joint
-        self.vocab_with_blank = int(vocab_size) + 1
-        self.num_extra = int(num_extra)
-
-    def forward(self, encoder_outputs: torch.Tensor, decoder_outputs: torch.Tensor):
-        logits = self.joint(encoder_outputs, decoder_outputs)
-        token_logits = logits[..., : self.vocab_with_blank]
-        duration_logits = logits[..., -self.num_extra :]
-
-        # Token selection
-        token_ids = torch.argmax(token_logits, dim=-1).to(dtype=torch.int32)
-        token_probs_all = torch.softmax(token_logits, dim=-1)
-        # gather expects int64 (long) indices; cast only for gather
-        token_prob = torch.gather(
-            token_probs_all, dim=-1, index=token_ids.long().unsqueeze(-1)
-        ).squeeze(-1)
-
-        # Duration prediction (duration bucket argmax → frame increments)
-        duration = torch.argmax(duration_logits, dim=-1).to(dtype=torch.int32)
-        return token_ids, token_prob, duration
-
-
-class JointDecisionSingleStep(torch.nn.Module):
-    """Single-step variant for streaming: encoder_step [1, 1024, 1] -> [1,1,1].
-
-    Inputs:
-      - encoder_step: [B=1, D=1024, T=1]
-      - decoder_step: [B=1, D=640, U=1]
-
-    Returns:
-      - token_id: [1, 1, 1] int32
-      - token_prob: [1, 1, 1] float32
-      - duration: [1, 1, 1] int32
-    """
-    def __init__(self, joint: JointWrapper, vocab_size: int, num_extra: int) -> None:
-        super().__init__()
-        self.joint = joint
-        self.vocab_with_blank = int(vocab_size) + 1
-        self.num_extra = int(num_extra)
-
-    def forward(self, encoder_step: torch.Tensor, decoder_step: torch.Tensor):
-        # Reuse JointWrapper which expects [B, D, T] and [B, D, U]
-        logits = self.joint(encoder_step, decoder_step)  # [1, 1, 1, V+extra]
-        token_logits = logits[..., : self.vocab_with_blank]
-        duration_logits = logits[..., -self.num_extra :]
-
-        token_ids = torch.argmax(token_logits, dim=-1, keepdim=False).to(dtype=torch.int32)
-        token_probs_all = torch.softmax(token_logits, dim=-1)
-        token_prob = torch.gather(
-            token_probs_all, dim=-1, index=token_ids.long().unsqueeze(-1)
-        ).squeeze(-1)
-        duration = torch.argmax(duration_logits, dim=-1, keepdim=False).to(dtype=torch.int32)
-        return token_ids, token_prob, duration
-
-
-def _coreml_convert(
-    traced: torch.jit.ScriptModule,
-    inputs,
-    outputs,
-    settings: ExportSettings,
-    compute_units_override: Optional[ct.ComputeUnit] = None,
-) -> ct.models.MLModel:
-    cu = compute_units_override if compute_units_override is not None else settings.compute_units
-    kwargs = {
-        "convert_to": "mlprogram",
-        "inputs": inputs,
-        "outputs": outputs,
-        "compute_units": cu,
-    }
-    print("Converting:", traced.__class__.__name__)
-    print("Conversion kwargs:", kwargs)
-    if settings.deployment_target is not None:
-        kwargs["minimum_deployment_target"] = settings.deployment_target
-    if settings.compute_precision is not None:
-        kwargs["compute_precision"] = settings.compute_precision
-    return ct.convert(traced, **kwargs)

convert/parakeet-tdt-v2-0.6b/coreml/quantize_coreml.py

@@ -1,1207 +0,0 @@
-#!/usr/bin/env python3
-"""Quantize CoreML mlpackages and compare quality, compression, latency, and compile time.
-
-This script focuses on the fused models:
-  - parakeet_mel_encoder.mlpackage (waveform -> encoder)
-  - parakeet_joint_decision.mlpackage (joint + softmax/argmax)
-
-It also quantizes the rest of the components in the directory for completeness.
-
-Variants tried by default (examples):
-  - int8-linear (per-channel)
-  - palettize 6-bit (Mel-only)
-  - jd-only variants (int8 and palette)
-  - prune + int8
-
-Outputs:
-  - A new directory per variant under <output_root>/<variant>/ with quantized mlpackages
-  - quantization_summary.json with aggregate metrics (baseline vs each variant)
-  - Plots saved under <output_root>/plots/ and mirrored to <repo>/plots/quantize/<compute_units_lower>/
-    - fused_quality.png / fused_latency.png / fused_compression.png / fused_size.png (latency chart also shows compile time)
-    - all_components_quality.png / all_components_latency.png / all_components_compression.png / all_components_size.png (latency chart also shows compile time)
-
-Notes:
-  - Uses CoreMLTools optimize.coreml linear/palettize/prune for MLProgram models.
-  - Keeps the fixed 15-second window shapes as per context/coreml_component_io.md.
-  - Run via `uv run` to ensure reproducible environment.
-  - Sets minimum deployment target to iOS 17 for all outputs and loads models with `compute_units=ALL` by default to enable ANE.
-  - Tracks offline compile time (mlpackage->mlmodelc) and includes that metric in plots.
-"""
-from __future__ import annotations
-
-import json
-import platform
-import shutil
-import subprocess
-import time
-from dataclasses import dataclass
-from pathlib import Path
-from typing import Dict, List, Optional, Set, Tuple
-
-import numpy as np
-import soundfile as sf
-import typer
-
-import coremltools as ct
-from coremltools.optimize.coreml import (
-    OptimizationConfig,
-    OpLinearQuantizerConfig,
-    OpPalettizerConfig,
-    OpThresholdPrunerConfig,
-    linear_quantize_weights,
-    palettize_weights,
-    prune_weights,
-)
-
-
-# Optional plotting
-try:
-    import matplotlib
-    matplotlib.use("Agg")
-    import matplotlib.pyplot as plt
-    HAS_MPL = True
-except Exception:
-    HAS_MPL = False
-
-
-BASE_DIR = Path(__file__).resolve().parent
-BYTES_IN_MB = 1024 * 1024
-
-
-DISPLAY_LABEL_OVERRIDES = {
-    "mel6bit-palettize": "mel6bit",
-    "enc6bit-palettize": "enc6bit",
-    "int8-linear": "int8 linear",
-}
-
-
-# When no explicit components are provided, we derive the set of
-# components to quantize from the selected variants' whitelists.
-# If any selected variant has no whitelist (i.e., applies globally),
-# we quantize all components.
-DEFAULT_TARGET_COMPONENTS: Tuple[str, str] = ()
-
-
-# Formatting helpers -------------------------------------------------------
-
-
-def _format_labels(labels: List[str]) -> List[str]:
-    formatted: List[str] = []
-    for label in labels:
-        pretty = DISPLAY_LABEL_OVERRIDES.get(label, label)
-        pretty = pretty.replace("_", " ")
-        pretty = pretty.replace("-", "\n")
-        formatted.append(pretty)
-    return formatted
-
-
-def _friendly_component_name(name: str) -> str:
-    return name.replace("_", " ").title()
-
-
-def _get_metric_series(
-    metrics: Dict[str, Dict[str, List[float]]],
-    component: str,
-    key: str,
-    length: int,
-) -> List[float]:
-    values = metrics.get(component, {}).get(key)
-    if values is None:
-        return [float("nan")] * length
-    if len(values) >= length:
-        return list(values[:length])
-    padded = list(values)
-    padded.extend([float("nan")] * (length - len(values)))
-    return padded
-
-
-def _plot_bar_rows(
-    out_path: Path,
-    display_labels: List[str],
-    rows: List[Dict[str, object]],
-    title_suffix: str,
-) -> None:
-    if not HAS_MPL:
-        return
-    if not rows or not display_labels:
-        return
-
-    n_labels = len(display_labels)
-    fig_width = max(8.0, 1.3 * n_labels)
-    fig_height = max(2.6 * len(rows), 3.0)
-
-    fig, axes = plt.subplots(len(rows), 1, figsize=(fig_width, fig_height), squeeze=False)
-    axes = axes.flatten()
-    x = np.arange(n_labels)
-
-    for ax, row in zip(axes, rows):
-        values = np.asarray(row.get("values", [float("nan")] * n_labels), dtype=np.float64)
-        color = row.get("color")
-        bars = ax.bar(x, values, width=0.6, color=color)
-        ax.set_title(str(row.get("title", "")))
-        ax.set_xticks(x)
-        ax.set_xticklabels(display_labels, rotation=25, ha="right")
-        ylim = row.get("ylim")
-        if isinstance(ylim, tuple) and len(ylim) == 2:
-            ax.set_ylim(ylim)
-        ax.grid(axis="y", linestyle="--", alpha=0.3)
-        # Add value labels for bars
-        for bar in bars:
-            h = bar.get_height()
-            if np.isnan(h):
-                continue
-            ax.annotate(
-                f"{h:.2f}",
-                xy=(bar.get_x() + bar.get_width() / 2, h),
-                xytext=(0, 3),
-                textcoords="offset points",
-                ha="center",
-                va="bottom",
-                fontsize=8,
-            )
-
-    if title_suffix:
-        fig.suptitle(title_suffix, fontsize=12)
-        fig.tight_layout(rect=(0, 0, 1, 0.95))
-    else:
-        fig.tight_layout()
-
-    out_path.parent.mkdir(parents=True, exist_ok=True)
-    plt.savefig(out_path)
-    plt.close(fig)
-
-
-def _plot_fused_category_charts(
-    plot_dir: Path,
-    labels: List[str],
-    mel_quality: List[float],
-    mel_latency_ms: List[float],
-    mel_compression: List[float],
-    mel_size_mb: List[float],
-    jd_acc: List[float],
-    jd_latency_ms: List[float],
-    jd_compression: List[float],
-    jd_size_mb: List[float],
-    mel_compile_ms: Optional[List[float]] = None,
-    jd_compile_ms: Optional[List[float]] = None,
-    title_suffix: str = "",
-) -> List[Path]:
-    if not HAS_MPL:
-        return []
-    outputs: List[Path] = []
-    display = _format_labels(labels)
-    prefix = f"Fused Components — {title_suffix}" if title_suffix else "Fused Components"
-
-    quality_path = plot_dir / "fused_quality.png"
-    _plot_bar_rows(
-        quality_path,
-        display,
-        [
-            {
-                "title": "MelEncoder quality (1 - norm err)",
-                "values": mel_quality,
-                "color": "C0",
-                "ylim": (0.0, 1.05),
-            },
-            {
-                "title": "JointDecision token-id match rate",
-                "values": jd_acc,
-                "color": "C0",
-                "ylim": (0.0, 1.05),
-            },
-        ],
-        f"{prefix} — Quality",
-    )
-    outputs.append(quality_path)
-
-    compression_path = plot_dir / "fused_compression.png"
-    _plot_bar_rows(
-        compression_path,
-        display,
-        [
-            {
-                "title": "MelEncoder compression ratio",
-                "values": mel_compression,
-                "color": "C2",
-            },
-            {
-                "title": "JointDecision compression ratio",
-                "values": jd_compression,
-                "color": "C2",
-            },
-        ],
-        f"{prefix} — Compression",
-    )
-    outputs.append(compression_path)
-
-    size_path = plot_dir / "fused_size.png"
-    _plot_bar_rows(
-        size_path,
-        display,
-        [
-            {
-                "title": "MelEncoder size (MB)",
-                "values": mel_size_mb,
-                "color": "C4",
-            },
-            {
-                "title": "JointDecision size (MB)",
-                "values": jd_size_mb,
-                "color": "C4",
-            },
-        ],
-        f"{prefix} — Size",
-    )
-    outputs.append(size_path)
-
-    latency_rows = [
-        {
-            "title": "MelEncoder latency (ms)",
-            "values": mel_latency_ms,
-            "color": "C1",
-        },
-        {
-            "title": "JointDecision latency (ms)",
-            "values": jd_latency_ms,
-            "color": "C1",
-        },
-    ]
-    if mel_compile_ms is not None and jd_compile_ms is not None:
-        latency_rows.extend(
-            [
-                {
-                    "title": "MelEncoder compile (ms)",
-                    "values": mel_compile_ms,
-                    "color": "C3",
-                },
-                {
-                    "title": "JointDecision compile (ms)",
-                    "values": jd_compile_ms,
-                    "color": "C3",
-                },
-            ]
-        )
-
-    latency_path = plot_dir / "fused_latency.png"
-    _plot_bar_rows(
-        latency_path,
-        display,
-        latency_rows,
-        f"{prefix} — Latency",
-    )
-    outputs.append(latency_path)
-
-    return outputs
-
-
-def _plot_all_component_category_charts(
-    plot_dir: Path,
-    labels: List[str],
-    metrics: Dict[str, Dict[str, List[float]]],
-    title_suffix: str = "",
-) -> List[Path]:
-    if not HAS_MPL:
-        return []
-    outputs: List[Path] = []
-    display = _format_labels(labels)
-    prefix = f"Component Breakdown — {title_suffix}" if title_suffix else "Component Breakdown"
-    comp_order = [
-        "preprocessor",
-        "encoder",
-        "mel_encoder",
-        "decoder",
-        "joint",
-        "joint_decision",
-    ]
-    n = len(labels)
-
-    quality_rows: List[Dict[str, object]] = []
-    for comp in comp_order:
-        friendly = _friendly_component_name(comp)
-        if comp == "joint_decision":
-            values = _get_metric_series(metrics, comp, "acc", n)
-            title = f"{friendly} token-id match rate"
-        else:
-            values = _get_metric_series(metrics, comp, "quality", n)
-            title = f"{friendly} quality (1 - norm err)"
-        quality_rows.append({"title": title, "values": values, "color": "C0", "ylim": (0.0, 1.05)})
-
-    compression_rows: List[Dict[str, object]] = []
-    for comp in comp_order:
-        friendly = _friendly_component_name(comp)
-        compression_rows.append(
-            {
-                "title": f"{friendly} compression ratio",
-                "values": _get_metric_series(metrics, comp, "compression", n),
-                "color": "C2",
-            }
-        )
-
-    size_rows: List[Dict[str, object]] = []
-    for comp in comp_order:
-        friendly = _friendly_component_name(comp)
-        size_rows.append(
-            {
-                "title": f"{friendly} size (MB)",
-                "values": _get_metric_series(metrics, comp, "size_mb", n),
-                "color": "C4",
-            }
-        )
-
-    latency_rows: List[Dict[str, object]] = []
-    for comp in comp_order:
-        friendly = _friendly_component_name(comp)
-        latency_rows.append(
-            {
-                "title": f"{friendly} latency (ms)",
-                "values": _get_metric_series(metrics, comp, "latency_ms", n),
-                "color": "C1",
-            }
-        )
-
-    compile_rows: List[Dict[str, object]] = []
-    for comp in comp_order:
-        friendly = _friendly_component_name(comp)
-        compile_rows.append(
-            {
-                "title": f"{friendly} compile (ms)",
-                "values": _get_metric_series(metrics, comp, "compile_ms", n),
-                "color": "C3",
-            }
-        )
-
-    quality_path = plot_dir / "all_components_quality.png"
-    _plot_bar_rows(quality_path, display, quality_rows, f"{prefix} — Quality")
-    outputs.append(quality_path)
-
-    compression_path = plot_dir / "all_components_compression.png"
-    _plot_bar_rows(compression_path, display, compression_rows, f"{prefix} — Compression")
-    outputs.append(compression_path)
-
-    size_path = plot_dir / "all_components_size.png"
-    _plot_bar_rows(size_path, display, size_rows, f"{prefix} — Size")
-    outputs.append(size_path)
-
-    latency_path = plot_dir / "all_components_latency.png"
-    _plot_bar_rows(latency_path, display, latency_rows, f"{prefix} — Latency")
-    outputs.append(latency_path)
-
-    compile_path = plot_dir / "all_components_compile.png"
-    _plot_bar_rows(compile_path, display, compile_rows, f"{prefix} — Compile")
-    outputs.append(compile_path)
-
-    return outputs
-
-
-app = typer.Typer(add_completion=False, pretty_exceptions_show_locals=False)
-
-
-@dataclass
-class VariantConfig:
-    name: str
-    # List of (op_kind, opt_config) steps applied in order; op_kind in {'linear','palettize','prune'}
-    steps: List[Tuple[str, OptimizationConfig]]
-    category: str
-    whitelist: Optional[List[str]] = None  # component names to apply; others copied
-
-
-def _dir_size_bytes(path: Path) -> int:
-    total = 0
-    for p in path.rglob("*"):
-        if p.is_file():
-            try:
-                total += p.stat().st_size
-            except OSError:
-                pass
-    return total
-
-
-def _load_metadata(input_dir: Path) -> Dict[str, object]:
-    meta_path = input_dir / "metadata.json"
-    if not meta_path.exists():
-        raise typer.BadParameter(f"Expected metadata.json in {input_dir}")
-    return json.loads(meta_path.read_text())
-
-
-def _prepare_audio(
-    seconds: float,
-    sample_rate: int,
-    audio_path: Optional[Path],
-) -> Tuple[np.ndarray, np.ndarray]:
-    max_samples = int(round(seconds * sample_rate))
-    if audio_path is None:
-        # random but deterministic sample
-        rng = np.random.default_rng(1234)
-        audio = rng.standard_normal(size=(1, max_samples), dtype=np.float32).astype(np.float32)
-    else:
-        data, sr = sf.read(str(audio_path), dtype="float32")
-        if sr != sample_rate:
-            raise typer.BadParameter(
-                f"Validation audio sample rate {sr} != expected {sample_rate}"
-            )
-        if data.ndim > 1:
-            data = data[:, 0]
-        if data.size < max_samples:
-            data = np.pad(data, (0, max_samples - data.size))
-        elif data.size > max_samples:
-            data = data[:max_samples]
-        audio = data.reshape(1, -1).astype(np.float32, copy=False)
-    length = np.array([max_samples], dtype=np.int32)
-    return audio, length
-
-
-def _predict_latency(model: ct.models.MLModel, inputs: Dict[str, np.ndarray], runs: int = 10, warmup: int = 3) -> Tuple[float, float]:
-    # Warmup
-    for _ in range(max(0, warmup)):
-        _ = model.predict(inputs)
-    # Timed runs
-    times: List[float] = []
-    for _ in range(max(1, runs)):
-        t0 = time.perf_counter()
-        _ = model.predict(inputs)
-        t1 = time.perf_counter()
-        times.append((t1 - t0) * 1000.0)  # ms
-    arr = np.array(times, dtype=np.float64)
-    return float(arr.mean()), float(arr.std(ddof=1) if arr.size > 1 else 0.0)
-
-
-def _max_abs_rel(a: np.ndarray, b: np.ndarray) -> Tuple[float, float]:
-    a = np.asarray(a, dtype=np.float32)
-    b = np.asarray(b, dtype=np.float32)
-    if a.size == 0:
-        return 0.0, 0.0
-    diff = np.abs(a - b)
-    max_abs = float(diff.max())
-    denom = np.maximum(np.abs(a), np.abs(b))
-    with np.errstate(divide="ignore", invalid="ignore"):
-        rel = np.where(denom == 0.0, 0.0, diff / denom)
-    max_rel = float(rel.max())
-    return max_abs, max_rel
-
-
-def _save_mlpackage(model: ct.models.MLModel, path: Path, description: str) -> None:
-    path.parent.mkdir(parents=True, exist_ok=True)
-    # Ensure iOS 17 target for proper MLProgram ops (e.g., blockwise shift/scale)
-    try:
-        model.minimum_deployment_target = ct.target.iOS17
-    except Exception:
-        pass
-    model.short_description = description
-    model.author = "Fluid Inference"
-    model.save(str(path))
-
-
-def _offline_compile_time_ms(model_path: Path) -> float:
-    """Compile mlpackage -> mlmodelc and return wall time in ms (host offline compile).
-
-    Returns NaN on failure.
-    """
-    compiled_dir: Optional[Path] = None
-    expected_dir = model_path.with_suffix(".mlmodelc")
-    try:
-        # Delete any prior compile artifact in-place so we can measure a fresh build
-        if expected_dir.exists():
-            shutil.rmtree(expected_dir, ignore_errors=True)
-
-        t0 = time.perf_counter()
-        compiled_path = ct.utils.compile_model(str(model_path))
-        t1 = time.perf_counter()
-        compiled_dir = Path(compiled_path)
-        return (t1 - t0) * 1000.0
-    except Exception:
-        return float("nan")
-    finally:
-        if (
-            compiled_dir is not None
-            and compiled_dir.exists()
-            and compiled_dir == expected_dir
-        ):
-            shutil.rmtree(compiled_dir, ignore_errors=True)
-
-def _chip_spec_string(compute_units: str) -> str:
-    try:
-        chip = subprocess.check_output(["sysctl", "-n", "machdep.cpu.brand_string"]).decode().strip()
-    except Exception:
-        chip = platform.processor() or platform.machine()
-    mac_ver = platform.mac_ver()[0] or platform.platform()
-    return f"Host: {chip} • macOS {mac_ver} • CoreMLTools {ct.__version__} • ComputeUnits={compute_units} • Min Target: iOS17"
-
-
-def _quantize_dir(
-    input_dir: Path,
-    output_dir: Path,
-    variant: VariantConfig,
-    global_whitelist: Optional[Set[str]] = None,
-) -> Dict[str, str]:
-    """Quantize all mlpackages in input_dir into output_dir using given variant config.
-
-    Returns a map of component name -> saved relative path.
-    """
-    meta = _load_metadata(input_dir)
-    comps = meta.get("components", {})
-    saved: Dict[str, str] = {}
-    for name, cfg in comps.items():
-        src_name = cfg.get("path")
-        if not src_name:
-            continue
-        src_path = input_dir / src_name
-        if not src_path.exists():
-            continue
-        dst_path = output_dir / src_name
-        # Use CPU+GPU for preprocessor to avoid NE preprocessor input size issues; others use CPU+NE
-        cu = ct.ComputeUnit.CPU_AND_GPU if name == "preprocessor" else ct.ComputeUnit.CPU_AND_NE
-        base_model = ct.models.MLModel(str(src_path), compute_units=cu)
-        # Target iOS17 when running optimizations so the right ops are chosen
-        try:
-            base_model.minimum_deployment_target = ct.target.iOS17
-        except Exception:
-            pass
-
-        if name == "decoder":
-            typer.echo(f"[{variant.name}] Skipping decoder quantization; copying baseline: {src_name}")
-            _save_mlpackage(base_model, dst_path, "Baseline copy (decoder quantization disabled) - decoder")
-            saved[name] = dst_path.name
-            continue
-
-        skip_reasons: List[str] = []
-        if variant.whitelist is not None and name not in variant.whitelist:
-            skip_reasons.append("not targeted by variant")
-        if global_whitelist is not None and name not in global_whitelist:
-            skip_reasons.append("not in requested components")
-
-        if skip_reasons:
-            reason = "; ".join(skip_reasons)
-            typer.echo(f"[{variant.name}] Skipping {name} ({reason}); copying baseline: {src_name}")
-            _save_mlpackage(base_model, dst_path, f"Baseline copy ({reason}) - {name}")
-            saved[name] = dst_path.name
-            continue
-
-        typer.echo(f"[{variant.name}] Quantizing {name}: {src_name}")
-
-        try:
-            q_model = base_model
-            for step_kind, step_cfg in variant.steps:
-                if step_kind == 'linear':
-                    q_model = linear_quantize_weights(q_model, step_cfg)
-                elif step_kind == 'palettize':
-                    q_model = palettize_weights(q_model, step_cfg)
-                elif step_kind == 'prune':
-                    q_model = prune_weights(q_model, step_cfg)
-                else:
-                    raise ValueError(f"Unknown variant step: {step_kind}")
-        except Exception as e:
-            # If quantization fails (e.g., unsupported op), fall back to copying baseline.
-            typer.echo(f"  ! Failed to quantize {name} with {variant.name}: {e}. Copying baseline.")
-            _save_mlpackage(base_model, dst_path, f"Baseline copy (failed to quantize) - {name}")
-        else:
-            _save_mlpackage(q_model, dst_path, f"{variant.name} quantized - {name}")
-        saved[name] = dst_path.name
-    # Persist a variant metadata shim
-    out_meta = {
-        "variant": variant.name,
-        "base_dir": str(input_dir.resolve()),
-        "components": saved,
-    }
-    (output_dir / "quantization_metadata.json").write_text(json.dumps(out_meta, indent=2))
-    return saved
-
-
-@app.command()
-def quantize(
-    input_dir: Path = typer.Option(Path("parakeet_coreml"), help="Directory containing baseline mlpackages + metadata.json"),
-    output_root: Path = typer.Option(Path("parakeet_coreml_quantized"), help="Root output dir for quantized variants"),
-    validation_audio: Optional[Path] = typer.Option(None, exists=True, resolve_path=True, help="Optional 15s, 16kHz wav for evaluation (defaults to bundled audio if present)"),
-    compute_units: str = typer.Option("CPU_AND_NE", help="Compute units for evaluation of non-preprocessor models. Preprocessor is forced to CPU_AND_GPU."),
-    runs: int = typer.Option(10, help="Timed runs per model for latency measurement"),
-    categories: Optional[List[str]] = typer.Option(
-        None,
-        "--category",
-        "-c",
-        help="Only run quantization variants in these categories (e.g., linear, mel-palettize). Can be repeated.",
-    ),
-    components: Optional[List[str]] = typer.Option(
-        None,
-        "--component",
-        "-m",
-        help="Component names to quantize. Defaults to mel_encoder and joint_decision. Use 'all' to keep every component enabled.",
-    ),
-) -> None:
-    """Quantize models, then compare quality, compression, latency, and compile time.
-
-    Variants include int8-linear (per-channel/per-tensor/block), palettization (6-bit),
-    jd-only probes, and prune+int8. Baseline is the pre-converted models.
-    """
-    meta = _load_metadata(input_dir)
-    sr = int(meta.get("sample_rate", 16000))
-    seconds = float(meta.get("max_audio_seconds", 15.0))
-
-    components_meta = meta.get("components", {})
-    component_lookup = {name.lower(): name for name in components_meta.keys()}
-    # Defer computing component_filter until after variant/category selection so
-    # that, by default, we quantize exactly the components targeted by the
-    # selected variants (instead of a hard-coded subset).
-    component_filter: Optional[Set[str]] = None
-
-    # Default audio if present
-    default_audio = (BASE_DIR / "audio" / "yc_first_minute_16k_15s.wav").resolve()
-    audio_path = validation_audio if validation_audio is not None else (default_audio if default_audio.exists() else None)
-    if audio_path is not None and validation_audio is None:
-        typer.echo(f"Using default validation audio: {audio_path}")
-    audio, audio_len = _prepare_audio(seconds, sr, audio_path)
-
-    # Load baseline models and helpers for inputs
-    # Baseline models for input preparation
-    # Force preprocessor to CPU+GPU and all other components to CPU+NE for evaluation
-    pre_base = ct.models.MLModel(str(input_dir / "parakeet_preprocessor.mlpackage"), compute_units=ct.ComputeUnit.CPU_AND_GPU)
-    enc_base = ct.models.MLModel(str(input_dir / "parakeet_encoder.mlpackage"), compute_units=ct.ComputeUnit.CPU_AND_NE)
-    mel_encoder_base = ct.models.MLModel(str(input_dir / "parakeet_mel_encoder.mlpackage"), compute_units=ct.ComputeUnit.CPU_AND_NE)
-    decoder_base = ct.models.MLModel(str(input_dir / "parakeet_decoder.mlpackage"), compute_units=ct.ComputeUnit.CPU_AND_NE)
-    joint_decision_base = ct.models.MLModel(str(input_dir / "parakeet_joint_decision.mlpackage"), compute_units=ct.ComputeUnit.CPU_AND_NE)
-    joint_base = ct.models.MLModel(str(input_dir / "parakeet_joint.mlpackage"), compute_units=ct.ComputeUnit.CPU_AND_NE)
-
-    # Prepare typical inputs once using baseline models
-    pre_out = pre_base.predict({"audio_signal": audio, "audio_length": audio_len})
-    mel_ref = np.array(pre_out["mel"], dtype=np.float32, copy=True)
-    mel_len = np.array(pre_out["mel_length"], dtype=np.int32, copy=True)
-
-    enc_out = enc_base.predict({"mel": mel_ref, "mel_length": mel_len})
-    encoder_ref = np.array(enc_out["encoder"], dtype=np.float32, copy=True)
-    encoder_len = np.array(enc_out["encoder_length"], dtype=np.int32, copy=True)
-
-    # Decoder inputs from metadata
-    dec_in = meta["components"]["decoder"]["inputs"]
-    targets_shape = tuple(int(x) for x in dec_in["targets"])  # e.g., (1, 256)
-    h_shape = tuple(int(x) for x in dec_in["h_in"])         # e.g., (2, 1, 640)
-    # Use zeros as targets for reproducibility without needing blank idx
-    targets = np.zeros(targets_shape, dtype=np.int32)
-    target_length = np.array([targets_shape[1]], dtype=np.int32)
-    h0 = np.zeros(h_shape, dtype=np.float32)
-    c0 = np.zeros(h_shape, dtype=np.float32)
-    dec_out = decoder_base.predict({
-        "targets": targets,
-        "target_length": target_length,
-        "h_in": h0,
-        "c_in": c0,
-    })
-    decoder_ref = np.array(dec_out["decoder"], dtype=np.float32, copy=True)
-
-    # Baseline sizes per component
-    pre_base_size = _dir_size_bytes(input_dir / "parakeet_preprocessor.mlpackage")
-    enc_base_size = _dir_size_bytes(input_dir / "parakeet_encoder.mlpackage")
-    mel_base_size = _dir_size_bytes(input_dir / "parakeet_mel_encoder.mlpackage")
-    dec_base_size = _dir_size_bytes(input_dir / "parakeet_decoder.mlpackage")
-    joint_base_size = _dir_size_bytes(input_dir / "parakeet_joint.mlpackage")
-    jd_base_size = _dir_size_bytes(input_dir / "parakeet_joint_decision.mlpackage")
-
-    # Baseline latencies
-    pre_base_inputs = {"audio_signal": audio, "audio_length": audio_len}
-    enc_base_inputs = {"mel": mel_ref, "mel_length": mel_len}
-    mel_base_inputs = {"audio_signal": audio, "audio_length": audio_len}
-    dec_base_inputs = {"targets": targets, "target_length": target_length, "h_in": h0, "c_in": c0}
-    joint_base_inputs = {"encoder": encoder_ref, "decoder": decoder_ref}
-    jd_base_inputs = {"encoder": encoder_ref, "decoder": decoder_ref}
-    pre_base_ms, _ = _predict_latency(pre_base, pre_base_inputs, runs=runs)
-    enc_base_ms, _ = _predict_latency(enc_base, enc_base_inputs, runs=runs)
-    mel_base_ms, _ = _predict_latency(mel_encoder_base, mel_base_inputs, runs=runs)
-    dec_base_ms, _ = _predict_latency(decoder_base, dec_base_inputs, runs=runs)
-    joint_base_ms, _ = _predict_latency(joint_base, joint_base_inputs, runs=runs)
-    jd_base_ms, _ = _predict_latency(joint_decision_base, jd_base_inputs, runs=runs)
-    # Cache baseline joint logits for comparisons
-    logits_base = np.array(joint_base.predict(joint_base_inputs)["logits"], dtype=np.float32, copy=True)
-
-    # Variants
-    variants: List[VariantConfig] = [
-        VariantConfig(
-            name="int8-linear",
-            steps=[(
-                "linear",
-                OptimizationConfig(global_config=OpLinearQuantizerConfig(mode="linear", granularity="per_channel")),
-            )],
-            category="linear",
-        ),
-        # 6-bit palettization for MelEncoder only
-        VariantConfig(
-            name="mel6bit-palettize",
-            steps=[(
-                "palettize",
-                OptimizationConfig(
-                    global_config=OpPalettizerConfig(mode="kmeans", nbits=6)
-                ),
-            )],
-            category="mel-palettize",
-            whitelist=["mel_encoder"],
-        ),
-        # 6-bit palettization for Encoder only
-        VariantConfig(
-            name="enc6bit-palettize",
-            steps=[(
-                "palettize",
-                OptimizationConfig(
-                    global_config=OpPalettizerConfig(mode="kmeans", nbits=6)
-                ),
-            )],
-            category="encoder-palettize",
-            whitelist=["encoder"],
-        ),
-        # (removed) Global palettization variants
-    ]
-
-    available_categories = {variant.category for variant in variants}
-    category_lookup = {cat.lower(): cat for cat in available_categories}
-    selected_categories: Set[str]
-    if categories:
-        normalized_categories = [cat.strip().lower() for cat in categories if cat.strip()]
-        if normalized_categories:
-            invalid_categories = [cat for cat in normalized_categories if cat not in category_lookup]
-            if invalid_categories:
-                available_str = ", ".join(sorted(available_categories)) or "none"
-                bad = ", ".join(sorted(set(invalid_categories)))
-                raise typer.BadParameter(
-                    f"Unknown category (--category): {bad}. Available categories: {available_str}."
-                )
-            selected_categories = {category_lookup[cat] for cat in normalized_categories}
-            variants = [variant for variant in variants if variant.category in selected_categories]
-        else:
-            selected_categories = available_categories
-    else:
-        selected_categories = available_categories
-
-    if not variants:
-        typer.echo("No quantization variants matched the requested categories; nothing to do.")
-        raise typer.Exit(code=0)
-
-    typer.echo("Running variant categories: " + ", ".join(sorted(selected_categories)))
-
-    # Resolve the component whitelist now that variants are known.
-    if components:
-        normalized_components = [comp.strip().lower() for comp in components if comp.strip()]
-        if any(comp == "all" for comp in normalized_components):
-            component_filter = None
-        else:
-            resolved: List[str] = []
-            invalid: List[str] = []
-            for comp in normalized_components:
-                match = component_lookup.get(comp)
-                if match is None:
-                    invalid.append(comp)
-                else:
-                    resolved.append(match)
-            if invalid:
-                available = ", ".join(sorted(component_lookup.values())) or "none"
-                bad = ", ".join(sorted(set(invalid)))
-                raise typer.BadParameter(
-                    f"Unknown component(s) for --component: {bad}. Available components: {available}."
-                )
-            component_filter = set(resolved)
-    else:
-        # Derive from selected variants' whitelists
-        derived: Set[str] = set()
-        has_global = False
-        for v in variants:
-            if v.whitelist is None:
-                has_global = True
-                break
-            derived.update(v.whitelist)
-        component_filter = None if has_global or not derived else derived
-
-    if component_filter is None:
-        typer.echo("Quantizing components: all components (derived)")
-    else:
-        typer.echo("Quantizing components: " + ", ".join(sorted(component_filter)) + " (derived)")
-
-    # Aggregate results (baseline + variants)
-    summary: Dict[str, Dict[str, object]] = {}
-    variants_names: List[str] = []
-    # Build arrays including baseline as first label
-    fused_labels: List[str] = ["baseline"]
-    mel_quality_scores: List[float] = [1.0]
-    mel_latency_ms: List[float] = [mel_base_ms]
-    mel_compression: List[float] = [1.0]
-    mel_size_mb: List[float] = [float(mel_base_size) / BYTES_IN_MB]
-    # Offline compile time (host) for fused models
-    mel_compile_ms: List[float] = [_offline_compile_time_ms(input_dir / "parakeet_mel_encoder.mlpackage")]
-    jd_accuracy: List[float] = [1.0]
-    jd_latency_ms: List[float] = [jd_base_ms]
-    jd_compression: List[float] = [1.0]
-    jd_size_mb: List[float] = [float(jd_base_size) / BYTES_IN_MB]
-    jd_compile_ms: List[float] = [_offline_compile_time_ms(input_dir / "parakeet_joint_decision.mlpackage")]
-
-    # For the all-components chart, collect per-component metrics similarly
-    all_metrics: Dict[str, Dict[str, List[float]]] = {
-        "preprocessor": {
-            "quality": [1.0],
-            "compression": [1.0],
-            "latency_ms": [pre_base_ms],
-            "compile_ms": [_offline_compile_time_ms(input_dir / "parakeet_preprocessor.mlpackage")],
-            "size_mb": [float(pre_base_size) / BYTES_IN_MB],
-        },
-        "encoder": {
-            "quality": [1.0],
-            "compression": [1.0],
-            "latency_ms": [enc_base_ms],
-            "compile_ms": [_offline_compile_time_ms(input_dir / "parakeet_encoder.mlpackage")],
-            "size_mb": [float(enc_base_size) / BYTES_IN_MB],
-        },
-        "mel_encoder": {
-            "quality": [1.0],
-            "compression": [1.0],
-            "latency_ms": [mel_base_ms],
-            "compile_ms": [mel_compile_ms[0]],
-            "size_mb": [float(mel_base_size) / BYTES_IN_MB],
-        },
-        "decoder": {
-            "quality": [1.0],
-            "compression": [1.0],
-            "latency_ms": [dec_base_ms],
-            "compile_ms": [_offline_compile_time_ms(input_dir / "parakeet_decoder.mlpackage")],
-            "size_mb": [float(dec_base_size) / BYTES_IN_MB],
-        },
-        "joint": {
-            "quality": [1.0],
-            "compression": [1.0],
-            "latency_ms": [joint_base_ms],
-            "compile_ms": [_offline_compile_time_ms(input_dir / "parakeet_joint.mlpackage")],
-            "size_mb": [float(joint_base_size) / BYTES_IN_MB],
-        },
-        "joint_decision": {
-            "acc": [1.0],
-            "compression": [1.0],
-            "latency_ms": [jd_base_ms],
-            "compile_ms": [jd_compile_ms[0]],
-            "size_mb": [float(jd_base_size) / BYTES_IN_MB],
-        },
-    }
-
-    # Populate baseline entry
-    summary["baseline"] = {
-        "components": {
-            "preprocessor": {
-                "quality": 1.0,
-                "latency_ms": pre_base_ms,
-                "size_bytes": float(pre_base_size),
-                "size_mb": float(pre_base_size) / BYTES_IN_MB,
-                "compression_ratio": 1.0,
-                "compile_ms": all_metrics["preprocessor"]["compile_ms"][0],
-            },
-            "encoder": {
-                "quality": 1.0,
-                "latency_ms": enc_base_ms,
-                "size_bytes": float(enc_base_size),
-                "size_mb": float(enc_base_size) / BYTES_IN_MB,
-                "compression_ratio": 1.0,
-                "compile_ms": all_metrics["encoder"]["compile_ms"][0],
-            },
-            "mel_encoder": {
-                "quality": 1.0,
-                "latency_ms": mel_base_ms,
-                "size_bytes": float(mel_base_size),
-                "size_mb": float(mel_base_size) / BYTES_IN_MB,
-                "compression_ratio": 1.0,
-                "compile_ms": all_metrics["mel_encoder"]["compile_ms"][0],
-            },
-            "decoder": {
-                "quality": 1.0,
-                "latency_ms": dec_base_ms,
-                "size_bytes": float(dec_base_size),
-                "size_mb": float(dec_base_size) / BYTES_IN_MB,
-                "compression_ratio": 1.0,
-                "compile_ms": all_metrics["decoder"]["compile_ms"][0],
-            },
-            "joint": {
-                "quality": 1.0,
-                "latency_ms": joint_base_ms,
-                "size_bytes": float(joint_base_size),
-                "size_mb": float(joint_base_size) / BYTES_IN_MB,
-                "compression_ratio": 1.0,
-                "compile_ms": all_metrics["joint"]["compile_ms"][0],
-            },
-            "joint_decision": {
-                "acc": 1.0,
-                "latency_ms": jd_base_ms,
-                "size_bytes": float(jd_base_size),
-                "size_mb": float(jd_base_size) / BYTES_IN_MB,
-                "compression_ratio": 1.0,
-                "compile_ms": all_metrics["joint_decision"]["compile_ms"][0],
-            },
-        }
-    }
-
-    for var in variants:
-        variants_names.append(var.name)
-        out_dir = output_root / var.name
-        out_dir_exists = out_dir.exists()
-        out_dir.mkdir(parents=True, exist_ok=True)
-
-        expected_components = []
-        for comp_cfg in meta.get("components", {}).values():
-            rel = comp_cfg.get("path")
-            if rel:
-                expected_components.append(out_dir / rel)
-
-        missing = [p for p in expected_components if not p.exists()]
-        if out_dir_exists and not missing:
-            typer.echo(f"[{var.name}] Output already present at {out_dir}; skipping quantization step.")
-        else:
-            if out_dir_exists and missing:
-                missing_names = ", ".join(sorted(p.name for p in missing)) or "unknown"
-                typer.echo(f"[{var.name}] Output directory exists but is incomplete (missing: {missing_names}). Re-quantizing.")
-                shutil.rmtree(out_dir, ignore_errors=True)
-                out_dir.mkdir(parents=True, exist_ok=True)
-            _quantize_dir(input_dir, out_dir, var, component_filter)
-
-        # Load quantized models and pre-compute offline compile time
-        pre_q_path = out_dir / "parakeet_preprocessor.mlpackage"
-        enc_q_path = out_dir / "parakeet_encoder.mlpackage"
-        mel_q_path = out_dir / "parakeet_mel_encoder.mlpackage"
-        dec_q_path = out_dir / "parakeet_decoder.mlpackage"
-        joint_q_path = out_dir / "parakeet_joint.mlpackage"
-        jd_q_path = out_dir / "parakeet_joint_decision.mlpackage"
-
-        pre_compile_ms = _offline_compile_time_ms(pre_q_path)
-        enc_compile_ms = _offline_compile_time_ms(enc_q_path)
-        mel_compile_q_ms = _offline_compile_time_ms(mel_q_path)
-        dec_compile_ms = _offline_compile_time_ms(dec_q_path)
-        joint_compile_ms = _offline_compile_time_ms(joint_q_path)
-        jd_compile_q_ms = _offline_compile_time_ms(jd_q_path)
-
-        # Match compute units for quantized artifacts: preprocessor on CPU+GPU; others on CPU+NE
-        pre_q = ct.models.MLModel(str(pre_q_path), compute_units=ct.ComputeUnit.CPU_AND_GPU)
-        enc_q = ct.models.MLModel(str(enc_q_path), compute_units=ct.ComputeUnit.CPU_AND_NE)
-        mel_q = ct.models.MLModel(str(mel_q_path), compute_units=ct.ComputeUnit.CPU_AND_NE)
-        dec_q = ct.models.MLModel(str(dec_q_path), compute_units=ct.ComputeUnit.CPU_AND_NE)
-        joint_q = ct.models.MLModel(str(joint_q_path), compute_units=ct.ComputeUnit.CPU_AND_NE)
-        jd_q = ct.models.MLModel(str(jd_q_path), compute_units=ct.ComputeUnit.CPU_AND_NE)
-
-        # Preprocessor quality vs baseline
-        pre_q_out = pre_q.predict(pre_base_inputs)
-        mel_q_in = np.array(pre_q_out["mel"], dtype=np.float32, copy=True)
-        a_pre, r_pre = _max_abs_rel(mel_ref, mel_q_in)
-        l2_pre_ref = float(np.linalg.norm(mel_ref))
-        l2_pre_err = float(np.linalg.norm(mel_ref - mel_q_in))
-        pre_norm_err = (l2_pre_err / (l2_pre_ref + 1e-8)) if l2_pre_ref > 0 else 0.0
-        pre_quality = float(max(0.0, 1.0 - pre_norm_err))
-        pre_q_ms, _ = _predict_latency(pre_q, pre_base_inputs, runs=runs)
-        pre_q_size = _dir_size_bytes(out_dir / "parakeet_preprocessor.mlpackage")
-        all_metrics["preprocessor"]["quality"].append(pre_quality)
-        all_metrics["preprocessor"]["latency_ms"].append(pre_q_ms)
-        all_metrics["preprocessor"]["compression"].append(float(pre_base_size) / float(pre_q_size if pre_q_size > 0 else 1))
-        all_metrics["preprocessor"].setdefault("compile_ms", []).append(pre_compile_ms)
-        all_metrics["preprocessor"]["size_mb"].append(float(pre_q_size) / BYTES_IN_MB)
-
-        # Encoder quality vs baseline (feed baseline mel to both)
-        enc_q_out = enc_q.predict({"mel": mel_ref, "mel_length": mel_len})
-        encoder_q = np.array(enc_q_out["encoder"], dtype=np.float32, copy=True)
-        l2_enc_ref = float(np.linalg.norm(encoder_ref))
-        l2_enc_err = float(np.linalg.norm(encoder_ref - encoder_q))
-        enc_norm_err = (l2_enc_err / (l2_enc_ref + 1e-8)) if l2_enc_ref > 0 else 0.0
-        enc_quality = float(max(0.0, 1.0 - enc_norm_err))
-        enc_q_ms, _ = _predict_latency(enc_q, enc_base_inputs, runs=runs)
-        enc_q_size = _dir_size_bytes(out_dir / "parakeet_encoder.mlpackage")
-        all_metrics["encoder"]["quality"].append(enc_quality)
-        all_metrics["encoder"]["latency_ms"].append(enc_q_ms)
-        all_metrics["encoder"]["compression"].append(float(enc_base_size) / float(enc_q_size if enc_q_size > 0 else 1))
-        all_metrics["encoder"].setdefault("compile_ms", []).append(enc_compile_ms)
-        all_metrics["encoder"]["size_mb"].append(float(enc_q_size) / BYTES_IN_MB)
-
-        # MelEncoder quality
-        mel_q_out = mel_q.predict(mel_base_inputs)
-        enc_q_fused = np.array(mel_q_out["encoder"], dtype=np.float32, copy=True)
-        a_mel, r_mel = _max_abs_rel(encoder_ref, enc_q_fused)
-        # Normalize error into a [0,1] quality score: 1 / (1 + normalized L2)
-        # Use relative measure derived from L2 norms as more stable than max.
-        l2_ref = float(np.linalg.norm(encoder_ref))
-        l2_err = float(np.linalg.norm(encoder_ref - enc_q_fused))
-        norm_err = (l2_err / (l2_ref + 1e-8)) if l2_ref > 0 else 0.0
-        mel_quality = float(max(0.0, 1.0 - norm_err))
-        mel_q_ms, _ = _predict_latency(mel_q, mel_base_inputs, runs=runs)
-        mel_q_size = _dir_size_bytes(out_dir / "parakeet_mel_encoder.mlpackage")
-        mel_ratio = float(mel_base_size) / float(mel_q_size if mel_q_size > 0 else 1)
-        mel_size_mb.append(float(mel_q_size) / BYTES_IN_MB)
-        all_metrics["mel_encoder"]["quality"].append(mel_quality)
-        all_metrics["mel_encoder"]["latency_ms"].append(mel_q_ms)
-        all_metrics["mel_encoder"]["compression"].append(float(mel_base_size) / float(mel_q_size if mel_q_size > 0 else 1))
-        all_metrics["mel_encoder"].setdefault("compile_ms", []).append(mel_compile_q_ms)
-        all_metrics["mel_encoder"]["size_mb"].append(float(mel_q_size) / BYTES_IN_MB)
-
-        # JointDecision quality: token-id and duration match rates
-        jd_base_out = joint_decision_base.predict(jd_base_inputs)
-        token_id_base = np.array(jd_base_out["token_id"], dtype=np.int32, copy=True)
-        duration_base = np.array(jd_base_out["duration"], dtype=np.int32, copy=True)
-        token_prob_base = np.array(jd_base_out["token_prob"], dtype=np.float32, copy=True)
-
-        jd_q_out = jd_q.predict(jd_base_inputs)
-        token_id_q = np.array(jd_q_out["token_id"], dtype=np.int32, copy=True)
-        duration_q = np.array(jd_q_out["duration"], dtype=np.int32, copy=True)
-        token_prob_q = np.array(jd_q_out["token_prob"], dtype=np.float32, copy=True)
-
-        # Accuracy metrics
-        id_match = float((token_id_q == token_id_base).mean())
-        dur_match = float((duration_q == duration_base).mean())
-        # Aggregate a single "accuracy" number as token-id match rate (primary)
-        jd_acc = id_match
-        jd_q_ms, _ = _predict_latency(jd_q, jd_base_inputs, runs=runs)
-        jd_q_size = _dir_size_bytes(out_dir / "parakeet_joint_decision.mlpackage")
-        jd_ratio = float(jd_base_size) / float(jd_q_size if jd_q_size > 0 else 1)
-        jd_size_mb.append(float(jd_q_size) / BYTES_IN_MB)
-        all_metrics["joint_decision"].setdefault("acc", []).append(jd_acc)
-        all_metrics["joint_decision"]["latency_ms"].append(jd_q_ms)
-        all_metrics["joint_decision"]["compression"].append(float(jd_base_size) / float(jd_q_size if jd_q_size > 0 else 1))
-        all_metrics["joint_decision"].setdefault("compile_ms", []).append(jd_compile_q_ms)
-        all_metrics["joint_decision"]["size_mb"].append(float(jd_q_size) / BYTES_IN_MB)
-
-        # Decoder quality vs baseline
-        dec_q_out = dec_q.predict(dec_base_inputs)
-        decoder_q = np.array(dec_q_out["decoder"], dtype=np.float32, copy=True)
-        l2_dec_ref = float(np.linalg.norm(decoder_ref))
-        l2_dec_err = float(np.linalg.norm(decoder_ref - decoder_q))
-        dec_norm_err = (l2_dec_err / (l2_dec_ref + 1e-8)) if l2_dec_ref > 0 else 0.0
-        dec_quality = float(max(0.0, 1.0 - dec_norm_err))
-        dec_q_ms, _ = _predict_latency(dec_q, dec_base_inputs, runs=runs)
-        dec_q_size = _dir_size_bytes(out_dir / "parakeet_decoder.mlpackage")
-        all_metrics["decoder"]["quality"].append(dec_quality)
-        all_metrics["decoder"]["latency_ms"].append(dec_q_ms)
-        all_metrics["decoder"]["compression"].append(float(dec_base_size) / float(dec_q_size if dec_q_size > 0 else 1))
-        all_metrics["decoder"].setdefault("compile_ms", []).append(dec_compile_ms)
-        all_metrics["decoder"]["size_mb"].append(float(dec_q_size) / BYTES_IN_MB)
-
-        # Joint quality vs baseline (compare logits)
-        joint_q_out = joint_q.predict(joint_base_inputs)
-        logits_q = np.array(joint_q_out["logits"], dtype=np.float32, copy=True)
-        l2_joint_ref = float(np.linalg.norm(logits_base))
-        l2_joint_err = float(np.linalg.norm(logits_base - logits_q))
-        joint_norm_err = (l2_joint_err / (l2_joint_ref + 1e-8)) if l2_joint_ref > 0 else 0.0
-        joint_quality = float(max(0.0, 1.0 - joint_norm_err))
-        joint_q_ms, _ = _predict_latency(joint_q, joint_base_inputs, runs=runs)
-        joint_q_size = _dir_size_bytes(out_dir / "parakeet_joint.mlpackage")
-        all_metrics["joint"]["quality"].append(joint_quality)
-        all_metrics["joint"]["latency_ms"].append(joint_q_ms)
-        all_metrics["joint"]["compression"].append(float(joint_base_size) / float(joint_q_size if joint_q_size > 0 else 1))
-        all_metrics["joint"].setdefault("compile_ms", []).append(joint_compile_ms)
-        all_metrics["joint"]["size_mb"].append(float(joint_q_size) / BYTES_IN_MB)
-
-        # Decoder deltas for JSON
-        a_dec, r_dec = _max_abs_rel(decoder_ref, decoder_q)
-        # Joint deltas for JSON
-        a_joint, r_joint = _max_abs_rel(logits_base, logits_q)
-
-        # Store metrics
-        summary[var.name] = {
-            "components": {
-                "preprocessor": {
-                    "quality": pre_quality,
-                    "latency_ms": pre_q_ms,
-                    "size_bytes": float(pre_q_size),
-                    "size_mb": float(pre_q_size) / BYTES_IN_MB,
-                    "compression_ratio": float(pre_base_size) / float(pre_q_size if pre_q_size > 0 else 1),
-                    "max_abs": a_pre,
-                    "max_rel": r_pre,
-                    "compile_ms": pre_compile_ms,
-                },
-                "encoder": {
-                    "quality": enc_quality,
-                    "latency_ms": enc_q_ms,
-                    "size_bytes": float(enc_q_size),
-                    "size_mb": float(enc_q_size) / BYTES_IN_MB,
-                    "compression_ratio": float(enc_base_size) / float(enc_q_size if enc_q_size > 0 else 1),
-                    "compile_ms": enc_compile_ms,
-                },
-                "mel_encoder": {
-                    "quality": mel_quality,
-                    "latency_ms": mel_q_ms,
-                    "size_bytes": float(mel_q_size),
-                    "size_mb": float(mel_q_size) / BYTES_IN_MB,
-                    "compression_ratio": mel_ratio,
-                    "max_abs": a_mel,
-                    "max_rel": r_mel,
-                    "compile_ms": mel_compile_q_ms,
-                },
-                "decoder": {
-                    "quality": dec_quality,
-                    "latency_ms": dec_q_ms,
-                    "size_bytes": float(dec_q_size),
-                    "size_mb": float(dec_q_size) / BYTES_IN_MB,
-                    "compression_ratio": float(dec_base_size) / float(dec_q_size if dec_q_size > 0 else 1),
-                    "max_abs": a_dec,
-                    "max_rel": r_dec,
-                    "compile_ms": dec_compile_ms,
-                },
-                "joint": {
-                    "quality": joint_quality,
-                    "latency_ms": joint_q_ms,
-                    "size_bytes": float(joint_q_size),
-                    "size_mb": float(joint_q_size) / BYTES_IN_MB,
-                    "compression_ratio": float(joint_base_size) / float(joint_q_size if joint_q_size > 0 else 1),
-                    "max_abs": a_joint,
-                    "max_rel": r_joint,
-                    "compile_ms": joint_compile_ms,
-                },
-                "joint_decision": {
-                    "acc": jd_acc,
-                    "duration_match": dur_match,
-                    "prob_mae": float(np.mean(np.abs(token_prob_q - token_prob_base))),
-                    "latency_ms": jd_q_ms,
-                    "size_bytes": float(jd_q_size),
-                    "size_mb": float(jd_q_size) / BYTES_IN_MB,
-                    "compression_ratio": jd_ratio,
-                    "compile_ms": jd_compile_q_ms,
-                },
-            }
-        }
-
-        fused_labels.append(var.name)
-        mel_quality_scores.append(mel_quality)
-        mel_latency_ms.append(mel_q_ms)
-        mel_compression.append(mel_ratio)
-        jd_accuracy.append(jd_acc)
-        jd_latency_ms.append(jd_q_ms)
-        jd_compression.append(jd_ratio)
-        mel_compile_ms.append(mel_compile_q_ms)
-        jd_compile_ms.append(jd_compile_q_ms)
-
-    # Write summary JSON
-    out_root = output_root
-    out_root.mkdir(parents=True, exist_ok=True)
-    (out_root / "quantization_summary.json").write_text(json.dumps(summary, indent=2))
-
-    # Plot
-    plot_dir = out_root / "plots"
-    title_suffix = _chip_spec_string(compute_units)
-    fused_paths = _plot_fused_category_charts(
-        plot_dir,
-        fused_labels,
-        mel_quality_scores,
-        mel_latency_ms,
-        mel_compression,
-        mel_size_mb,
-        jd_accuracy,
-        jd_latency_ms,
-        jd_compression,
-        jd_size_mb,
-        mel_compile_ms,
-        jd_compile_ms,
-        title_suffix,
-    )
-    component_paths = _plot_all_component_category_charts(
-        plot_dir,
-        fused_labels,
-        all_metrics,
-        title_suffix,
-    )
-
-    typer.echo(f"Wrote summary JSON: {out_root / 'quantization_summary.json'}")
-    if HAS_MPL:
-        all_plot_paths = fused_paths + component_paths
-        repo_plot_dir = BASE_DIR / "plots" / "quantize" / compute_units.lower()
-        repo_plot_dir.mkdir(parents=True, exist_ok=True)
-        for path in all_plot_paths:
-            typer.echo(f"Wrote plot: {path}")
-            if path.exists():
-                dest = repo_plot_dir / path.name
-                shutil.copy2(path, dest)
-                typer.echo(f"Mirrored plot: {dest}")
-    else:
-        typer.echo("matplotlib unavailable; skipped plotting.")
-
-
-if __name__ == "__main__":
-    app()

convert/parakeet-tdt-v2-0.6b/coreml/speech_to_text_streaming_infer_rnnt.py

@@ -1,341 +0,0 @@
-"""Streaming inference helper that stitches Parakeet CoreML components using the RNNT greedy loop from Nemo."""
-from __future__ import annotations
-
-import argparse
-import json
-import logging
-from pathlib import Path
-from typing import Iterable, List, Optional, Sequence
-
-import coremltools as ct
-import librosa
-import numpy as np
-import torch
-
-from parakeet_components import CoreMLModelBundle
-
-LOGGER = logging.getLogger("parakeet_streaming")
-
-
-class BatchedHyps:
-    """Minimal port of Nemo's batched hypothesis buffer."""
-
-    def __init__(
-        self,
-        batch_size: int,
-        init_length: int,
-        device: torch.device,
-        float_dtype: torch.dtype,
-    ) -> None:
-        if init_length <= 0:
-            raise ValueError("init_length must be > 0")
-        self._max_length = init_length
-        self.current_lengths = torch.zeros(batch_size, device=device, dtype=torch.long)
-        self.transcript = torch.zeros((batch_size, self._max_length), device=device, dtype=torch.long)
-        self.timestamps = torch.zeros((batch_size, self._max_length), device=device, dtype=torch.long)
-        self.scores = torch.zeros(batch_size, device=device, dtype=float_dtype)
-        self.last_timestamp = torch.full((batch_size,), -1, device=device, dtype=torch.long)
-        self.last_timestamp_lasts = torch.zeros(batch_size, device=device, dtype=torch.long)
-        self._batch_indices = torch.arange(batch_size, device=device)
-        self._ones = torch.ones_like(self._batch_indices)
-
-    def add_results(
-        self,
-        active_mask: torch.Tensor,
-        labels: torch.Tensor,
-        time_indices: torch.Tensor,
-        scores: torch.Tensor,
-    ) -> None:
-        self.scores = torch.where(active_mask, self.scores + scores, self.scores)
-        self.transcript[self._batch_indices, self.current_lengths] = labels
-        self.timestamps[self._batch_indices, self.current_lengths] = time_indices
-        torch.where(
-            torch.logical_and(active_mask, self.last_timestamp == time_indices),
-            self.last_timestamp_lasts + 1,
-            self.last_timestamp_lasts,
-            out=self.last_timestamp_lasts,
-        )
-        torch.where(
-            torch.logical_and(active_mask, self.last_timestamp != time_indices),
-            self._ones,
-            self.last_timestamp_lasts,
-            out=self.last_timestamp_lasts,
-        )
-        torch.where(active_mask, time_indices, self.last_timestamp, out=self.last_timestamp)
-        self.current_lengths += active_mask
-
-
-class CoreMLStreamingDecoder:
-    """Use exported decoder and joint CoreML models with Nemo's greedy RNNT loop."""
-
-    def __init__(
-        self,
-        decoder_model: ct.models.MLModel,
-        joint_model: ct.models.MLModel,
-        *,
-        vocab_size: int,
-        blank_id: int,
-        num_layers: int,
-        hidden_size: int,
-        durations: Sequence[int] = (0, 1, 2, 3, 4),
-        max_symbols: int = 10,
-        device: torch.device = torch.device("cpu"),
-    ) -> None:
-        self.decoder_model = decoder_model
-        self.joint_model = joint_model
-        self.vocab_size = vocab_size
-        self.blank_id = blank_id
-        self.num_layers = num_layers
-        self.hidden_size = hidden_size
-        self.durations = torch.tensor(durations, dtype=torch.long, device=device)
-        self.max_symbols = max_symbols
-        self.device = device
-
-    def _predict_decoder(self, labels: torch.Tensor, h_in: np.ndarray, c_in: np.ndarray) -> tuple[torch.Tensor, np.ndarray, np.ndarray]:
-        outputs = self.decoder_model.predict(
-            {
-                "targets": np.array([labels.cpu().numpy()], dtype=np.int32),
-                "target_lengths": np.array([labels.numel()], dtype=np.int32),
-                "h_in": h_in,
-                "c_in": c_in,
-            }
-        )
-        decoder_output = torch.from_numpy(outputs["decoder_output"]).to(self.device)
-        return decoder_output, outputs["h_out"], outputs["c_out"]
-
-    def _predict_joint(self, encoder_frame: torch.Tensor, decoder_output: torch.Tensor) -> torch.Tensor:
-        outputs = self.joint_model.predict(
-            {
-                "encoder_outputs": encoder_frame.unsqueeze(1).cpu().numpy().astype(np.float32),
-                "decoder_outputs": decoder_output.cpu().numpy().astype(np.float32),
-            }
-        )
-        logits = torch.from_numpy(outputs["logits"]).to(self.device)
-        return logits.squeeze(1).squeeze(1)
-
-    def decode(self, encoder_output: torch.Tensor, encoder_lengths: torch.Tensor) -> List[List[int]]:
-        batch_size, max_time, _ = encoder_output.shape
-        encoder_output = encoder_output.to(self.device)
-        encoder_lengths = encoder_lengths.to(self.device)
-
-        float_dtype = encoder_output.dtype
-        batch_indices = torch.arange(batch_size, device=self.device)
-        labels = torch.full((batch_size,), fill_value=self.blank_id, device=self.device, dtype=torch.long)
-        time_indices = torch.zeros_like(labels)
-        safe_time_indices = torch.zeros_like(labels)
-        time_indices_current = torch.zeros_like(labels)
-        last_timesteps = encoder_lengths - 1
-        active_mask = encoder_lengths > 0
-        advance_mask = torch.empty_like(active_mask)
-        active_mask_prev = torch.empty_like(active_mask)
-        became_inactive = torch.empty_like(active_mask)
-
-        hyps = BatchedHyps(
-            batch_size=batch_size,
-            init_length=max_time * self.max_symbols if self.max_symbols else max_time,
-            device=self.device,
-            float_dtype=float_dtype,
-        )
-
-        h_in = np.zeros((self.num_layers, batch_size, self.hidden_size), dtype=np.float32)
-        c_in = np.zeros((self.num_layers, batch_size, self.hidden_size), dtype=np.float32)
-
-        while active_mask.any():
-            active_mask_prev.copy_(active_mask)
-            decoder_output, h_in, c_in = self._predict_decoder(labels, h_in, c_in)
-            logits = self._predict_joint(encoder_output[batch_indices, safe_time_indices], decoder_output)
-
-            scores, labels = logits[:, : self.vocab_size].max(dim=-1)
-            duration_indices = logits[:, self.vocab_size : self.vocab_size + len(self.durations)].argmax(dim=-1)
-            durations = self.durations[duration_indices]
-
-            blank_mask = labels == self.blank_id
-            durations.masked_fill_(torch.logical_and(durations == 0, blank_mask), 1)
-            time_indices_current.copy_(time_indices)
-            time_indices += durations
-            torch.minimum(time_indices, last_timesteps, out=safe_time_indices)
-            torch.less(time_indices, encoder_lengths, out=active_mask)
-            torch.logical_and(active_mask, blank_mask, out=advance_mask)
-
-            while advance_mask.any():
-                torch.where(advance_mask, time_indices, time_indices_current, out=time_indices_current)
-                logits = self._predict_joint(encoder_output[batch_indices, safe_time_indices], decoder_output)
-                more_scores, more_labels = logits[:, : self.vocab_size].max(dim=-1)
-                labels = torch.where(advance_mask, more_labels, labels)
-                scores = torch.where(advance_mask, more_scores, scores)
-                duration_indices = logits[:, self.vocab_size : self.vocab_size + len(self.durations)].argmax(dim=-1)
-                durations = self.durations[duration_indices]
-                blank_mask = labels == self.blank_id
-                durations.masked_fill_(torch.logical_and(durations == 0, blank_mask), 1)
-                torch.where(advance_mask, time_indices + durations, time_indices, out=time_indices)
-                torch.minimum(time_indices, last_timesteps, out=safe_time_indices)
-                torch.less(time_indices, encoder_lengths, out=active_mask)
-                torch.logical_and(active_mask, blank_mask, out=advance_mask)
-
-            torch.ne(active_mask, active_mask_prev, out=became_inactive)
-            hyps.add_results(active_mask, labels, time_indices_current, scores)
-
-            if self.max_symbols is not None:
-                force_blank = torch.logical_and(
-                    active_mask,
-                    torch.logical_and(
-                        torch.logical_and(labels != self.blank_id, hyps.last_timestamp_lasts >= self.max_symbols),
-                        hyps.last_timestamp == time_indices,
-                    ),
-                )
-                time_indices += force_blank
-                torch.minimum(time_indices, last_timesteps, out=safe_time_indices)
-                torch.less(time_indices, encoder_lengths, out=active_mask)
-
-        results: List[List[int]] = []
-        for hyp in hyps.transcript:
-            tokens = [int(token) for token in hyp.tolist() if 0 < token < self.vocab_size]
-            results.append(tokens)
-        return results
-
-
-class StreamingTranscriber:
-    def __init__(
-        self,
-        bundle: CoreMLModelBundle,
-        *,
-        blank_id: Optional[int] = None,
-        num_layers: int = 2,
-        hidden_size: int = 640,
-        durations: Sequence[int] = (0, 1, 2, 3, 4),
-    ) -> None:
-        self.preprocessor = ct.models.MLModel(str(bundle.preprocessor), compute_units=ct.ComputeUnit.CPU_ONLY)
-        self.encoder = ct.models.MLModel(str(bundle.encoder), compute_units=ct.ComputeUnit.CPU_ONLY)
-        self.decoder = ct.models.MLModel(str(bundle.decoder), compute_units=ct.ComputeUnit.CPU_ONLY)
-        self.joint = ct.models.MLModel(str(bundle.joint), compute_units=ct.ComputeUnit.CPU_ONLY)
-        self.tokenizer = self._load_tokenizer(bundle.tokenizer)
-
-        vocab_size = max(self.tokenizer.keys()) + 1
-        if blank_id is None:
-            blank_id = vocab_size - 1
-        self.decoder_helper = CoreMLStreamingDecoder(
-            self.decoder,
-            self.joint,
-            vocab_size=vocab_size,
-            blank_id=blank_id,
-            num_layers=num_layers,
-            hidden_size=hidden_size,
-            durations=durations,
-        )
-        self.blank_id = blank_id
-
-    @staticmethod
-    def _load_tokenizer(tokenizer_path: Optional[Path]) -> dict[int, str]:
-        if tokenizer_path is None:
-            raise ValueError("Tokenizer JSON is required")
-        with Path(tokenizer_path).open() as f:
-            data = json.load(f)
-        return {int(k): v for k, v in data.items()}
-
-    def _tokens_to_text(self, tokens: Iterable[int]) -> str:
-        pieces: List[str] = []
-        for token in tokens:
-            piece = self.tokenizer.get(token)
-            if piece is None:
-                continue
-            if piece.startswith("▁"):
-                if pieces:
-                    pieces.append(" ")
-                pieces.append(piece[1:])
-            else:
-                pieces.append(piece)
-        return "".join(pieces).strip()
-
-    def _preprocess(self, audio: np.ndarray) -> tuple[np.ndarray, int]:
-        audio_2d = audio.reshape(1, -1).astype(np.float32)
-        length = np.array([audio_2d.shape[-1]], dtype=np.int32)
-        outputs = self.preprocessor.predict({
-            "audio_signal": audio_2d,
-            "audio_length": length,
-        })
-        return outputs["melspectrogram"], int(outputs["melspectrogram_length"][0])
-
-    def _encode(self, mel: np.ndarray, mel_length: int) -> tuple[torch.Tensor, torch.Tensor]:
-        outputs = self.encoder.predict({
-            "melspectrogram": mel.astype(np.float32),
-            "melspectrogram_length": np.array([mel_length], dtype=np.int32),
-        })
-        encoder_output = outputs["encoder_output"]
-        if encoder_output.ndim == 3:
-            encoder_output = np.transpose(encoder_output, (0, 2, 1))
-        length = torch.tensor(outputs["encoder_output_length"], dtype=torch.long)
-        return torch.from_numpy(encoder_output.astype(np.float32)), length
-
-    def transcribe(self, audio_path: Path) -> str:
-        audio, _ = librosa.load(str(audio_path), sr=16000)
-        mel, mel_length = self._preprocess(audio)
-        encoder_output, encoder_length = self._encode(mel, mel_length)
-        token_ids = self.decoder_helper.decode(encoder_output, encoder_length)[0]
-        return self._tokens_to_text(token_ids)
-
-    def transcribe_many(self, audio_paths: Sequence[Path]) -> List[str]:
-        results: List[str] = []
-        for path in audio_paths:
-            LOGGER.info("Transcribing %s", path)
-            results.append(self.transcribe(path))
-        return results
-
-
-def _resolve_bundle(args: argparse.Namespace) -> CoreMLModelBundle:
-    base = Path(args.model_dir) if args.model_dir else None
-    if base is None and not all([args.preprocessor, args.encoder, args.decoder, args.joint, args.tokenizer]):
-        raise ValueError("Either --model-dir or explicit model paths are required")
-    return CoreMLModelBundle(
-        preprocessor=Path(args.preprocessor) if args.preprocessor else base / "Melspectrogram.mlpackage",
-        encoder=Path(args.encoder) if args.encoder else base / "ParakeetEncoder.mlpackage",
-        decoder=Path(args.decoder) if args.decoder else base / "ParakeetDecoder.mlpackage",
-        joint=Path(args.joint) if args.joint else base / "RNNTJoint.mlpackage",
-        tokenizer=Path(args.tokenizer) if args.tokenizer else base / "tokenizer.json",
-    )
-
-
-def _build_parser() -> argparse.ArgumentParser:
-    parser = argparse.ArgumentParser(description="Streaming RNNT inference with CoreML components")
-    parser.add_argument("--model-dir", type=Path, help="Directory containing exported CoreML models")
-    parser.add_argument("--preprocessor", type=Path, help="Path to the preprocessor .mlpackage")
-    parser.add_argument("--encoder", type=Path, help="Path to the encoder .mlpackage")
-    parser.add_argument("--decoder", type=Path, help="Path to the decoder .mlpackage")
-    parser.add_argument("--joint", type=Path, help="Path to the joint .mlpackage")
-    parser.add_argument("--tokenizer", type=Path, help="Path to tokenizer JSON")
-    parser.add_argument("audio", nargs="+", help="Audio files to transcribe")
-    parser.add_argument("--blank-id", type=int, help="Blank token id")
-    parser.add_argument("--num-layers", type=int, default=2, help="Prediction network layer count")
-    parser.add_argument("--hidden-size", type=int, default=640, help="Prediction network hidden size")
-    parser.add_argument("--durations", type=int, nargs="+", default=[0, 1, 2, 3, 4], help="RNNT duration bucket values")
-    parser.add_argument("--verbose", "-v", action="count", default=0, help="Increase log verbosity")
-    return parser
-
-
-def _configure_logging(verbosity: int) -> None:
-    level = logging.WARNING - (10 * verbosity)
-    logging.basicConfig(level=max(logging.DEBUG, level), format="[%(levelname)s] %(message)s")
-
-
-def main(argv: Optional[Sequence[str]] = None) -> None:
-    parser = _build_parser()
-    args = parser.parse_args(argv)
-    _configure_logging(args.verbose)
-
-    try:
-        bundle = _resolve_bundle(args)
-        transcriber = StreamingTranscriber(
-            bundle,
-            blank_id=args.blank_id,
-            num_layers=args.num_layers,
-            hidden_size=args.hidden_size,
-            durations=tuple(args.durations),
-        )
-        transcripts = transcriber.transcribe_many([Path(p) for p in args.audio])
-        for path, text in zip(args.audio, transcripts):
-            print(f"{path}: {text}")
-    except ValueError as exc:
-        parser.error(str(exc))
-
-
-if __name__ == "__main__":  # pragma: no cover
-    main()

scripts/run_benchmarks.py

@@ -1,273 +0,0 @@
-#!/usr/bin/env python3
-"""
-FluidAudio Benchmark Suite
-
-Runs ASR, VAD, and Diarization benchmarks and saves results to JSON.
-Compare results against Documentation/Benchmarks.md baselines.
-
-Usage:
-    python run_benchmarks.py              # Run all benchmarks
-    python run_benchmarks.py --quick      # Quick smoke test
-    python run_benchmarks.py --asr-only   # ASR benchmark only
-    python run_benchmarks.py --vad-only   # VAD benchmark only
-    python run_benchmarks.py --diar-only  # Diarization only
-"""
-
-import argparse
-import json
-import subprocess
-import sys
-from datetime import datetime
-from pathlib import Path
-
-
-# Baseline values from Documentation/Benchmarks.md
-BASELINES = {
-    "asr": {
-        "wer_percent": 5.8,
-        "rtfx_min": 200,  # M4 Pro: ~210x
-        "description": "LibriSpeech test-clean, Parakeet TDT 0.6B"
-    },
-    "vad": {
-        "f1_percent": 85.0,
-        "rtfx_min": 500,
-        "description": "VOiCES dataset, Silero VAD"
-    },
-    "diarization": {
-        "der_percent": 17.7,
-        "rtfx_min": 1.0,
-        "description": "AMI SDM, pyannote-based"
-    }
-}
-
-
-def run_command(cmd: list[str], output_file: Path | None = None) -> tuple[int, str]:
-    """Run a command and optionally save output."""
-    print(f"Running: {' '.join(cmd)}")
-
-    result = subprocess.run(
-        cmd,
-        capture_output=True,
-        text=True
-    )
-
-    output = result.stdout + result.stderr
-
-    if output_file:
-        output_file.write_text(output)
-
-    return result.returncode, output
-
-
-def build_release() -> bool:
-    """Build the project in release mode."""
-    print("\n" + "=" * 60)
-    print("Building release...")
-    print("=" * 60)
-
-    returncode, _ = run_command(["swift", "build", "-c", "release"])
-
-    if returncode != 0:
-        print("ERROR: Build failed!")
-        return False
-
-    print("Build successful.")
-    return True
-
-
-def run_asr_benchmark(output_dir: Path, quick: bool = False) -> dict | None:
-    """Run ASR benchmark on LibriSpeech test-clean."""
-    print("\n" + "=" * 60)
-    print("ASR Benchmark (LibriSpeech test-clean)")
-    print("=" * 60)
-
-    max_files = "100" if quick else "all"
-    output_json = output_dir / f"asr_results.json"
-
-    cmd = [
-        "swift", "run", "-c", "release", "fluidaudio", "asr-benchmark",
-        "--subset", "test-clean",
-        "--max-files", max_files,
-        "--output", str(output_json)
-    ]
-
-    returncode, output = run_command(cmd, output_dir / "asr_log.txt")
-
-    if returncode != 0:
-        print(f"ERROR: ASR benchmark failed!")
-        return None
-
-    if output_json.exists():
-        return json.loads(output_json.read_text())
-
-    return None
-
-
-def run_vad_benchmark(output_dir: Path, quick: bool = False) -> dict | None:
-    """Run VAD benchmark."""
-    print("\n" + "=" * 60)
-    print("VAD Benchmark")
-    print("=" * 60)
-
-    dataset = "mini50" if quick else "voices-subset"
-    output_json = output_dir / f"vad_results.json"
-
-    cmd = [
-        "swift", "run", "-c", "release", "fluidaudio", "vad-benchmark",
-        "--dataset", dataset,
-        "--all-files",
-        "--threshold", "0.5",
-        "--output", str(output_json)
-    ]
-
-    returncode, output = run_command(cmd, output_dir / "vad_log.txt")
-
-    if returncode != 0:
-        print(f"ERROR: VAD benchmark failed!")
-        return None
-
-    if output_json.exists():
-        return json.loads(output_json.read_text())
-
-    return None
-
-
-def run_diarization_benchmark(output_dir: Path, quick: bool = False) -> dict | None:
-    """Run diarization benchmark on AMI SDM."""
-    print("\n" + "=" * 60)
-    print("Diarization Benchmark (AMI SDM)")
-    print("=" * 60)
-
-    output_json = output_dir / f"diarization_results.json"
-
-    cmd = [
-        "swift", "run", "-c", "release", "fluidaudio", "diarization-benchmark",
-        "--auto-download",
-        "--output", str(output_json)
-    ]
-
-    if quick:
-        cmd.extend(["--single-file", "ES2004a"])
-
-    returncode, output = run_command(cmd, output_dir / "diarization_log.txt")
-
-    if returncode != 0:
-        print(f"ERROR: Diarization benchmark failed!")
-        return None
-
-    if output_json.exists():
-        return json.loads(output_json.read_text())
-
-    return None
-
-
-def compare_results(results: dict) -> None:
-    """Compare results against baselines."""
-    print("\n" + "=" * 60)
-    print("Results vs Baselines (Documentation/Benchmarks.md)")
-    print("=" * 60)
-
-    if "asr" in results and results["asr"]:
-        asr = results["asr"]
-        baseline = BASELINES["asr"]
-        wer = asr.get("wer", asr.get("average_wer", 0)) * 100
-        rtfx = asr.get("rtfx", asr.get("median_rtfx", 0))
-
-        wer_status = "✓" if wer <= baseline["wer_percent"] * 1.1 else "✗"
-        rtfx_status = "✓" if rtfx >= baseline["rtfx_min"] * 0.8 else "✗"
-
-        print(f"\nASR ({baseline['description']}):")
-        print(f"  WER:  {wer:.1f}% (baseline: {baseline['wer_percent']}%) {wer_status}")
-        print(f"  RTFx: {rtfx:.1f}x (baseline: {baseline['rtfx_min']}x+) {rtfx_status}")
-
-    if "vad" in results and results["vad"]:
-        vad = results["vad"]
-        baseline = BASELINES["vad"]
-        f1 = vad.get("f1_score", 0)
-        rtfx = vad.get("rtfx", 0)
-
-        f1_status = "✓" if f1 >= baseline["f1_percent"] * 0.9 else "✗"
-        rtfx_status = "✓" if rtfx >= baseline["rtfx_min"] * 0.5 else "✗"
-
-        print(f"\nVAD ({baseline['description']}):")
-        print(f"  F1:   {f1:.1f}% (baseline: {baseline['f1_percent']}%+) {f1_status}")
-        print(f"  RTFx: {rtfx:.1f}x (baseline: {baseline['rtfx_min']}x+) {rtfx_status}")
-
-    if "diarization" in results and results["diarization"]:
-        diar = results["diarization"]
-        baseline = BASELINES["diarization"]
-        der = diar.get("der", diar.get("average_der", 0)) * 100
-        rtfx = diar.get("rtfx", diar.get("average_rtfx", 0))
-
-        der_status = "✓" if der <= baseline["der_percent"] * 1.2 else "✗"
-        rtfx_status = "✓" if rtfx >= baseline["rtfx_min"] else "✗"
-
-        print(f"\nDiarization ({baseline['description']}):")
-        print(f"  DER:  {der:.1f}% (baseline: {baseline['der_percent']}%) {der_status}")
-        print(f"  RTFx: {rtfx:.1f}x (baseline: {baseline['rtfx_min']}x+) {rtfx_status}")
-
-
-def main():
-    parser = argparse.ArgumentParser(description="FluidAudio Benchmark Suite")
-    parser.add_argument("--quick", action="store_true", help="Quick smoke test with smaller datasets")
-    parser.add_argument("--asr-only", action="store_true", help="Run ASR benchmark only")
-    parser.add_argument("--vad-only", action="store_true", help="Run VAD benchmark only")
-    parser.add_argument("--diar-only", action="store_true", help="Run diarization benchmark only")
-    parser.add_argument("--output-dir", type=str, help="Output directory for results")
-    args = parser.parse_args()
-
-    # Determine which benchmarks to run
-    run_all = not (args.asr_only or args.vad_only or args.diar_only)
-
-    # Setup output directory
-    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
-    if args.output_dir:
-        output_dir = Path(args.output_dir)
-    else:
-        output_dir = Path("benchmark-results") / timestamp
-
-    output_dir.mkdir(parents=True, exist_ok=True)
-
-    print("=" * 60)
-    print("FluidAudio Benchmark Suite")
-    print("=" * 60)
-    print(f"Mode: {'Quick' if args.quick else 'Full'}")
-    print(f"Output: {output_dir}")
-    print(f"Time: {timestamp}")
-
-    # Build first
-    if not build_release():
-        sys.exit(1)
-
-    results = {}
-
-    # Run benchmarks
-    if run_all or args.asr_only:
-        results["asr"] = run_asr_benchmark(output_dir, args.quick)
-
-    if run_all or args.vad_only:
-        results["vad"] = run_vad_benchmark(output_dir, args.quick)
-
-    if run_all or args.diar_only:
-        results["diarization"] = run_diarization_benchmark(output_dir, args.quick)
-
-    # Save combined results
-    combined_output = output_dir / "benchmark_results.json"
-    combined_output.write_text(json.dumps({
-        "timestamp": timestamp,
-        "mode": "quick" if args.quick else "full",
-        "baselines": BASELINES,
-        "results": results
-    }, indent=2))
-
-    # Compare against baselines
-    compare_results(results)
-
-    print("\n" + "=" * 60)
-    print("Benchmark complete!")
-    print("=" * 60)
-    print(f"Results saved to: {combined_output}")
-
-
-if __name__ == "__main__":
-    main()