#!/usr/bin/env python3
"""Exports Qwen2.5-{1.5B,3B}-Instruct to int4 ONNX models that ONNX Runtime runs on Android and iOS.

The ONNX Runtime GenAI model builder writes the graph: GroupQueryAttention with a KV cache, 4-bit MatMulNBits
(blocks of 32, accuracy_level 4) and the tied embedding as a 4-bit GatherBlockQuantized that shares its weights
with the LM head. The LM head only computes the last position's logits. Next to the model it writes the tokenizer
as vocab.txt and merges.txt, the plain files the Kotlin and Swift tokenizers read.

Usage:
  .venv/bin/python -I export.py [--size 1.5B|3B ...]
"""

import argparse
import json
import shutil
import subprocess
import sys
from pathlib import Path

import onnx

sys.path.insert(0, str(Path(__file__).resolve().parent))  # -I drops the script's own directory
from common import HERE, SIZES, out_dir, src_dir  # noqa: E402

EXTRA_OPTIONS = [
    "shared_embeddings=true",
    "prune_lm_head=true",
    "block_size=32",
    "accuracy_level=4",
    "hf_remote=false",
]


def build(size: str) -> Path:
    out = out_dir(size)
    onnx_dir = out / "onnx"
    if onnx_dir.exists():
        shutil.rmtree(onnx_dir)
    cmd = [
        sys.executable, "-m", "onnxruntime_genai.models.builder",
        "-i", str(src_dir(size)), "-o", str(onnx_dir), "-p", "int4", "-e", "cpu",
        "-c", str(HERE / "data" / "cache"), "--extra_options", *EXTRA_OPTIONS,
    ]
    subprocess.run(cmd, check=True)
    return onnx_dir


def write_tokenizer(size: str) -> None:
    """vocab.txt: one byte-level token per line, by id (added tokens included). merges.txt: one 'a b' pair per line, by rank."""
    tok = json.loads((src_dir(size) / "tokenizer.json").read_text(encoding="utf-8"))
    model = tok["model"]
    assert model["type"] == "BPE" and not model.get("byte_fallback"), model["type"]
    vocab = dict(model["vocab"])
    for added in tok["added_tokens"]:
        vocab[added["content"]] = added["id"]
    ids = sorted(vocab.values())
    assert ids == list(range(len(ids))), "token ids must be contiguous"
    by_id = {i: t for t, i in vocab.items()}
    for token in by_id.values():
        assert "\n" not in token and "\r" not in token, repr(token)
    out = out_dir(size)
    (out / "vocab.txt").write_text("".join(by_id[i] + "\n" for i in ids), encoding="utf-8")
    merges = [m if isinstance(m, str) else " ".join(m) for m in model["merges"]]
    for m in merges:
        assert m.count(" ") == 1, repr(m)
    (out / "merges.txt").write_text("".join(m + "\n" for m in merges), encoding="utf-8")
    shutil.copy(src_dir(size) / "LICENSE", out / "LICENSE")


def describe(onnx_dir: Path) -> None:
    model = onnx.load(onnx_dir / "model.onnx", load_external_data=False)

    def shape(v):
        return [d.dim_param or d.dim_value for d in v.type.tensor_type.shape.dim]

    for v in list(model.graph.input)[:5] + list(model.graph.output)[:3]:
        print(f"  {v.name}: {onnx.TensorProto.DataType.Name(v.type.tensor_type.elem_type)} {shape(v)}")
    ops = sorted({n.op_type for n in model.graph.node})
    print("  ops:", ", ".join(ops))


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--size", choices=SIZES, nargs="+", default=list(SIZES))
    args = parser.parse_args()
    for size in args.size:
        onnx_dir = build(size)
        write_tokenizer(size)
        print(size)
        describe(onnx_dir)


if __name__ == "__main__":
    main()