Solomon / src /solomon /serving.py
Archer Hume
Solomon v1.1
105f9ef
Raw
History Blame Contribute Delete
4.01 kB
"""Preserve five-task trained-head confidence/routing with optional evidence."""
import hashlib,json
from pathlib import Path
import numpy as np
from solomon.engine import SolomonEngine,sha
from solomon.engine_reasoning import ReasoningEngine
from solomon.service_heads import HeadsService
from solomon.service_answers import Service as ConfidenceService
from solomon.service_evidence import EvidenceService,CONTRACT
class ServiceEngine(SolomonEngine,ReasoningEngine):
def __init__(self,*args,correctness_report_sha256=None,**kwargs):
super().__init__(*args,**kwargs)
self.identity.update(trained_heads_report_sha256=correctness_report_sha256,serving_sha256=sha(__file__))
self.identity.pop('fingerprint',None)
self.identity['fingerprint']=hashlib.sha256(json.dumps(self.identity,sort_keys=True).encode()).hexdigest()
def estimate(self,state,block,n,stage='fast',max_new_tokens=512):
if not any('image' in p for p in state['parts']):
return ReasoningEngine.estimate(self,state,block,n,stage=stage,max_new_tokens=max_new_tokens)
if stage!='fast':raise ValueError('image reasoning is unqualified')
with self._lock,self.torch.inference_mode():
ids,_,_,_=self._encode(state['parts'],self._render(state['parts'],block),state['features'])
p=state['prefix_tokens']
if not self.torch.equal(ids[:,:p],state['prefix_ids']):raise ValueError('image estimate prefix mismatch')
return {'input_tokens':int(ids.shape[1])-p,'generated_tokens':0,'branches':1,
'full_prompt_tokens':int(ids.shape[1]),'reused_prefix_tokens':p}
class ConfidenceHeadsService(HeadsService):
def expected_identity(self,stage):
return {**super().expected_identity(stage),'contract':CONTRACT,'serving_sha256':sha(__file__)}
def _validate_artifact(self,stage,artifact):
ConfidenceService._validate_artifact(self,stage,artifact)
if artifact.get('schema')!='solomon-heads-envelope-v1':raise ValueError('the evidence layer confidence envelope required')
selected=artifact.get('trained_head_identity',{});runtime=self.runtime_identity
if (selected.get('weights_sha256')!=runtime.get('trained_heads_sha256') or
selected.get('report_sha256')!=runtime.get('trained_heads_report_sha256') or
selected.get('feature_identity',{}).get('runtime_fingerprint')!=runtime.get('base_fingerprint') or
selected.get('feature_identity',{}).get('adapter_sha256')!=runtime.get('adapter_sha256')):
raise ValueError('the evidence layer correctness/head/backbone identity mismatch')
class FreshDecider:
def __init__(self,engine,design,store):self.engine=engine;self.design=design;self.store=store
def __call__(self,parts,task,request):
from solomon.service_states import Service as Readout
engine=self.engine
class Capture:
def __init__(self):self.records=[]
def ask(self,*args,**kwargs):
result=engine.ask(*args,**kwargs);self.records.append(result);return result
capture=Capture();state=engine.prefill(parts or [{'text':''}])
class Fixed(Readout):
def _warm(self,key):return state
readout=Fixed(self.store,capture,self.design)
with engine.task_context(task):answer=readout.ask('fresh-evidence',task,**request)
return {'answer':answer,'cost':{'input_tokens':state['prefix_tokens']+sum(r['branch_tokens'] for r in capture.records),
'branches':len(capture.records),'prefill_tokens':state['prefix_tokens']}}
def service(store,engine,design,*,artifact=None,views_artifact=None,views_tasks=(),selector=None,page_selector=None):
answers=ConfidenceHeadsService(store,engine,design,artifact=artifact,views_artifact=views_artifact,views_tasks=views_tasks)
return EvidenceService(answers,FreshDecider(engine,design,Path(store)/'.evidence-readout'),selector,page_selector)