import torch

def export_model():
    print("正在读取庞大的训练存档...")
    # 读取你那个 2.7GB 的存档
    huge_ckpt = torch.load("checkpoints/latest_ckpt.pt", map_location="cpu")
    
    # 提取出纯纯的模型权重（只有 450 KB）
    pure_model_weights = huge_ckpt["model_state_dict"]
    
    # 将纯权重保存给 Botzone 用
    deploy_path = "botzone_cfr_net.pt"
    torch.save(pure_model_weights, deploy_path)
    
    print(f"瘦身成功！纯模型已保存为: {deploy_path}")
    
    import os
    size_kb = os.path.getsize(deploy_path) / 1024
    print(f"真实模型体积: {size_kb:.2f} KB")

if __name__ == "__main__":
    export_model()