evaluation.py

1"""Frozen retention and memorization diagnostics; all training inputs stay on disk."""
2import json,math,sys
3from pathlib import Path
4import numpy as np
5import torch
6from reader import Records,SmallRecords,H
7sys.path.append(str(H.parent/'cm6-curriculum-v1'))
8from composite import phase,tensors
9from exact_metrics import summarize
10
11def raw(net,rows,device):
12    predictions=[]
13    for start in range(0,len(rows),64):
14        x,_,m=[v.to(device) for v in tensors(rows[start:start+64])];x=x[:,:120];m=m[:,:120];parts=[]
15        for _ in range(3):
16            with torch.autocast('cuda',dtype=torch.bfloat16,enabled=device.type=='cuda'):logits=net.next_logits(x,m)
17            if not torch.isfinite(logits).all():raise FloatingPointError('nonfinite raw diagnostic logits')
18            t=logits.argmax(-1)
19            parts.append(t.cpu().tolist());x=torch.cat((x,t[:,None]),1);m=torch.cat((m,torch.ones_like(t[:,None],dtype=m.dtype)),1)
20        predictions.extend(zip(*parts))
21    return [list(x) for x in predictions]
22
23def losses(net,rows,device):
24    values=[];hits=np.zeros(3,dtype=np.int64)
25    for start in range(0,len(rows),64):
26        x,y,m=[v.to(device) for v in tensors(rows[start:start+64])]
27        with torch.autocast('cuda',dtype=torch.bfloat16,enabled=device.type=='cuda'):logits=net(x,m)
28        values.extend(torch.nn.functional.cross_entropy(logits.float().reshape(-1,512),y.reshape(-1),reduction='none').reshape(-1,3).sum(-1).cpu().tolist())
29        hits+=(logits.argmax(-1)==y).sum(0).cpu().numpy()
30    return np.asarray(values),hits
31
32def evaluate(net,cfg,device):
33    net.eval();groups={};exact_outcomes=[];total_loss=0.;total_n=0
34    with torch.inference_mode():
35        for spec in cfg['validation_groups']:
36            source=(Records if spec['kind']==0 else SmallRecords)(spec['root'],spec['source'],spec['split'],verify=True)
37            rows=source.take(range(len(source)));value,correct=losses(net,rows,device);p=phase(rows)
38            groups[spec['name']]={'decisions':len(rows),'loss_sum':float(value.sum()),'nll':float(value.sum()/(3*len(rows))),
39                'phase_nll':[float(value[p==k].mean()/3) if np.any(p==k) else None for k in (0,1,2)],'phase_accuracy':(correct/len(rows)).tolist()}
40            total_loss+=float(value.sum());total_n+=len(rows)
41            if spec['kind']==1:
42                metadata=[json.loads(line) for line in Path(spec['tie_metadata']).read_text().splitlines()]
43                for row,prediction in zip(rows,raw(net,rows,device)):
44                    item=metadata[int(row['game'])];exact_outcomes.append({'key':item['position_key'],'game':item['source_sequence_sha256'],'correct':int(prediction in item['tied_target_triples'])})
45            if hasattr(source,'close'):source.close()
46        root=Path(cfg['prepared_validation']);book_targets=json.loads((root/'book-probe-targets.json').read_text());book={}
47        for name in ('book-training-probe','book-validation-states'):
48            rows=np.load(root/(name+'.npy'),allow_pickle=False);predictions=raw(net,rows,device)
49            hits=[int(p in book_targets[str(int(row['game']))]) for row,p in zip(rows,predictions)]
50            with_history=np.any(rows['attention'][:,2:50],axis=1)
51            book[name]={'decisions':len(rows),'accuracy':sum(hits)/len(hits),
52                'history_accuracy':float(np.mean(np.asarray(hits)[with_history])) if with_history.any() else None,
53                'fen_only_accuracy':float(np.mean(np.asarray(hits)[~with_history])) if (~with_history).any() else None}
54        pairs=np.load(root/'book-validation-pairs.npy',allow_pickle=False);book_loss,_=losses(net,pairs,device);pair_nll=float(book_loss.mean()/3)
55        legal_targets=json.loads((root/'legality-targets.json').read_text());legal_rows=np.concatenate([np.load(root/f'legality-source{s}.npy',allow_pickle=False) for s in (0,1)])
56        legal_hits=[int(p in target['triples']) for p,target in zip(raw(net,legal_rows,device),legal_targets)]
57    ordinary=sum(groups[f'0-{s}']['loss_sum'] for s in (0,1))/(3*sum(groups[f'0-{s}']['decisions'] for s in (0,1)))
58    metrics={'ordinary_nll':ordinary,'source_nll':{str(s):groups[f'0-{s}']['nll'] for s in (0,1)},
59        'phase_nll':{f'{s}/{k}':groups[f'0-{s}']['phase_nll'][k] for s in (0,1) for k in (0,1,2)},
60        'exact_accuracy':sum(x['correct'] for x in exact_outcomes)/len(exact_outcomes),'legality':sum(legal_hits)/len(legal_hits),
61        'book_training_recall':book['book-training-probe']['accuracy'],'book_validation_pair_nll':pair_nll}
62    metrics['finite']=all(math.isfinite(x) for x in [ordinary,pair_nll,*metrics['source_nll'].values(),*(v for v in metrics['phase_nll'].values() if v is not None)])
63    net.train();return {'nll':total_loss/(3*total_n),'decisions':total_n,'groups':groups,'exact_outcomes':exact_outcomes,'exact_cluster_accuracy':summarize(exact_outcomes),
64        'retention':metrics,'book':book,'book_validation_pair_nll':pair_nll,'raw_legality':metrics['legality'],
65        'definition':'selection is minimum pooled canonical-label NLL on unchanged ordinary/exact slice; book recall on 2048 training states and all 229 held-out pairs are diagnostic only; raw legality on 1024 frozen ordinary states'}
66