File size: 4,445 Bytes
5dc80b3 | 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 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 | #!/usr/bin/env python3
"""
Convert HRM checkpoints to safetensors and push to Hugging Face.
Usage:
python push_to_hf.py \
--baseline "checkpoints/Sudoku-extreme-1k-aug-1000 ACT-torch/HierarchicalReasoningModel_ACTV1 belligerent-squirrel/step_52080" \
--tiered "checkpoints/Sudoku-extreme-1k-aug-1000 ACT-torch/HRM_Tiered realistic-dalmatian/step_52080" \
--repo "Code2aum/HRM_based_evolution"
"""
import argparse
import os
import shutil
import torch
from safetensors.torch import save_file
from huggingface_hub import login, HfApi, upload_folder
def convert_checkpoint(ckpt_path, output_dir, model_name):
"""Load .pt checkpoint, save as safetensors + config."""
os.makedirs(output_dir, exist_ok=True)
# Load weights
state_dict = torch.load(ckpt_path, map_location="cpu", weights_only=True)
# Strip torch.compile prefix if present
cleaned = {}
for k, v in state_dict.items():
key = k.removeprefix("_orig_mod.")
cleaned[key] = v
# Save as safetensors
st_path = os.path.join(output_dir, "model.safetensors")
save_file(cleaned, st_path)
print(f" ✓ Saved {st_path} ({os.path.getsize(st_path)/1e6:.1f} MB)")
# Copy config
ckpt_dir = os.path.dirname(ckpt_path)
config_src = os.path.join(ckpt_dir, "all_config.yaml")
if os.path.exists(config_src):
shutil.copy2(config_src, os.path.join(output_dir, "config.yaml"))
print(f" ✓ Copied config.yaml")
# Copy model source
for src_file in ["hrm_act_v1.py", "hrm_tiered.py", "losses.py"]:
src_path = os.path.join(ckpt_dir, src_file)
if os.path.exists(src_path):
shutil.copy2(src_path, os.path.join(output_dir, src_file))
# Write a model card
card = f"""---
tags:
- hrm
- hierarchical-reasoning
- sudoku
- pytorch
license: mit
---
# {model_name}
Hierarchical Reasoning Model trained on Sudoku-Extreme-1K (20,000 epochs).
## Architecture
- **Type**: {model_name}
- **Hidden Size**: 512
- **Heads**: 8
- **H/L Layers**: 4/4
- **H/L Cycles**: 2/2
- **Parameters**: ~27.3M
## Training
- **Dataset**: Sudoku-Extreme-1K (1000 puzzles, 1000 augmentations each)
- **Epochs**: 20,000
- **Batch Size**: 384
- **Learning Rate**: 7e-5
- **GPU**: NVIDIA RTX 4090
## Usage
```python
from safetensors.torch import load_file
state_dict = load_file("{model_name}/model.safetensors")
```
"""
with open(os.path.join(output_dir, "README.md"), "w") as f:
f.write(card)
print(f" ✓ Created README.md")
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--baseline", type=str, required=True)
parser.add_argument("--tiered", type=str, required=True)
parser.add_argument("--repo", type=str, default="Code2aum/HRM_based_evolution")
parser.add_argument("--output-dir", type=str, default="hf_upload")
parser.add_argument("--no-push", action="store_true", help="Convert only, skip push")
args = parser.parse_args()
print("=" * 60)
print(" HRM → SafeTensors → Hugging Face")
print("=" * 60)
# Convert baseline
base_dir = os.path.join(args.output_dir, "baseline_hrm_v1")
print(f"\n Converting Baseline...")
convert_checkpoint(args.baseline, base_dir, "HRM_Baseline_V1")
# Convert tiered
tier_dir = os.path.join(args.output_dir, "tiered_hrm_sram_dram")
print(f"\n Converting Tiered...")
convert_checkpoint(args.tiered, tier_dir, "HRM_Tiered_SRAM_DRAM")
# Copy benchmark results if available
bench_dir = "benchmark_results"
if os.path.exists(bench_dir):
dest = os.path.join(args.output_dir, "benchmark_results")
if not os.path.exists(dest):
shutil.copytree(bench_dir, dest)
print(f"\n ✓ Copied benchmark_results/")
if args.no_push:
print(f"\n Files ready at: {args.output_dir}/")
print(" Run without --no-push to upload to HF.")
return
# Push to HF
print(f"\n Logging in to Hugging Face...")
login()
print(f"\n Uploading to {args.repo}...")
api = HfApi()
api.create_repo(repo_id=args.repo, repo_type="model", exist_ok=True)
upload_folder(
folder_path=args.output_dir,
repo_id=args.repo,
repo_type="model",
)
print(f"\n ✓ Uploaded to https://huggingface.co/{args.repo}")
print("\n" + "=" * 60)
print(" Done!")
print("=" * 60)
if __name__ == "__main__":
main()
|