| """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') |
| |
| 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) |
|
|
