training/export_checkpoint.py

"""Export a trained Ares checkpoint and tokenizer for the next browser-runtime stage."""
import argparse, json, shutil
from pathlib import Path
import torch

def main():
 p=argparse.ArgumentParser(); p.add_argument('--checkpoint',default='checkpoints_modern/best.pt'); p.add_argument('--tokenizer',default='ares_tokenizer.json'); p.add_argument('--out',default='ares_export'); a=p.parse_args(); out=Path(a.out); out.mkdir(parents=True,exist_ok=True)
 ck=torch.load(a.checkpoint,map_location='cpu',weights_only=False); torch.save(ck['model'],out/'model_state_dict.pt')
 if Path(a.tokenizer).exists(): shutil.copy2(a.tokenizer,out/'tokenizer.json')
 manifest={'format':'ares-checkpoint-v1','source_checkpoint':str(a.checkpoint),'config':ck.get('config',{}),'step':ck.get('step'),'best_validation_loss':ck.get('best'),'tokenizer':('tokenizer.json' if Path(a.tokenizer).exists() else None),'browser_runtime':'pending-export-adapter'}
 (out/'manifest.json').write_text(json.dumps(manifest,indent=2),encoding='utf8'); print('Exported',out)
if __name__=='__main__': main()