"""Convert the latest checkpoint to bf16 to halve file size for Colab upload."""
import torch
import json
import os

def main():
    ckpt = torch.load("checkpoints/latest.pt", map_location="cpu", weights_only=False)
    print(f"Original checkpoint: step={ckpt['step']}, loss={ckpt['loss']}")
    print(f"  Embedding shape: {ckpt['model_state_dict']['token_embedding.weight'].shape}")

    # Convert all float32 tensors to bf16
    state = ckpt["model_state_dict"]
    for key in state:
        if state[key].dtype == torch.float32:
            state[key] = state[key].to(torch.bfloat16)

    ckpt["model_state_dict"] = state
    out_path = "checkpoints/latest_bf16.pt"
    torch.save(ckpt, out_path)
    size_mb = os.path.getsize(out_path) / 1e6
    print(f"bf16 checkpoint saved: {out_path} ({size_mb:.1f} MB)")

if __name__ == "__main__":
    main()