File size: 6,090 Bytes
093306a | 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 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 | """Sixth-tranche fixed evaluation (no training, no generation). FP32 + math SDPA, as in the fifth evaluator.
Arms: CANDIDATE (sixth endpoint) and PARENT (step 4,126, from the fifth recovery). Retention baselines
A0/English93 and step3403/English599 are the recorded per-unit rows from the fifth T4 FP32 evaluation
(sha-pinned in FROZEN_GATES.json); the candidate is scored live with the same arithmetic."""
import argparse, hashlib, json, math, os, random, sys, time
import torch
from torch.nn.attention import sdpa_kernel, SDPBackend
sys.path.insert(0,os.path.dirname(os.path.abspath(__file__)))
from frozen_stasis.config import ModelConfig
from frozen_stasis.model import StasisDecoder
def fsha(p):
h=hashlib.sha256(); f=open(p,'rb')
for b in iter(lambda:f.read(1<<24),b''): h.update(b)
return h.hexdigest()
def load_model(state,dev):
m=StasisDecoder(ModelConfig()); m.load_state_dict(state,strict=True); return m.to(dev).eval()
@torch.no_grad()
def nll_window(m,tokens,dev):
x=torch.tensor(tokens,dtype=torch.long,device=dev)[None]
with sdpa_kernel(SDPBackend.MATH):
logits,_=m(x[:,:-1])
return float(torch.nn.functional.cross_entropy(logits.float().reshape(-1,logits.shape[-1]),x[0,1:],reduction='sum').double())
@torch.no_grad()
def nll_english93(m,rows,dev):
per={}
for r in rows:
x=torch.tensor(r['input_ids'],dtype=torch.long,device=dev)[None]
with sdpa_kernel(SDPBackend.MATH):
logits,_=m(x)
l=torch.nn.functional.cross_entropy(logits.float()[0],torch.tensor(r['target_ids'],device=dev),reduction='none')
mask=torch.tensor(r['score_mask'],device=dev)
per[r['document_id']]=per.get(r['document_id'],0.0)+float(l[mask].double().sum())
return per
def boot_delta(base,cand,seed=1337,draws=1000):
ids=sorted(base); rng=random.Random(seed); out=[]
for _ in range(draws):
s=rng.choices(ids,k=len(ids)); out.append(sum(cand[i][0]-base[i][0] for i in s)/sum(base[i][1] for i in s))
return sorted(out)[int(0.95*draws)-1]
def boot_ratio(base,cand,seed=1337,draws=1000):
ids=sorted(base); rng=random.Random(seed); out=[]
for _ in range(draws):
s=rng.choices(ids,k=len(ids)); out.append(sum(cand[i] for i in s)/sum(base[i] for i in s))
return sorted(out)[int(0.95*draws)-1]
def main():
ap=argparse.ArgumentParser(); ap.add_argument('--gates',required=True); ap.add_argument('--candidate',required=True)
ap.add_argument('--parent',required=True); ap.add_argument('--english599',default=None); ap.add_argument('--out',required=True)
ap.add_argument('--device',default='cuda:0'); ap.add_argument('--limit',type=int,default=None); a=ap.parse_args()
t0=time.time(); g=json.load(open(a.gates)); base=os.path.dirname(os.path.abspath(a.gates))
torch.backends.cuda.matmul.allow_tf32=False; torch.backends.cudnn.allow_tf32=False
for k,p in g['pinned_inputs'].items():
if os.path.exists(os.path.join(base,p['path'])): assert fsha(os.path.join(base,p['path']))==p['sha256'],k
dev=torch.device(a.device); L=a.limit
cand=load_model(torch.load(a.candidate,map_location='cpu',weights_only=True)['model'],dev)
par=load_model(torch.load(a.parent,map_location='cpu',weights_only=True)['model'],dev)
res={'gates_sha256':fsha(a.gates),'gates':{}}; rows={}
def windows(path): return json.load(open(os.path.join(base,path)))[:L]
for panel in ('sixth_final','final126'):
spec=g[panel]; ws=windows(spec['path']); bc={};cc={}
for w in ws:
bc[w['component_id']]=[nll_window(par,w['tokens'],dev),2048]; cc[w['component_id']]=[nll_window(cand,w['tokens'],dev),2048]
bce=sum(v[0] for v in bc.values())/sum(v[1] for v in bc.values()); cce=sum(v[0] for v in cc.values())/sum(v[1] for v in cc.values())
imp=1-cce/bce; ub=boot_delta(bc,cc)
ok=(imp>=spec['min_relative_ce_improvement']) and ub<spec['bootstrap_upper95_delta_strictly_below']
res['gates'][panel]=dict(parent_ce=bce,candidate_ce=cce,relative_ce_improvement=imp,bootstrap_upper95_delta=ub,units=len(bc),passed=ok)
rows[panel]={'parent':bc,'candidate':cc}
# English93 vs recorded A0
spec=g['retention_english93']; ew=[json.loads(l) for l in open(os.path.join(base,spec['windows_path']))]
if L: keep={r['document_id'] for r in ew[:L]}; ew=[r for r in ew if r['document_id'] in keep]
a0={(r:=json.loads(l))['id']:r['total_nll_nats'] for l in open(os.path.join(base,spec['baseline_rows_path']))}
c93=nll_english93(cand,ew,dev); b93={k:a0[k] for k in c93}
ub=boot_ratio(b93,c93); res['gates']['retention_english93']=dict(ratio=sum(c93.values())/sum(b93.values()),bootstrap_upper95_ratio=ub,units=len(c93),passed=ub<=spec['max_upper95_ratio'])
rows['english93']=c93
# English599 observed panel vs recorded step3403
spec=g['retention_english599']
if a.english599 and os.path.exists(a.english599) and fsha(a.english599)==spec['windows_sha256']:
ws=json.load(open(a.english599))[:L]
b={(r:=json.loads(l))['component_id']:r['nll_nats'] for l in open(os.path.join(base,spec['baseline_rows_path']))}
c={w['component_id']:nll_window(cand,w['tokens'],dev) for w in ws}; bb={k:b[k] for k in c}
ub=boot_ratio(bb,c); res['gates']['retention_english599']=dict(ratio=sum(c.values())/sum(bb.values()),bootstrap_upper95_ratio=ub,units=len(c),passed=ub<=spec['max_upper95_ratio'])
rows['english599']=c
else:
res['gates']['retention_english599']=dict(passed=False,status='PANEL_NOT_MOUNTED_OR_SHA_MISMATCH')
res['all_pass']=all(v.get('passed') for v in res['gates'].values())
res['decision']='RESEARCH_GAIN_ESTABLISHED_NO_AUTOMATIC_PROMOTION' if res['all_pass'] else 'NO_JOINT_RESEARCH_VALUE_ESTABLISHED'
res['limit']=L; res['elapsed_seconds']=round(time.time()-t0,1)
os.makedirs(a.out,exist_ok=True)
json.dump(res,open(os.path.join(a.out,'EVAL_RESULT.json'),'w'),indent=1); json.dump(rows,open(os.path.join(a.out,'EVAL_ROWS.json'),'w'))
print(json.dumps(res,indent=1))
if __name__=='__main__': main()
|