# /// script
# requires-python = ">=3.11,<3.14"
# dependencies = [
#   "torch>=2.6,<2.8",
#   "numpy<3",
#   "tokenizers>=0.20",
#   "onnx>=1.17",
#   "onnxruntime>=1.20",
#   "onnx-ir",
#   "safetensors>=0.4",
#   "huggingface_hub>=0.30",
# ]
# ///
"""Re-exports a trained tinynext checkpoint with train.py's export (float ONNX and the compressed browser format in
web/), without retraining.

  uv run export.py --weights v2-release/v2 --out v2-release/browser
"""

import argparse
import importlib.util
import json
import sys
from pathlib import Path

import torch
from safetensors.torch import load_file
from tokenizers import Tokenizer

HERE = Path(__file__).parent


def load(name: str):
    spec = importlib.util.spec_from_file_location(name, HERE / f"{'eval' if name == 'tinynext_eval' else name}.py")
    module = importlib.util.module_from_spec(spec)
    sys.modules[name] = module
    spec.loader.exec_module(module)
    return module


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--weights", required=True, help="Folder with model.safetensors, config.json, tokenizer.json")
    parser.add_argument("--out", required=True)
    args = parser.parse_args()
    train, ev = load("train"), load("tinynext_eval")
    weights, out = Path(args.weights), Path(args.out)
    out.mkdir(parents=True, exist_ok=True)

    config = json.loads((weights / "config.json").read_text())
    model = train.TinyNext(config["vocab"], config["dim"], config["layers"], config["heads"], config["ffn"])
    model.load_state_dict(load_file(weights / "model.safetensors"))
    tokenizer = Tokenizer.from_file(str(weights / "tokenizer.json"))
    (out / "tokenizer.json").write_bytes((weights / "tokenizer.json").read_bytes())
    report = train.export(model, out, torch.device("cpu"), train.parity_prompts(ev, tokenizer), load("compress"))
    (out / "web" / "config.json").write_bytes((weights / "config.json").read_bytes())
    (out / "export.json").write_text(json.dumps(report, indent=1))
    print(json.dumps(report, indent=1))


if __name__ == "__main__":
    main()