| |
| """ |
| 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) |
|
|
| |
| state_dict = torch.load(ckpt_path, map_location="cpu", weights_only=True) |
|
|
| |
| cleaned = {} |
| for k, v in state_dict.items(): |
| key = k.removeprefix("_orig_mod.") |
| cleaned[key] = v |
|
|
| |
| 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)") |
|
|
| |
| 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") |
|
|
| |
| 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)) |
|
|
| |
| 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) |
|
|
| |
| 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") |
|
|
| |
| 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") |
|
|
| |
| 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 |
|
|
| |
| 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() |
|
|
