File size: 13,978 Bytes
1d2de8a | 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 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 | """Final normalized feature extraction, preserving the qualified CUDA engine."""
import hashlib,json
from pathlib import Path
import numpy as np
def sha(path):return hashlib.sha256(Path(path).read_bytes()).hexdigest()
def token_ids_for(engine,state,block,n):
text=engine._render(state['parts'],block);raw=engine.t.encode(text,add_special_tokens=False);ids=[]
for c in 'ABCDEFGHIJ'[:n]:
ext=engine.t.encode(text+c,add_special_tokens=False)
if ext[:-1]!=raw or len(ext)!=len(raw)+1:raise ValueError('unstable answer token continuation')
ids.append(ext[-1])
return ids
def capture(engine,state,block,n):
"""Prehook sees precisely the post-final-norm input used by the vocabulary head."""
captured=[]
def hook(module,args):
h=args[0]
if h.ndim!=2 or h.shape[0]!=1:raise ValueError('unexpected final head input shape')
captured.append(h.detach().float().cpu().numpy()[0].copy())
handle=engine.model.lm_head.register_forward_pre_hook(hook)
try:result=engine.ask(state,block,n,execution='cached')
finally:handle.remove()
if len(captured)!=1:raise ValueError('expected exactly one vocabulary projection')
ids=token_ids_for(engine,state,block,n)
if not np.isfinite(captured[0]).all():raise ValueError('nonfinite hidden features')
return captured[0],ids,result
def copy_letter_head(engine):
ids=[]
for c in 'ABCDEFGHIJ':
pieces=engine.t.encode(c,add_special_tokens=False)
if len(pieces)!=1:raise ValueError('letters must be one token')
ids.append(pieces[0])
head=engine.model.lm_head;w=head.weight.detach()[ids].float().cpu().numpy().copy();b=getattr(head,'bias',None)
return {'weight':w,'bias':np.zeros(10,np.float32) if b is None else b.detach()[ids].float().cpu().numpy().copy(),'token_ids':np.array(ids,np.int64),'vocabulary_size':np.array(head.weight.shape[0],np.int64),'bias_present':np.array(b is not None)}
def atomic_json(path,value):
path=Path(path);tmp=path.with_suffix(path.suffix+'.tmp');tmp.write_text(json.dumps(value,sort_keys=True,indent=2));tmp.replace(path)
def write_document(path,features,token_ids,results,metadata,identity):
"""NPZ first, then checksummed completion JSON. No cumulative shard rewrite."""
path=Path(path);n=len(features)
if not n or not(len(token_ids)==len(results)==len(metadata)==n):raise ValueError('document arrays must align')
h=np.asarray(features,np.float32);logits=np.zeros((n,10),np.float32);tokens=np.full((n,10),-1,np.int64);widths=[]
if h.ndim!=2 or not np.isfinite(h).all():raise ValueError('invalid document feature matrix')
rows=[]
for i,(ids,r,m) in enumerate(zip(token_ids,results,metadata)):
values=np.asarray(r['letter_logits'],np.float32);width=len(ids)
if not 2<=width<=10 or values.shape!=(width,) or not np.isfinite(values).all():raise ValueError('invalid teacher row')
logits[i,:width]=values;tokens[i,:width]=ids;widths.append(width)
rows.append({**m,'n':width,'mass':float(r['mass']),'top_is_letter':bool(r['top_is_letter']),'prompt_tokens':int(r['prompt_tokens']),'branch_tokens':int(r['branch_tokens']),'reused_prefix_tokens':int(r['reused_prefix_tokens'])})
tmp=path.with_suffix('.npz.tmp')
with tmp.open('wb') as f:np.savez(f,hidden=h,teacher_logits=logits,token_ids=tokens,widths=np.asarray(widths,np.int64))
tmp.replace(path)
report={'complete':True,'identity':identity,'rows':rows,'npz_sha256':sha(path),'feature_shape':list(h.shape),'feature_kind':'final-normalized-lm-head-input','storage':'uncompressed-npz-float32'}
atomic_json(path.with_suffix('.json'),report)
return report
def validate_document(path,identity,expected_ids):
path=Path(path);meta=path.with_suffix('.json')
if not path.exists() or not meta.exists():return False
r=json.loads(meta.read_text())
if r.get('complete')is not True or r.get('identity')!=identity or r.get('npz_sha256')!=sha(path) or [x['id'] for x in r['rows']]!=list(expected_ids):raise ValueError('resumed document identity mismatch')
return True
class TrainedHeads:
"""Validated saved answer/correctness heads. No backbone or label access."""
def __init__(self,directory,*,expected_report_sha256=None,expected_weights_sha256=None):
directory=Path(directory);report_path=directory/'report.json';weights_path=directory/'weights.npz'
self.report_sha256=sha(report_path);self.weights_sha256=sha(weights_path)
if expected_report_sha256 is not None and self.report_sha256!=expected_report_sha256:raise ValueError('selected head report changed')
if expected_weights_sha256 is not None and self.weights_sha256!=expected_weights_sha256:raise ValueError('selected head weights changed')
self.report=json.loads(report_path.read_text())
if self.report['weights_sha256']!=self.weights_sha256:raise ValueError('trained head weight checksum mismatch')
self.models={}
with np.load(weights_path,allow_pickle=False) as archive:
for kind,entries in self.report['models'].items():
self.models[kind]={}
for key,fields in entries.items():
model={name:archive[value['array']].astype(np.float64) if isinstance(value,dict) and set(value)=={'array'} else value for name,value in fields.items()}
if any(not np.isfinite(value).all() for value in model.values() if isinstance(value,np.ndarray)):raise ValueError('nonfinite trained head')
if kind=='heads':
if model['weight'].ndim!=2 or model['weight'].shape[0]!=10 or model['bias'].shape!=(10,):raise ValueError('invalid answer head shape')
elif kind=='correctness':
if model['mean'].ndim!=1 or model['scale'].shape!=model['mean'].shape or np.any(model['scale']<=0) or model['weights'].shape!=(len(model['mean'])+1,):raise ValueError('invalid correctness head shape')
else:raise ValueError('unknown trained model kind')
self.models[kind][key]=model
if not self.models.get('heads') or not self.models.get('correctness'):raise ValueError('both trained answer and correctness heads required')
def answer(self,hidden,key,n):
if key not in self.models['heads']:raise ValueError('no trained answer head for '+str(key))
if type(n)is not int or not 2<=n<=10:raise ValueError('invalid head output width')
h=np.asarray(hidden,np.float64);model=self.models['heads'][key]
if h.shape!=(model['weight'].shape[1],) or not np.isfinite(h).all():raise ValueError('invalid hidden state for answer head')
return (model['weight']@h+model['bias'])[:n]
def correctness(self,task,branches):
from scipy.special import expit
if task not in self.models['correctness'] or not branches:raise ValueError('no trained whole-answer correctness head')
hidden=np.asarray([r['hidden'] for r in branches],np.float64);probs=[]
if hidden.ndim!=2 or not np.isfinite(hidden).all():raise ValueError('invalid correctness hidden representations')
for row in branches:
p=np.asarray(row['probabilities'],float)
if p.ndim!=1 or len(p)<2 or not np.isfinite(p).all() or np.any(p<0) or abs(p.sum()-1)>1e-6:raise ValueError('invalid correctness answer distribution')
probs.append(p)
feature=np.r_[hidden.mean(0),min(float(p.max()) for p in probs),min(float(np.sort(p)[-1]-np.sort(p)[-2]) for p in probs),max(float(-np.sum(p*np.log(np.maximum(p,1e-300)))/np.log(len(p))) for p in probs),float(len(branches)),float(sum(len(p) for p in probs))]
model=self.models['correctness'][task]
if feature.shape!=model['mean'].shape:raise ValueError('whole-answer confidence feature width mismatch')
value=float(expit(np.r_[1.,(feature-model['mean'])/model['scale']]@model['weights']))
return {'reliability':value,'features':feature.tolist(),'task':task,'target':'whole_return_exact_correctness','calibration_established':False,'heads_sha256':self.weights_sha256}
def semantic_head_key(task,block):
from solomon.engine_contract import SUFFICIENCY,SINGLE_R,SINGLE_S,ORDERED_R,ORDERED_S,LABEL,BOOLEAN_TASK
if task not in ('boolean','single','ordered','multilabel','entity'):raise ValueError('explicit answering task required')
if block.startswith(SUFFICIENCY) and task in ('single','ordered'):return task+'/sufficiency3'
for prefix,t,form in ((SINGLE_R,'single','R'),(SINGLE_S,'single','S'),(ORDERED_R,'ordered','R'),(ORDERED_S,'ordered','S')):
if block.startswith(prefix) and task==t:return task+'/choice'+form
if block.startswith(LABEL) and task=='multilabel':return 'multilabel/state4'
if block.startswith(BOOLEAN_TASK):
if task in ('boolean','entity'):return task+'/state4'
if task=='ordered':return 'ordered/threshold4'
raise ValueError('task block has no trained semantic head mapping')
from solomon.engine_reasoning import ReasoningEngine
class HeadsEngine(ReasoningEngine):
"""Question-only backbone with trained small heads; native LM head retained for reasoning.
Wrap a complete readout in task_context(task), or supply ask(head_key=...).
Answer projection is deliberately evaluated in CPU float64 on the exact copied
float32 hidden vector to match saved-feature evaluation; no speedup is claimed.
"""
def __init__(self,heads_dir,*,expected_base_fingerprint,expected_report_sha256=None,expected_weights_sha256=None,**kwargs):
import threading
self.heads=TrainedHeads(heads_dir,expected_report_sha256=expected_report_sha256,expected_weights_sha256=expected_weights_sha256)
super().__init__(**kwargs)
if self.identity['fingerprint']!=expected_base_fingerprint:raise ValueError('trained head backbone/adapter runtime mismatch')
self.base_identity=dict(self.identity);self._head_task=None;self._head_lock=threading.RLock()
self.identity={**self.identity,'base_fingerprint':expected_base_fingerprint,'trained_heads_sha256':self.heads.weights_sha256,'trained_heads_report_sha256':self.heads.report_sha256,'heads_runtime_sha256':sha(__file__),'answer_projection':'trained-semantic-head-float64-cpu'}
self.identity.pop('fingerprint',None);self.identity['fingerprint']=hashlib.sha256(json.dumps(self.identity,sort_keys=True).encode()).hexdigest()
def task_context(self,task):
from contextlib import contextmanager
@contextmanager
def context():
with self._head_lock:
previous=self._head_task;self._head_task=task
try:yield self
finally:self._head_task=previous
return context()
def prefill(self,document):
with self._head_lock:return super().prefill(document)
def ask(self,state,block,n_letters,execution='cached',head_key=None):
import copy
with self._head_lock:
key=head_key or semantic_head_key(self._head_task,block)
if key not in self.heads.models['heads']:raise ValueError('missing trained head '+key)
if self.placement!='question':raise ValueError('trained heads require question-only adaptation')
if execution not in ('cached','full'):raise ValueError('invalid execution')
torch=self.torch;p=state['prefix_tokens'];previous=self.ctx['start']
try:
with torch.inference_mode():
text=self._render(state['parts'],block);ids,embeds,pos,_=self._encode(state['parts'],text,state['features'])
if not torch.equal(ids[:,:p],state['prefix_ids']):raise ValueError('head question prefix mismatch')
self.ctx['start']=(0 if execution=='cached' or self.placement=='full' else p) if self.adapter else None
if execution=='cached':h=self.lm(inputs_embeds=embeds[:,p:],position_ids=pos[...,p:],past_key_values=copy.deepcopy(state['cache']),use_cache=True).last_hidden_state
else:h=self.lm(inputs_embeds=embeds,position_ids=pos,use_cache=False).last_hidden_state
hidden=h[0,-1].float().cpu().numpy().copy()
logits=self.heads.answer(hidden,key,n_letters);probs=np.exp(logits-logits.max());probs/=probs.sum()
return {'letter_logits':logits,'probabilities':probs,'prediction':int(probs.argmax()),'hidden':hidden.tolist(),'head_key':key,'mass':None,'top_is_letter':None,'mass_available':False,'execution':execution,'fallback':'','prompt_tokens':int(ids.shape[1]),'branch_tokens':int(ids.shape[1])-p,'input_tokens':int(ids.shape[1])-p,'reused_prefix_tokens':p if execution=='cached' else 0,'generation_tokens':0,'fingerprint':self.identity['fingerprint']}
finally:self.ctx['start']=previous
def estimate(self,state,block,n,stage='fast',max_new_tokens=512):
if stage!='fast' or not any('image' in part for part in state['parts']):
return super().estimate(state,block,n,stage=stage,max_new_tokens=max_new_tokens)
if type(n)is not int or not 2<=n<=10:raise ValueError('invalid letter count')
# Reuse prefetched vision features; no transformer branch is evaluated.
with self._head_lock,self.torch.inference_mode():
ids,_,_,_=self._encode(state['parts'],self._render(state['parts'],block),state['features'])
prefix=state['prefix_tokens']
if not self.torch.equal(ids[:,:prefix],state['prefix_ids']):raise ValueError('image budget prefix mismatch')
return {'input_tokens':int(ids.shape[1])-prefix,'generated_tokens':0,'branches':1,
'full_prompt_tokens':int(ids.shape[1]),'reused_prefix_tokens':prefix}
def predict_correctness(self,task,branches):return self.heads.correctness(task,branches)
def reason(self,*args,**kwargs):
with self._head_lock:return super().reason(*args,**kwargs)
|