# /// script
# requires-python = ">=3.11,<3.14"
# dependencies = ["onnx>=1.17", "numpy<3", "onnxruntime>=1.20"]
# ///
"""Shrinks a float ONNX model for download without changing how it runs.

The browser computes in float32 (onnxruntime-web's int8 kernels are several times slower than its float ones), so
only the download is compressed: every large weight is stored as int8 with a float32 scale per block of 64 values.
The page expands them back to float32 and hands them to onnxruntime-web as the model's external data.

Writes, next to each other:
  graph.onnx      the model graph; large weights point into "weights.f32" (external data)
  weights.q8      int8 values, then float32 scales, tensor by tensor
  weights.json    for each tensor: where its int8 values and scales start, its shape, and where it goes in weights.f32

  uv run compress.py v2-release/browser/model.onnx v2-release/web
"""

import json
import sys
from pathlib import Path

import numpy as np
import onnx
from onnx import numpy_helper
from onnx.external_data_helper import set_external_data

BLOCK = 64
MIN_ELEMENTS = 4096  # smaller tensors (norm weights) stay inside the graph as float32
EXTERNAL = "weights.f32"


def quantize(values: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
    """Symmetric int8 with one scale per BLOCK consecutive values."""
    flat = values.astype(np.float32).ravel()
    padded = np.pad(flat, (0, (-len(flat)) % BLOCK)).reshape(-1, BLOCK)
    scales = np.abs(padded).max(axis=1) / 127
    scales[scales == 0] = 1
    q = np.clip(np.round(padded / scales[:, None]), -127, 127).astype(np.int8)
    return q.ravel()[: len(flat)], scales.astype(np.float32)


def dequantize(q: np.ndarray, scales: np.ndarray, count: int) -> np.ndarray:
    padded = np.pad(q.astype(np.float32), (0, (-count) % BLOCK)).reshape(-1, BLOCK)
    return (padded * scales[:, None]).ravel()[:count]


def dedupe_transposes(model) -> None:
    """The output layer reuses the embedding table, but export stores a transposed copy. Keep one and transpose it
    in the graph instead, which saves about 2.6 MB of download."""
    from onnx import helper

    graph = model.graph
    arrays = {t.name: numpy_helper.to_array(t) for t in graph.initializer if t.dims and np.prod(t.dims) >= MIN_ELEMENTS}
    for name, values in list(arrays.items()):
        if values.ndim != 2:
            continue
        for other, candidate in arrays.items():
            if other != name and candidate.shape == values.T.shape and np.array_equal(candidate, values.T):
                transposed = f"{other}_transposed"
                graph.node.insert(0, helper.make_node("Transpose", [other], [transposed], perm=[1, 0], name=transposed))
                for node in graph.node:
                    node.input[:] = [transposed if i == name else i for i in node.input]
                graph.initializer.remove(next(t for t in graph.initializer if t.name == name))
                print(f"stored {name} as a transpose of {other}")
                return


def compress(src: Path, out: Path) -> dict:
    out.mkdir(parents=True, exist_ok=True)
    model = onnx.load(str(src))
    dedupe_transposes(model)
    manifest, q8 = [], bytearray()
    f32_offset = 0
    for tensor in model.graph.initializer:
        values = numpy_helper.to_array(tensor)
        if values.dtype != np.float32 or values.size < MIN_ELEMENTS:
            continue
        q, scales = quantize(values)
        entry = {"name": tensor.name, "shape": list(values.shape), "count": int(values.size), "f32_offset": f32_offset,
                 "q_offset": len(q8), "scale_offset": len(q8) + len(q) + (-len(q)) % 4}
        q8 += q.tobytes()
        q8 += b"\0" * ((-len(q8)) % 4)  # keep scales 4-byte aligned for Float32Array views
        q8 += scales.tobytes()
        manifest.append(entry)
        # Point the graph at weights.f32, which the browser builds by expanding weights.q8.
        if not tensor.HasField("raw_data"):
            tensor.raw_data = values.astype(np.float32).tobytes()
            tensor.ClearField("float_data")
        set_external_data(tensor, location=EXTERNAL, offset=f32_offset, length=values.size * 4)
        tensor.ClearField("raw_data")
        tensor.data_location = onnx.TensorProto.EXTERNAL
        f32_offset += values.size * 4
    onnx.save_model(model, str(out / "graph.onnx"))
    (out / "weights.q8").write_bytes(bytes(q8))
    (out / "weights.json").write_text(json.dumps({"block": BLOCK, "f32_bytes": f32_offset, "tensors": manifest}))
    return {"graph_mb": round((out / "graph.onnx").stat().st_size / 1e6, 3), "weights_q8_mb": round(len(q8) / 1e6, 2),
            "expanded_mb": round(f32_offset / 1e6, 2), "tensors": len(manifest)}


def expand(out: Path) -> Path:
    """What the browser does: rebuild weights.f32 from weights.q8. Used here to check accuracy."""
    meta = json.loads((out / "weights.json").read_text())
    raw = (out / "weights.q8").read_bytes()
    f32 = np.zeros(meta["f32_bytes"] // 4, np.float32)
    for t in meta["tensors"]:
        q = np.frombuffer(raw, np.int8, t["count"], t["q_offset"])
        blocks = (t["count"] + meta["block"] - 1) // meta["block"]
        scales = np.frombuffer(raw, np.float32, blocks, t["scale_offset"])
        f32[t["f32_offset"] // 4: t["f32_offset"] // 4 + t["count"]] = dequantize(q, scales, t["count"])
    path = out / EXTERNAL
    f32.tofile(path)
    return path


if __name__ == "__main__":
    src, out = Path(sys.argv[1]), Path(sys.argv[2])
    print(json.dumps(compress(src, out), indent=1))
    expand(out)
    print(f"expanded {out / EXTERNAL} for checking")