| """The confidence layer text reasoning on immutable answer-contract base-prefix states.""" |
| import copy,hashlib,json |
| from pathlib import Path |
| from solomon.engine_numerics import CudaEngine |
|
|
| FORCE='\n</think>\n\nGive only the answer letter: ' |
|
|
| def validate_request(state,n_letters,max_tokens): |
| if not isinstance(max_tokens,int) or isinstance(max_tokens,bool) or not 1<=max_tokens<=3072:raise ValueError('reasoning cap must be 1..3072') |
| if not isinstance(n_letters,int) or isinstance(n_letters,bool) or not 2<=n_letters<=10:raise ValueError('invalid letter count') |
| if any('image' in p for p in state['parts']):raise ValueError('the confidence layer reasoning confidence supports text only') |
|
|
| class ReasoningEngine(CudaEngine): |
| def reasoning_identity(self): |
| return {'base_runtime':self.identity,'reasoning_code_sha256':hashlib.sha256(Path(__file__).read_bytes()).hexdigest(),'adapter':'disabled','forced_suffix':FORCE} |
|
|
| def estimate(self,state,block,n,stage='fast',max_new_tokens=512): |
| if stage not in ('fast','reasoning'):raise ValueError('unknown stage') |
| if any('image' in p for p in state['parts']):raise ValueError('the confidence layer text-only budget estimator') |
| text=self._render(state['parts'],block) |
| if stage=='reasoning': |
| validate_request(state,n,max_new_tokens) |
| marker='<|im_start|>assistant\n' |
| if marker not in text:raise ValueError('unknown assistant template') |
| text=text[:text.rfind(marker)]+marker+'<think>\n' |
| ids=self.t.encode(text,add_special_tokens=False) |
| prefix=state['prefix_ids'][0].tolist() |
| if ids[:len(prefix)]!=prefix:raise ValueError('budget prefix mismatch') |
| extra=len(self.t.encode(FORCE,add_special_tokens=False)) if stage=='reasoning' else 0 |
| return {'input_tokens':len(ids)-len(prefix)+extra,'generated_tokens':max_new_tokens if stage=='reasoning' else 0,'branches':1,'full_prompt_tokens':len(ids),'reused_prefix_tokens':len(prefix)} |
|
|
| def _reason_inputs(self,state,block): |
| text=self._render(state['parts'],block) |
| marker='<|im_start|>assistant\n' |
| if marker not in text:raise ValueError('unknown assistant template') |
| text=text[:text.rfind(marker)]+marker+'<think>\n' |
| ids,embeds,pos,_=self._encode(state['parts'],text,state['features']) |
| if not self.torch.equal(ids[:,:state['prefix_tokens']],state['prefix_ids']):raise ValueError('reasoning prefix mismatch') |
| return text,ids,embeds,pos |
|
|
| def reason(self,state,block,n_letters,max_tokens=512,execution='cached',replay_tokens=None,cap_checkpoints=None): |
| """Greedy base reasoning then bounded forced letter readout; replay is qualification only.""" |
| validate_request(state,n_letters,max_tokens) |
| checkpoints=tuple(cap_checkpoints or ()) |
| if any(not isinstance(c,int) or isinstance(c,bool) or not 1<=c<=max_tokens for c in checkpoints):raise ValueError('invalid cap checkpoints') |
| if checkpoints and (replay_tokens is not None or execution!='cached'):raise ValueError('checkpoints require cached generation') |
| if self.placement!='question':raise ValueError('reasoning needs base document state') |
| if execution not in ('cached','full'):raise ValueError('bad execution') |
| torch=self.torch;p=state['prefix_tokens'];previous=self.ctx['start'] |
| from transformers import DynamicCache |
| try: |
| self.ctx['start']=None |
| with torch.inference_mode(): |
| text,ids,embeds,pos=self._reason_inputs(state,block) |
| cache=copy.deepcopy(state['cache']) if execution=='cached' else DynamicCache(config=self.lm.config) |
| start=p if execution=='cached' else 0 |
| h=self.lm(inputs_embeds=embeds[:,start:],position_ids=pos[...,start:],past_key_values=cache,use_cache=True).last_hidden_state |
| generated=[];snapshots={};stop='token_cap';lastpos=pos[...,-1:];close=self.t.encode('</think>',add_special_tokens=False) |
| eos=self.t.eos_token_id |
| for i in range(0 if execution=='full' and replay_tokens is not None else max_tokens): |
| token=int(replay_tokens[i]) if replay_tokens is not None and i<len(replay_tokens) else int(self.model.lm_head(h[:,-1]).argmax(-1)) |
| if replay_tokens is not None and i>=len(replay_tokens):break |
| generated.append(token) |
| one=torch.tensor([[token]],device=ids.device);lastpos=lastpos+1 |
| h=self.lm(input_ids=one,position_ids=lastpos,past_key_values=cache,use_cache=True).last_hidden_state |
| if len(generated) in checkpoints: |
| sf=self.t.encode(FORCE,add_special_tokens=False);sfids=torch.tensor([sf],device=ids.device) |
| sfpos=lastpos+torch.arange(1,len(sf)+1,device=ids.device).view(1,1,-1) |
| sh=self.lm(input_ids=sfids,position_ids=sfpos,past_key_values=copy.deepcopy(cache),use_cache=True).last_hidden_state |
| sl=self.model.lm_head(sh[:,-1]).float()[0] |
| li=[self.t.encode(c,add_special_tokens=False)[0] for c in 'ABCDEFGHIJ'[:n_letters]] |
| sc=sl[li];sp=sc.softmax(-1) |
| snapshots[str(len(generated))]={'letter_logits':sc.cpu().tolist(),'probabilities':sp.cpu().tolist(),'prediction':int(sp.argmax()),'mass':float(sl.softmax(-1)[li].sum()),'generation_tokens':len(generated),'finish_reason':'token_cap'} |
| if replay_tokens is None and (token==eos or generated[-len(close):]==close):stop='closed_reasoning' if token!=eos else 'eos';break |
| if execution=='full' and replay_tokens is not None: |
| generated=list(replay_tokens[:max_tokens]);lastpos=pos[...,-1:]+len(generated) |
| suffix=self.t.encode(FORCE,add_special_tokens=False);suffix_ids=torch.tensor([suffix],device=ids.device) |
| suffix_pos=lastpos+torch.arange(1,len(suffix)+1,device=ids.device).view(1,1,-1) |
| if execution=='full' and replay_tokens is not None: |
| extra=torch.tensor([generated+suffix],device=ids.device) |
| all_embeds=torch.cat((embeds,self.model.get_input_embeddings()(extra)),dim=1) |
| extra_pos=pos[...,-1:]+torch.arange(1,extra.shape[1]+1,device=ids.device).view(1,1,-1) |
| h=self.lm(inputs_embeds=all_embeds,position_ids=torch.cat((pos,extra_pos),dim=-1),use_cache=False).last_hidden_state |
| else: |
| h=self.lm(input_ids=suffix_ids,position_ids=suffix_pos,past_key_values=cache,use_cache=True).last_hidden_state |
| logits=self.model.lm_head(h[:,-1]).float()[0] |
| letter_ids=[] |
| for c in 'ABCDEFGHIJ'[:n_letters]: |
| token_ids=self.t.encode(c,add_special_tokens=False) |
| if len(token_ids)!=1:raise ValueError('non-single letter token') |
| letter_ids.append(token_ids[0]) |
| chosen=logits[letter_ids];probs=chosen.softmax(-1) |
| return {'cap_checkpoints':snapshots,'letter_logits':chosen.cpu().tolist(),'probabilities':probs.cpu().tolist(),'mass':float(logits.softmax(-1)[letter_ids].sum()),'prediction':int(probs.argmax()),'generated_tokens':generated,'generation_tokens':len(generated),'text':self.t.decode(generated),'finish_reason':stop,'forced_answer':True,'execution':execution,'prompt_tokens':int(ids.shape[1]),'input_tokens':int(ids.shape[1])-p+len(suffix),'branch_tokens':int(ids.shape[1])-p+len(suffix),'reused_prefix_tokens':p if execution=='cached' else 0,'adapter_enabled':False,'identity':self.reasoning_identity()} |
| finally:self.ctx['start']=previous |
|
|
| def exit_features(self,state,block,n_letters,layers=(31,47,55)): |
| """Full-depth feature capture only; explicitly does not execute early exit.""" |
| captured={};handles=[] |
| def hook(layer): |
| def capture(module,args,output): |
| value=output[0] if isinstance(output,tuple) else output |
| captured[layer]=value[:,-1].detach().clone() |
| return capture |
| try: |
| for layer in layers:handles.append(self.lm.layers[layer].register_forward_hook(hook(layer))) |
| full=self.ask(state,block,n_letters) |
| result={} |
| with self.torch.inference_mode(): |
| text=self._render(state['parts'],block);raw=self.t.encode(text,add_special_tokens=False) |
| letter_ids=[self.t.encode(text+c,add_special_tokens=False)[-1] for c in 'ABCDEFGHIJ'[:n_letters]] |
| for layer,value in captured.items(): |
| logits=self.model.lm_head(self.lm.norm(value)).float()[0,letter_ids] |
| result[str(layer)]={'feature':value.float().cpu().tolist()[0],'letter_logits':logits.cpu().tolist(),'prediction':int(logits.argmax())} |
| return {'full':{k:(v.tolist() if hasattr(v,'tolist') else v) for k,v in full.items()},'layers':result,'executed_truncated':False,'total_layers':len(self.lm.layers)} |
| finally: |
| for handle in handles:handle.remove() |
|
|
