export_decoder_fp32.py

from __future__ import annotations

import hashlib
import json
import sys
from collections import Counter
from pathlib import Path
from typing import Iterable, Sequence

import numpy as np
import onnx
import onnxruntime as ort
import torch
from safetensors.torch import load_file as load_safetensors


ROOT = Path(__file__).resolve().parent
MODEL_DIR = ROOT / "model"
OUTPUT_DIR = ROOT / "output"
REPORT_DIR = ROOT / "reports"
OPSET = 17
NUM_LAYERS = 6
NUM_KV_HEADS = 2
HEAD_DIM = 64
VISION_TOKENS = 256
HIDDEN_SIZE = 512
VOCAB_SIZE = 14_630
QUANTIZED_ONLY_OPS = {"DynamicQuantizeLinear", "MatMulInteger"}


def vectorized_causal_mask(
    q_len: int,
    kv_len: int,
    device: torch.device,
    dtype: torch.dtype,
) -> torch.Tensor:
    """Equivalent to Baberu's row-assignment loop without ScatterND expansion."""
    query_positions = torch.arange(q_len, device=device).unsqueeze(1) + (kv_len - q_len)
    key_positions = torch.arange(kv_len, device=device).unsqueeze(0)
    zero = torch.zeros((), device=device, dtype=dtype)
    blocked = torch.full((), float("-inf"), device=device, dtype=dtype)
    return torch.where(key_positions <= query_positions, zero, blocked)


def webgpu_rms_norm(self, value: torch.Tensor) -> torch.Tensor:
    """RMSNorm written without Pow(x, 2), which fails in ORT WebGPU 1.27 here."""
    input_dtype = value.dtype
    value_f32 = value.to(torch.float32)
    variance = (value_f32 * value_f32).mean(-1, keepdim=True)
    normalized = value_f32 * torch.rsqrt(variance + self.eps)
    return (self.weight * normalized).to(input_dtype)


def load_model():
    sys.path.insert(0, str(MODEL_DIR))
    from configuration_baberu import BaberuOCRConfig
    from modeling_baberu import BaberuOCRModel, BaberuRMSNorm

    BaberuRMSNorm.forward = webgpu_rms_norm

    config = BaberuOCRConfig.from_pretrained(MODEL_DIR)
    expected_architecture = {
        "num_hidden_layers": NUM_LAYERS,
        "hidden_size": HIDDEN_SIZE,
        "intermediate_size": 1536,
        "num_attention_heads": 8,
        "num_key_value_heads": NUM_KV_HEADS,
        "head_dim": HEAD_DIM,
        "vision_num_tokens": VISION_TOKENS,
        "vocab_size": VOCAB_SIZE,
    }
    mismatches = {
        name: {"expected": expected, "actual": getattr(config, name, None)}
        for name, expected in expected_architecture.items()
        if getattr(config, name, None) != expected
    }
    if mismatches:
        raise RuntimeError(
            "Checkpoint is not the complete native 121 MB Baberu architecture: "
            f"{mismatches}"
        )
    with torch.device("meta"):
        model = BaberuOCRModel(config)

    state = load_safetensors(MODEL_DIR / "model.safetensors", device="cpu")
    incompatible = model.load_state_dict(state, strict=False, assign=True)
    model.tie_weights()

    unexpected = list(incompatible.unexpected_keys)
    missing = [name for name in incompatible.missing_keys if name != "lm_head.weight"]
    if unexpected or missing:
        raise RuntimeError(
            f"Checkpoint mismatch: missing={missing}, unexpected={unexpected}"
        )

    meta_parameters = [name for name, value in model.named_parameters() if value.is_meta]
    if meta_parameters:
        raise RuntimeError(f"Parameters remained on meta device: {meta_parameters[:10]}")

    model = model.float().eval()
    # The upstream implementation uses a Python row-assignment loop. The
    # legacy ONNX exporter expands a 257-token prefill into 257 ScatterND nodes.
    # This vectorized expression is numerically identical and exports to the
    # small Range/LessOrEqual/Where form supported by ORT WebGPU.
    model.model._build_causal_mask = vectorized_causal_mask
    return model


def flatten_cache(past_key_values: Sequence[Sequence[torch.Tensor]]):
    keys = tuple(layer[0] for layer in past_key_values)
    values = tuple(layer[1] for layer in past_key_values)
    return keys + values


class DecoderBase(torch.nn.Module):
    def __init__(self, ocr_model):
        super().__init__()
        self.decoder = ocr_model.model
        self.lm_head = ocr_model.lm_head
        self.logit_cap = float(ocr_model.config.final_logit_softcap or 0.0)

    def project_logits(self, hidden_states: torch.Tensor) -> torch.Tensor:
        logits = self.lm_head(hidden_states)
        if self.logit_cap:
            logits = torch.tanh(logits / self.logit_cap) * self.logit_cap
        return logits


class DecoderPrefill(DecoderBase):
    def __init__(self, ocr_model):
        super().__init__(ocr_model)
        bos = torch.zeros(1, 1, VOCAB_SIZE, dtype=torch.float32)
        bos[0, 0, 1] = 1.0
        self.register_buffer("bos_one_hot", bos, persistent=False)

    def forward(self, vision_embeds: torch.Tensor):
        bos_embed = torch.matmul(self.bos_one_hot, self.lm_head.weight)
        inputs_embeds = torch.cat((vision_embeds, bos_embed), dim=1)
        outputs = self.decoder(
            inputs_embeds=inputs_embeds,
            past_key_values=None,
            use_cache=True,
            return_dict=True,
        )
        # Generation only consumes the final prefill position. Projecting all
        # 257 positions would needlessly copy about 14 MiB of logits to JS.
        final_hidden_state = outputs.last_hidden_state[:, -1:, :]
        return (self.project_logits(final_hidden_state),) + flatten_cache(
            outputs.past_key_values
        )


class DecoderStep(DecoderBase):
    def forward(
        self,
        token_one_hot: torch.Tensor,
        position_ids: torch.Tensor,
        past_k0: torch.Tensor,
        past_k1: torch.Tensor,
        past_k2: torch.Tensor,
        past_k3: torch.Tensor,
        past_k4: torch.Tensor,
        past_k5: torch.Tensor,
        past_v0: torch.Tensor,
        past_v1: torch.Tensor,
        past_v2: torch.Tensor,
        past_v3: torch.Tensor,
        past_v4: torch.Tensor,
        past_v5: torch.Tensor,
    ):
        keys = (past_k0, past_k1, past_k2, past_k3, past_k4, past_k5)
        values = (past_v0, past_v1, past_v2, past_v3, past_v4, past_v5)
        token_embed = torch.matmul(token_one_hot, self.lm_head.weight)
        outputs = self.decoder(
            inputs_embeds=token_embed,
            position_ids=position_ids,
            past_key_values=tuple(zip(keys, values)),
            use_cache=True,
            return_dict=True,
        )
        return (self.project_logits(outputs.last_hidden_state),) + flatten_cache(
            outputs.past_key_values
        )


def output_names() -> list[str]:
    return (
        ["logits"]
        + [f"present_k{index}" for index in range(NUM_LAYERS)]
        + [f"present_v{index}" for index in range(NUM_LAYERS)]
    )


def past_names() -> list[str]:
    return [f"past_k{index}" for index in range(NUM_LAYERS)] + [
        f"past_v{index}" for index in range(NUM_LAYERS)
    ]


def export_prefill(wrapper: DecoderPrefill, path: Path):
    torch.manual_seed(7)
    vision = torch.randn(1, VISION_TOKENS, HIDDEN_SIZE, dtype=torch.float32)
    torch.onnx.export(
        wrapper,
        (vision,),
        path,
        input_names=["vision_embeds"],
        output_names=output_names(),
        opset_version=OPSET,
        do_constant_folding=True,
        dynamo=False,
    )
    return vision


def make_past(length: int) -> tuple[torch.Tensor, ...]:
    torch.manual_seed(11 + length)
    return tuple(
        torch.randn(1, NUM_KV_HEADS, length, HEAD_DIM, dtype=torch.float32) * 0.02
        for _ in range(NUM_LAYERS * 2)
    )


def make_one_hot(token: int) -> torch.Tensor:
    value = torch.zeros(1, 1, VOCAB_SIZE, dtype=torch.float32)
    value[0, 0, token] = 1.0
    return value


def export_step(wrapper: DecoderStep, path: Path):
    token_one_hot = make_one_hot(4)
    position_ids = torch.tensor([[VISION_TOKENS + 1]], dtype=torch.int32)
    past = make_past(VISION_TOKENS + 1)
    input_names = ["token_one_hot", "position_ids"] + past_names()
    dynamic_axes = {
        name: {2: "past_len"} for name in past_names()
    } | {
        name: {2: "total_len"} for name in output_names()[1:]
    }
    torch.onnx.export(
        wrapper,
        (token_one_hot, position_ids, *past),
        path,
        input_names=input_names,
        output_names=output_names(),
        dynamic_axes=dynamic_axes,
        opset_version=OPSET,
        do_constant_folding=True,
        dynamo=False,
    )
    return token_one_hot, position_ids, past


def sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for chunk in iter(lambda: handle.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def tensor_shape(value: onnx.ValueInfoProto) -> list[int | str]:
    return [dimension.dim_value or dimension.dim_param for dimension in value.type.tensor_type.shape.dim]


def inspect_graph(path: Path) -> dict:
    model = onnx.load(path, load_external_data=False)
    onnx.checker.check_model(model)
    operators = Counter(node.op_type for node in model.graph.node)
    blocked = sorted(QUANTIZED_ONLY_OPS.intersection(operators))
    if blocked:
        raise RuntimeError(f"{path.name} still contains WASM-only quantized ops: {blocked}")
    return {
        "bytes": path.stat().st_size,
        "sha256": sha256(path),
        "opsets": {entry.domain or "ai.onnx": entry.version for entry in model.opset_import},
        "operators": dict(sorted(operators.items())),
        "inputs": {value.name: tensor_shape(value) for value in model.graph.input},
        "outputs": {value.name: tensor_shape(value) for value in model.graph.output},
    }


def numpy_inputs(values: Iterable[torch.Tensor]) -> list[np.ndarray]:
    return [value.detach().cpu().numpy() for value in values]


def compare_outputs(
    label: str,
    expected: Sequence[torch.Tensor],
    actual: Sequence[np.ndarray],
) -> dict:
    if len(expected) != len(actual):
        raise RuntimeError(f"{label}: output count differs")
    differences = []
    for expected_value, actual_value in zip(expected, actual):
        expected_array = expected_value.detach().cpu().numpy()
        if expected_array.shape != actual_value.shape:
            raise RuntimeError(
                f"{label}: shape differs: {expected_array.shape} != {actual_value.shape}"
            )
        differences.append(float(np.max(np.abs(expected_array - actual_value))))
    expected_token = int(expected[0][0, -1].argmax().item())
    actual_token = int(actual[0][0, -1].argmax())
    if expected_token != actual_token:
        raise RuntimeError(
            f"{label}: top token differs: PyTorch={expected_token}, ONNX={actual_token}"
        )
    return {
        "max_abs_by_output": differences,
        "max_abs": max(differences),
        "top_token": expected_token,
    }


def validate_prefill(
    wrapper: DecoderPrefill,
    path: Path,
    vision: torch.Tensor,
) -> dict:
    with torch.inference_mode():
        expected = wrapper(vision)
    session = ort.InferenceSession(str(path), providers=["CPUExecutionProvider"])
    actual = session.run(None, {"vision_embeds": vision.numpy()})
    return compare_outputs("prefill", expected, actual)


def validate_step(wrapper: DecoderStep, path: Path, lengths: Sequence[int]) -> dict:
    session = ort.InferenceSession(str(path), providers=["CPUExecutionProvider"])
    results = {}
    for length in lengths:
        token_one_hot = make_one_hot(4)
        position_ids = torch.tensor([[length]], dtype=torch.int32)
        past = make_past(length)
        with torch.inference_mode():
            expected = wrapper(token_one_hot, position_ids, *past)
        values = numpy_inputs((token_one_hot, position_ids, *past))
        actual = session.run(
            None,
            dict(zip((item.name for item in session.get_inputs()), values)),
        )
        results[str(length)] = compare_outputs(f"step[{length}]", expected, actual)
    return results


def main() -> None:
    if not (MODEL_DIR / "model.safetensors").exists():
        raise SystemExit("Model is missing. Run download_model.py first.")
    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
    REPORT_DIR.mkdir(parents=True, exist_ok=True)

    torch.manual_seed(0)
    torch.set_grad_enabled(False)
    model = load_model()
    prefill = DecoderPrefill(model).eval()
    step = DecoderStep(model).eval()

    prefill_path = OUTPUT_DIR / "decoder_prefill_fp32.onnx"
    step_path = OUTPUT_DIR / "decoder_step_fp32.onnx"
    print(f"Exporting {prefill_path}")
    vision = export_prefill(prefill, prefill_path)
    print(f"Exporting {step_path}")
    export_step(step, step_path)

    report = {
        "model_revision": json.loads(
            (MODEL_DIR / "source-revision.json").read_text(encoding="utf-8")
        ),
        "opset": OPSET,
        "graphs": {
            "prefill": inspect_graph(prefill_path),
            "step": inspect_graph(step_path),
        },
        "parity": {
            "prefill": validate_prefill(prefill, prefill_path, vision),
            "step": validate_step(step, step_path, (VISION_TOKENS + 1, 288)),
        },
    }
    report_path = REPORT_DIR / "export-report.json"
    report_path.write_text(json.dumps(report, indent=2) + "\n", encoding="utf-8")
    print(f"Wrote {report_path}")
    print(json.dumps(report["parity"], indent=2))


if __name__ == "__main__":
    main()