| """Freeze an owned training checkpoint and evaluate only development splits.""" |
| import argparse |
| import json |
| import os |
| from pathlib import Path |
| import subprocess |
| import sys |
| import torch |
| from safetensors.torch import save_file |
|
|
|
|
| def main(): |
| p=argparse.ArgumentParser();p.add_argument('--checkpoint',required=True);p.add_argument('--data',required=True) |
| p.add_argument('--out',required=True);args=p.parse_args();data=Path(args.data) |
| r=torch.load(args.checkpoint,map_location='cpu',weights_only=False) |
| dest=Path(args.out)/('step-'+str(r['step']));dest.mkdir(parents=True,exist_ok=True) |
| if (dest/'selection.json').exists():print((dest/'selection.json').read_text());return |
| save_file({k:v.to(torch.bfloat16).contiguous() for k,v in r['model'].items()},str(dest/'model.safetensors'), |
| metadata={'step':str(r['step']),'random_initialization':'true'}) |
| (dest/'config.json').write_text(json.dumps(r['config'],indent=2)) |
| info={k:r[k] for k in ['step','processed_tokens','response_tokens','training_seconds','random_initialization']} |
| (dest/'checkpoint-info.json').write_text(json.dumps(info,indent=2));del r |
| print(json.dumps({'frozen':str(dest),'step':info['step']}),flush=True);scores={};digest=None |
| for split in ['validation','development']: |
| if not (data/(split+'.jsonl')).exists():continue |
| with (dest/(split+'.log')).open('w') as log: |
| subprocess.run([sys.executable,'-u','-m','tinyquery.evaluate','--checkpoint',str(dest/'model.safetensors'), |
| '--tokenizer',str(data/'tokenizer.json'),'--data',str(data/(split+'.jsonl')), |
| '--out',str(dest/(split+'.jsonl'))],stdout=log,stderr=subprocess.STDOUT,check=True) |
| summary=json.loads((dest/(split+'.summary.json')).read_text());scores[split]=summary['metrics']['success'] |
| digest=summary['checkpoint_sha256'];print(json.dumps({'split':split,'metrics':summary['metrics']}),flush=True) |
| selection={'checkpoint':str(dest/'model.safetensors'),'step':info['step'],'checkpoint_sha256':digest, |
| 'development_scores':scores,'selection_score':sum(v['rate'] for v in scores.values())/len(scores), |
| 'rule':'Equal mean of full validation and development task success. Test/manual excluded.'} |
| (dest/'selection.json').write_text(json.dumps(selection,indent=2));print(json.dumps(selection),flush=True) |
|
|
|
|
| if __name__=='__main__':main() |
|
|
