File size: 2,087 Bytes
33bfb6d 4dece7f 33bfb6d 4dece7f 33bfb6d 4dece7f 33bfb6d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 | # /// 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()
|