File size: 2,441 Bytes
b296ad4 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 | """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()
|