| |
| """Run the declared quantization, discarded smoke, final training and paired evaluation.""" |
| import argparse |
| import hashlib |
| import json |
| import os |
| from pathlib import Path |
| import subprocess |
| import sys |
| import time |
|
|
| ROOT=Path(__file__).resolve().parents[1] |
|
|
| def sha(path): |
| h=hashlib.sha256() |
| with Path(path).open('rb') as f: |
| for block in iter(lambda:f.read(8<<20),b''): h.update(block) |
| return h.hexdigest() |
|
|
| def main(): |
| p=argparse.ArgumentParser(description=__doc__) |
| p.add_argument('--source-dir',type=Path,required=True) |
| p.add_argument('--prose-tokens',type=Path,required=True) |
| args=p.parse_args() |
| logs=ROOT/'logs';logs.mkdir(exist_ok=True) |
| receipt=logs/'experiment_v1.json' |
| if receipt.exists(): raise FileExistsError(receipt) |
| base=ROOT/'artifacts/w4_base_v1' |
| train=ROOT/'artifacts/w4_resurface_v1' |
| data=ROOT/'training_data/numeric_v1' |
| report=ROOT/'reports/w4_resurface_v1_confirm_full.json' |
| report.parent.mkdir(exist_ok=True) |
| common=['--source-dir',str(args.source_dir),'--w4-dir',str(base),'--data-root',str(data), |
| '--train-manifest-sha256',sha(data/'train/manifest.json'), |
| '--prose-manifest',str(ROOT/'docs/prose_train_manifest.json'), |
| '--prose-tokens',str(args.prose_tokens)] |
| stages=[ |
| ('quantize',['scripts/quantize_w4.py','--source',str(args.source_dir),'--output',str(base)]), |
| ('smoke',['scripts/train.py',*common,'--out-dir',str(ROOT/'artifacts/w4_smoke_v1'),'--smoke']), |
| ('train',['scripts/train.py',*common,'--out-dir',str(train)]), |
| ('evaluate',['scripts/evaluate.py','--source-dir',str(args.source_dir),'--w4-dir',str(base), |
| '--adapter',str(train/'adapter_fp16.pt'),'--train-report',str(train/'report.json'), |
| '--data-root',str(data),'--eval-manifest-sha256',sha(data/'confirm/manifest.json'), |
| '--report',str(report)]), |
| ] |
| state={'started_unix':time.time(),'complete':False,'stages':[], |
| 'protocol_sha256':sha(ROOT/'docs/PROTOCOL.md'), |
| 'code_sha256':{str(f.relative_to(ROOT)):sha(f) for folder in ['mamba2_recall','scripts'] |
| for f in sorted((ROOT/folder).glob('*.py'))}} |
| def save(): |
| tmp=receipt.with_suffix('.tmp');tmp.write_text(json.dumps(state,indent=2)+'\n');tmp.replace(receipt) |
| save() |
| env=dict(os.environ, PYTHONPATH=str(ROOT)) |
| for name,cmd in stages: |
| row={'name':name,'command':[sys.executable,*cmd],'started_unix':time.time()} |
| state['stages'].append(row);state['current_stage']=name;save() |
| log=logs/(name+'_v1.log') |
| with log.open('xb') as f: |
| result=subprocess.run(row['command'],cwd=ROOT,env=env,stdin=subprocess.DEVNULL, |
| stdout=f,stderr=subprocess.STDOUT) |
| row.update(returncode=result.returncode,finished_unix=time.time(),log=str(log)) |
| save() |
| if result.returncode: |
| state['failed_stage']=name;save();raise SystemExit(result.returncode) |
| state.update(complete=True,finished_unix=time.time());save() |
|
|
| if __name__=='__main__':main() |
|
|
