cfahlgren1 HF Staff
export the compressed browser format from train.py instead of int8 activations
4dece7f verified | # /// 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() | |