| """Convert a project ``.pt`` checkpoint into Hugging Face format. |
| |
| Example: |
| python convert_checkpoint.py --checkpoint ../../checkpoints/policy/on-policy/step-0006000.pt |
| """ |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| from pathlib import Path |
| import sys |
|
|
| import torch |
|
|
| HERE = Path(__file__).resolve().parent |
| ROOT = HERE.parents[1] |
| if str(HERE) not in sys.path: |
| sys.path.insert(0, str(HERE)) |
|
|
| from configuration_chess_policy import ChessPolicyConfig |
| from modeling_chess_policy import ChessTransitionPolicy |
|
|
|
|
| def main() -> None: |
| parser = argparse.ArgumentParser(description=__doc__) |
| parser.add_argument( |
| "--checkpoint", |
| type=Path, |
| default=ROOT |
| / "checkpoints" |
| / "policy" |
| / "on-policy" |
| / "step-0006000.pt", |
| ) |
| parser.add_argument("--output-dir", type=Path, default=HERE) |
| parser.add_argument( |
| "--safe-serialization", |
| action=argparse.BooleanOptionalAction, |
| default=True, |
| help="Write model.safetensors instead of pytorch_model.bin.", |
| ) |
| args = parser.parse_args() |
|
|
| checkpoint = torch.load(args.checkpoint, map_location="cpu", weights_only=False) |
| if "model" not in checkpoint or "model_config" not in checkpoint: |
| raise ValueError("Checkpoint must contain model and model_config entries") |
|
|
| config = ChessPolicyConfig(**checkpoint["model_config"]) |
| config.architectures = ["ChessTransitionPolicy"] |
| config.auto_map = { |
| "AutoConfig": "configuration_chess_policy.ChessPolicyConfig", |
| "AutoModel": "modeling_chess_policy.ChessTransitionPolicy", |
| } |
| model = ChessTransitionPolicy(config) |
| model.load_state_dict(checkpoint["model"]) |
| model.eval() |
| args.output_dir.mkdir(parents=True, exist_ok=True) |
| model.save_pretrained( |
| args.output_dir, |
| safe_serialization=args.safe_serialization, |
| ) |
| print(f"Saved Hugging Face model to {args.output_dir}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
