#!/usr/bin/env python3
"""
Export the Teensy PyTorch checkpoint to HuggingFace-compatible files.

Produces:
  - model.safetensors   (model weights only, smaller than the full .pt checkpoint)
  - config.json         (model hyperparameters)

Copyright (c) 2025 Pankaj Doharey
Modified from NanoGPT (Andrej Karpathy).
"""

import os
import json
import argparse

import torch
from safetensors.torch import save_file

from model import adapt_nanogpt_weights


def export(out_dir, export_dir):
    os.makedirs(export_dir, exist_ok=True)

    ckpt_path = os.path.join(out_dir, 'teensy-0.pt')
    print(f"Loading checkpoint: {ckpt_path}")
    checkpoint = torch.load(ckpt_path, map_location='cpu')

    cfg = checkpoint['model_args']
    with open(os.path.join(export_dir, 'config.json'), 'w') as f:
        json.dump(cfg, f, indent=2)
    print(f"Wrote config.json: {cfg}")

    state_dict = checkpoint['model']
    unwanted_prefix = '_orig_mod.'
    for k, v in list(state_dict.items()):
        if k.startswith(unwanted_prefix):
            state_dict[k[len(unwanted_prefix):]] = state_dict.pop(k)
    state_dict = adapt_nanogpt_weights(state_dict)

    # Move all tensors to CPU and convert to contiguous float32 for portability.
    # Clone tensors that may share storage (e.g. tied token/output embeddings) so
    # safetensors can write them independently.
    state_dict = {k: v.contiguous().to(torch.float32).cpu().clone() for k, v in state_dict.items()}

    save_file(state_dict, os.path.join(export_dir, 'model.safetensors'))
    print(f"Wrote model.safetensors to {export_dir}")

    pt_size = os.path.getsize(ckpt_path)
    st_size = os.path.getsize(os.path.join(export_dir, 'model.safetensors'))
    print(f"\nSize comparison:")
    print(f"  .pt checkpoint:         {pt_size / 1e6:.1f} MB")
    print(f"  model.safetensors:     {st_size / 1e6:.1f} MB")
    print(f"  savings:               {(1 - st_size / pt_size) * 100:.1f}%")


def main():
    parser = argparse.ArgumentParser(description='Export Teensy checkpoint to safetensors')
    parser.add_argument('--out_dir', default='checkpoints', help='Directory containing teensy-0.pt')
    parser.add_argument('--export_dir', default='exported', help='Directory to write exported files')
    args = parser.parse_args()
    export(args.out_dir, args.export_dir)


if __name__ == '__main__':
    main()