| """Convert rsac checkpoints (.pt) into the Hugging Face release layout (safetensors + JSON). |
| |
| For each NAME=CKPT it writes OUT/NAME/actor.safetensors (policy only, for deployment), |
| OUT/NAME/full.safetensors (actor.* and critics.* plus log_alpha, for fine-tuning) and |
| OUT/NAME/training.json (iteration, samples, run config, last validation row, reward scales), |
| then reloads both files with modeling_aio and checks the actor gives the same outputs as the |
| checkpoint. Copy a checkpoint before converting it if a run may still overwrite it. |
| |
| python release/huggingface/export.py --out DIR apc-v3-finetune=PATH/best.pt mixed-it6000=PATH/best.pt |
| """ |
| import argparse, json, sys |
| from pathlib import Path |
| import torch |
| from safetensors.torch import save_file |
|
|
| HERE = Path(__file__).resolve().parent |
| sys.path.insert(0, str(HERE)) |
| import modeling_aio as M |
|
|
|
|
| def tensors(state, prefix=''): |
| return {prefix+k: v.detach().cpu().clone().contiguous() for k, v in state.items()} |
|
|
|
|
| def load_pt(ckpt): |
| """torch.load that also works off Windows: run configs pickle pathlib.WindowsPath objects.""" |
| import pathlib |
| saved = pathlib.WindowsPath |
| if sys.platform != 'win32': |
| pathlib.WindowsPath = pathlib.PureWindowsPath |
| try: |
| return torch.load(ckpt, map_location='cpu', weights_only=False) |
| finally: |
| pathlib.WindowsPath = saved |
|
|
|
|
| def export(name, ckpt, out): |
| ck = load_pt(ckpt) |
| d = out/name; d.mkdir(parents=True, exist_ok=True) |
| meta = {'format': 'pt', 'source': Path(ckpt).name, 'iteration': str(ck.get('iteration')), |
| 'samples': str(ck.get('samples'))} |
| save_file(tensors(ck['model']['actor']), d/'actor.safetensors', metadata=meta) |
| full = tensors(ck['model']['actor'], 'actor.') |
| if 'critics' in ck['model']: |
| full.update(tensors(ck['model']['critics'], 'critics.')) |
| if 'log_alpha' in ck: |
| full['log_alpha'] = torch.tensor([float(ck['log_alpha'])]) |
| save_file(full, d/'full.safetensors', metadata=meta) |
| info = {k: ck.get(k) for k in ('iteration', 'samples', 'log_alpha', 'best', 'stale', 'reward_scale', 'val', 'config')} |
| (d/'training.json').write_text(json.dumps(info, indent=1, default=str)) |
|
|
| ref = M.Actor(); ref.load_state_dict(ck['model']['actor']); ref.eval() |
| g = torch.Generator().manual_seed(0) |
| obs = torch.randn(3, 7, M.SIZE, generator=g); prev = torch.rand(3, 7, M.NMV, generator=g)*2-1 |
| with torch.no_grad(): |
| a = ref(obs, prev)[0] |
| for f in ('actor.safetensors', 'full.safetensors'): |
| b = M.load_actor(d/f)(obs, prev)[0] |
| assert torch.equal(a, b), f'{name}/{f}: actor output differs' |
| if 'critics' in ck['model']: |
| M.load_full(d/'full.safetensors') |
| n = M.count(ref) |
| print(json.dumps(dict(name=name, iteration=ck.get('iteration'), samples=ck.get('samples'), actor_params=n, |
| files={p.name: p.stat().st_size for p in sorted(d.iterdir())}))) |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument('--out', type=Path, required=True) |
| ap.add_argument('ckpts', nargs='+', help='NAME=CKPT') |
| a = ap.parse_args() |
| for item in a.ckpts: |
| name, path = item.split('=', 1) |
| export(name, Path(path), a.out) |
|
|
|
|
| if __name__ == '__main__': |
| main() |
|
|
