| """Question-only CUDA runtime for immutable semantic-head checkpoints. |
| |
| v1.1 (Solomon v1.1): |
| * precision: 'bf16' (default: bf16 weights, fp32 recurrence, the trainer's mode), 'fp32' (the qualified v1 |
| numerics) or 'int8' (torchao weight-only int8 on the decoder linears). Recorded in the identity. |
| * head layout: any subset of the ten semantic heads loads, provided the yes/no head boolean/state4 is present. |
| The v1.1 heads file carries ONE merged yes/no head (boolean/state4) and no entity/multilabel heads; the |
| serving layer routes multi-label branches to it through the binding's head routing table. A non-v1 layout |
| is recorded in the identity (answer_heads), so a binding pins it. |
| * the document prefill keeps its final hidden states (and any tapped layer) for the trained evidence head. |
| """ |
| import copy |
| from contextlib import contextmanager |
| import hashlib |
| import json |
| from pathlib import Path |
| import threading |
| import numpy as np |
| from solomon.engine_numerics import CudaEngine |
| from solomon.heads import semantic_head_key |
|
|
| def sha(path):return hashlib.sha256(Path(path).read_bytes()).hexdigest() |
|
|
| SEMANTIC_HEADS=('boolean/state4','entity/state4','multilabel/state4','single/choiceR','single/choiceS','single/sufficiency3', |
| 'ordered/choiceR','ordered/choiceS','ordered/sufficiency3','ordered/threshold4') |
| REQUIRED_HEADS=('boolean/state4',) |
|
|
|
|
| def load_heads(heads_path): |
| """{head_key: {'weight','bias'}} from a semantic-heads .npz. Tolerant of the merged v1.1 layout.""" |
| with np.load(heads_path,allow_pickle=False) as arrays: |
| heads={k[:-7]:{'weight':arrays[k].copy(),'bias':arrays[k[:-7]+'/bias'].copy()} |
| for k in arrays.files if k.endswith('/weight')} |
| unknown=sorted(set(heads)-set(SEMANTIC_HEADS)) |
| if unknown:raise ValueError('unknown semantic heads: '+', '.join(unknown)) |
| missing=[k for k in REQUIRED_HEADS if k not in heads] |
| if missing:raise ValueError('semantic heads file lacks the yes/no head: '+', '.join(missing)) |
| return heads |
|
|
| class SolomonEngine(CudaEngine): |
| def __init__(self, adapter, heads_path, *, expected_adapter=None, expected_heads=None, correctness=None, precision='bf16', |
| model_dir=None, prefix_layers=()): |
| if expected_adapter and sha(adapter)!=expected_adapter:raise ValueError('adapter hash mismatch') |
| if expected_heads and sha(heads_path)!=expected_heads:raise ValueError('heads hash mismatch') |
| self.answer_heads=load_heads(heads_path) |
| self.keep_prefix_hidden=True;self.prefix_layers=tuple(int(i) for i in prefix_layers) |
| super().__init__(adapter=adapter,placement='question',precision=precision,**({'model_dir':model_dir} if model_dir else {})) |
| self.base_identity=dict(self.identity) |
| for head in self.answer_heads.values(): |
| if head['weight'].shape!=(10,5120) or head['bias'].shape!=(10,) or not all(np.isfinite(v).all() for v in head.values()): |
| raise ValueError('invalid semantic head') |
| self.correctness_models=correctness or {} |
| self._task=None;self._lock=threading.RLock() |
| if len(self.answer_heads)!=10:self.identity['answer_heads']=','.join(sorted(self.answer_heads)) |
| self.identity.update(trained_heads_sha256=sha(heads_path),answer_engine_sha256=sha(__file__), |
| answer_projection='trained-semantic-head-float32',base_fingerprint=self.base_identity['fingerprint']) |
| self.identity.pop('fingerprint',None) |
| self.identity['fingerprint']=hashlib.sha256(json.dumps(self.identity,sort_keys=True).encode()).hexdigest() |
|
|
| @contextmanager |
| def task_context(self,task): |
| with self._lock: |
| previous=self._task;self._task=task |
| try:yield self |
| finally:self._task=previous |
|
|
| def ask(self,state,block,n_letters,execution='cached',head_key=None,tap_layers=()): |
| with self._lock: |
| key=head_key or semantic_head_key(self._task,block) |
| if key not in self.answer_heads or not 2<=n_letters<=10:raise ValueError('unknown semantic head or width') |
| if execution not in ('cached','full'):raise ValueError('invalid execution') |
| torch=self.torch;p=state['prefix_tokens'];previous=self.ctx['start'];handles=[];taps={} |
| def hook(index): |
| def capture(module,args,output): |
| h=output[0] if isinstance(output,tuple) else output |
| |
| taps[str(index)]=self.lm.norm(h[:,-1:])[0,-1].detach().float().cpu().numpy().copy() |
| return capture |
| try: |
| for index in tap_layers: |
| if type(index)is not int or not 0<=index<len(self.lm.layers):raise ValueError('invalid layer tap') |
| handles.append(self.lm.layers[index].register_forward_hook(hook(index))) |
| 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('question prefix mismatch') |
| self.ctx['start']=0 if execution=='cached' else p |
| if execution=='cached':hidden=self.lm(inputs_embeds=embeds[:,p:],position_ids=pos[...,p:],past_key_values=copy.deepcopy(state['cache']),use_cache=True).last_hidden_state |
| else:hidden=self.lm(inputs_embeds=embeds,position_ids=pos,use_cache=False).last_hidden_state |
| head=self.answer_heads[key] |
| logits=torch.nn.functional.linear(hidden[:,-1].float(),torch.as_tensor(head['weight'],device=hidden.device),torch.as_tensor(head['bias'],device=hidden.device))[0,:n_letters] |
| probabilities=logits.softmax(-1).cpu().numpy() |
| feature=hidden[0,-1].float().cpu().numpy().copy() |
| return {'letter_logits':logits.cpu().numpy(),'probabilities':probabilities,'prediction':int(probabilities.argmax()), |
| 'hidden':feature,'taps':taps,'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 |
| for handle in handles:handle.remove() |
|
|
| def predict_correctness(self,task,branches): |
| from solomon.head_training import predict_correctness |
| if task not in self.correctness_models:return {'reliability':None} |
| hidden=[np.asarray(r['hidden'],float) for r in branches] |
| probs=[np.asarray(r['probabilities'],float) for r in branches] |
| feature=np.r_[np.mean(hidden,axis=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))] |
| return {'reliability':predict_correctness(self.correctness_models[task],feature)} |
|
|
