| """Model-side input/output codec and registered analytic circuits. |
| |
| The codec parses JSON and tokenizes text. It does not call Python solvers, |
| sympy, regex validators, SAT packages or external programs to compute answers. |
| """ |
| import json |
| import torch |
| from torch import nn |
| from .circuits import (SchemaCircuit,WorkflowCircuit,IndexCircuit,SyntaxScopeCircuit, |
| PolynomialCircuit,ConstraintCircuit,UnitCircuit,SATCircuit, |
| ClosureCircuit,BinaryCircuit,UNIT_NAMES) |
|
|
| FEATURES=('schema','workflow','index','syntax_scope','polynomial','constraints','units','sat','closure','binary') |
|
|
|
|
| class AnalyticEngine(nn.Module): |
| def __init__(self,spec=None): |
| super().__init__() |
| defaults=dict(fields=16,key_bytes=32,value_bytes=64,enums=8,nodes=32, |
| docs=64,doc_tokens=8192,query_tokens=64,text_bytes=65536,stack=32,scopes=8,symbols=64, |
| degree=8,coefficient_bound=1024,sat_variables=12,sat_clauses=32,bytes=512) |
| if spec and set(spec)-set(defaults):raise ValueError('Unknown analytic capacity') |
| self.spec={**defaults,**(spec or {})} |
| limits={**defaults,'fields':32,'key_bytes':128,'value_bytes':256,'enums':32, |
| 'text_bytes':1048576,'stack':64,'scopes':16,'symbols':256,'bytes':1024} |
| for key,value in self.spec.items(): |
| if type(value) is not int or not 1<=value<=limits[key]:raise ValueError('Unsupported analytic capacity: '+key) |
| s=self.spec |
| self.layers=nn.ModuleDict(dict(schema=SchemaCircuit(),workflow=WorkflowCircuit(),index=IndexCircuit(), |
| syntax_scope=SyntaxScopeCircuit(),polynomial=PolynomialCircuit(s['degree']), |
| constraints=ConstraintCircuit(),units=UnitCircuit(),sat=SATCircuit(s['sat_variables']), |
| closure=ClosureCircuit(),binary=BinaryCircuit())) |
| self.register_buffer('enabled',torch.ones(len(FEATURES),dtype=torch.bool)) |
| self.register_buffer('calls',torch.zeros(len(FEATURES),dtype=torch.long),persistent=False) |
|
|
| @property |
| def device(self):return self.enabled.device |
|
|
| def prepare(self): |
| torch.set_float32_matmul_precision('highest') |
| if self.device.type=='cuda':torch.backends.cuda.matmul.allow_tf32=False |
| for name,parameter in self.named_parameters(): |
| double=(name.startswith('layers.polynomial.') or name.startswith('layers.units.units.') or name.startswith('layers.units.combine.')) |
| expected=torch.float64 if double else torch.float32 |
| if parameter.dtype!=expected or parameter.requires_grad: |
| raise TypeError('Analytic coefficients require their compiled FP32/FP64 dtype and frozen gradients: '+name) |
|
|
| def set_enabled(self,names): |
| if set(names)-set(FEATURES):raise ValueError('Unknown analytic feature') |
| self.enabled.copy_(torch.tensor([name in names for name in FEATURES],device=self.device)) |
|
|
| def forward(self,feature,**kwargs): |
| self.prepare() |
| i=FEATURES.index(feature) |
| if not bool(self.enabled[i]):raise RuntimeError(f'Feature disabled: {feature}') |
| self.calls[i].add_(1) |
| return self.layers[feature](**kwargs) |
|
|
| def tensor(self,data,dtype=torch.float32):return torch.tensor(data,device=self.device,dtype=dtype) |
|
|
| def text(self,text,limit=None): |
| if not isinstance(text,str):raise TypeError('Expected text') |
| encoded=text.encode('utf-8') |
| if len(encoded)>(limit or self.spec['text_bytes']):raise ValueError('Text capacity exceeded') |
| return self.tensor(list(encoded),torch.long) |
|
|
| @staticmethod |
| def integer(value,lo,hi): |
| if type(value) is not int or not lo<=value<=hi:raise ValueError(f'Integer must be in [{lo},{hi}]') |
| return value |
|
|
| @torch.inference_mode() |
| def run(self,feature,payload,tokenizer=None): |
| self.prepare() |
| if feature not in FEATURES:raise ValueError('Unknown feature') |
| if not isinstance(payload,dict):raise TypeError('Payload must be an object') |
| if not bool(self.enabled[FEATURES.index(feature)]):raise RuntimeError('Feature disabled: '+feature) |
| from .interfaces import validate_payload |
| validate_payload(feature,payload) |
| s=self.spec |
| def primitive(value): |
| if isinstance(value,torch.Tensor): |
| return value.item() if value.numel()==1 else value.cpu().tolist() |
| return value |
| def result(values):return {k:primitive(v) for k,v in values.items()} |
|
|
| if feature=='schema':return self.schema(payload) |
| if feature=='workflow': |
| n=self.integer(payload['nodes'],1,s['nodes']);deps=torch.zeros(n,n,device=self.device) |
| for before,after in payload.get('dependencies',[]): |
| deps[self.integer(after,0,n-1),self.integer(before,0,n-1)]=1 |
| complete=torch.zeros(n,device=self.device) |
| for node in payload.get('completed',[]):complete[self.integer(node,0,n-1)]=1 |
| out=self(feature,dependencies=deps,completed=complete) |
| return dict(available=out['available'].nonzero().flatten().cpu().tolist(),cyclic=bool(out['cyclic'])) |
| if feature=='index': |
| records=payload['documents'] |
| if not 0<=len(records)<=s['docs']:raise ValueError('Document capacity exceeded') |
| def tokens(text,limit,empty=False): |
| ids=list(text.encode('utf-8')) |
| if not int(not empty)<=len(ids)<=limit:raise ValueError('Index byte capacity exceeded') |
| return ids |
| docs=torch.zeros(len(records),s['doc_tokens'],device=self.device,dtype=torch.long) |
| mask=torch.zeros_like(docs,dtype=torch.float32) |
| keys=[];versions=[] |
| for i,doc in enumerate(records): |
| ids=tokens(doc['text'],s['doc_tokens'],empty=True);docs[i,:len(ids)]=self.tensor(ids,torch.long);mask[i,:len(ids)]=1 |
| keys.append(self.integer(doc.get('key',doc['id']),0,2**32-1)) |
| versions.append(self.integer(doc.get('version',1),1,2**31-1)) |
| query=tokens(payload['query'],s['query_tokens']) |
| out=self(feature,documents=docs,doc_mask=mask,query=self.tensor(query,torch.long), |
| versions=self.tensor(versions,torch.long),keys=self.tensor(keys,torch.long)) |
| found=bool(out['found']);index=int(out['index']) |
| return dict(found=found,document_id=records[index]['id'] if found else None, |
| byte_offset=int(out['positions'][index]) if found else None,scores=out['scores'].cpu().tolist()) |
| if feature=='syntax_scope': |
| out={} |
| if 'text' in payload: |
| parsed=self(feature,text=self.text(payload['text']),max_depth=s['stack']) |
| out.update(balanced=bool(parsed['valid']),remaining_depth=int(parsed['depth'])) |
| if 'events' in payload: |
| events=[] |
| for event in payload['events']: |
| op=self.integer(event[0],0,3);identifier=self.integer(event[1],0,s['symbols']-1) |
| events.append([op,identifier]) |
| if len(events)>s['text_bytes']:raise ValueError('Scope event capacity exceeded') |
| |
| self.calls[FEATURES.index(feature)].add_(1) |
| if not bool(self.enabled[FEATURES.index(feature)]):raise RuntimeError('Feature disabled') |
| checked=self.layers[feature].scope(self.tensor(events,torch.long).reshape(-1,2),s['scopes'],s['symbols']) |
| out['scope_valid']=bool(checked['valid']) |
| if not out:raise ValueError('Provide text or explicit scope events') |
| return out |
| if feature=='polynomial': |
| def coefficients(values): |
| if not 1<=len(values)<=s['degree']+1:raise ValueError('Polynomial degree capacity exceeded') |
| for v in values:self.integer(v,-s['coefficient_bound'],s['coefficient_bound']) |
| return self.tensor(values+[0]*(s['degree']+1-len(values)),torch.float64) |
| left=coefficients(payload['left']);right=coefficients(payload.get('right',[0])) |
| value=self(feature,left=left,right=right,operation=payload['operation']).round().long().cpu().tolist() |
| while len(value)>1 and value[-1]==0:value.pop() |
| return {'coefficients':value} |
| if feature=='constraints': |
| words=self.integer(payload.get('words',-1),-1,10000) |
| paragraphs=self.integer(payload.get('paragraphs',-1),-1,10000) |
| forbidden=self.text(payload.get('forbidden',''),64) |
| out=self(feature,text=self.text(payload['text']),words=words,paragraphs=paragraphs,forbidden=forbidden) |
| return dict(valid=bool(out['valid']),words=int(out['words']),paragraphs=int(out['paragraphs']),forbidden=bool(out['forbidden'])) |
| if feature=='units': |
| op=payload.get('operation','compatible') |
| if op not in ('compatible','convert','add','multiply','divide'):raise ValueError('Unsupported unit operation') |
| if payload['left'] not in UNIT_NAMES or payload['right'] not in UNIT_NAMES:raise ValueError('Unit outside registry') |
| out=self(feature,left=self.tensor(UNIT_NAMES.index(payload['left']),torch.long), |
| right=self.tensor(UNIT_NAMES.index(payload['right']),torch.long),operation=op) |
| if op in ('convert','add') and not bool(out['compatible']):raise ValueError('Incompatible unit dimensions for '+op) |
| return dict(compatible=bool(out['compatible']),dimensions=out['dimensions'].long().cpu().tolist(), |
| ratio_num=int(out['scale_num']),ratio_den=int(out['scale_den'])) |
| if feature=='sat': |
| n=self.integer(payload['variables'],1,s['sat_variables']);clauses=payload['clauses'] |
| if not 1<=len(clauses)<=s['sat_clauses']:raise ValueError('Clause capacity exceeded') |
| padded=[] |
| for clause in clauses: |
| if not 1<=len(clause)<=3:raise ValueError('Only clauses of one to three literals supported') |
| for literal in clause: |
| self.integer(literal,-n,n) |
| if literal==0:raise ValueError('Literal zero is invalid') |
| padded.append(clause+[clause[-1]]*(3-len(clause))) |
| out=self(feature,clauses=self.tensor(padded,torch.long),n_variables=n) |
| return dict(satisfiable=bool(out['satisfiable']),assignment=out['assignment'].long().cpu().tolist() if bool(out['satisfiable']) else None) |
| if feature=='closure': |
| n=self.integer(payload['nodes'],1,s['nodes']);adj=torch.zeros(n,n,device=self.device) |
| for a,b in payload.get('edges',[]):adj[self.integer(a,0,n-1),self.integer(b,0,n-1)]=1 |
| closure=self(feature,adjacency=adj) |
| if 'query' in payload: |
| a,b=payload['query'];return {'reachable':bool(closure[self.integer(a,0,n-1),self.integer(b,0,n-1)])} |
| return {'adjacency':closure.long().cpu().tolist()} |
| if feature=='binary': |
| op=payload['operation'];module=self.layers[feature] |
| if not bool(self.enabled[FEATURES.index(feature)]):raise RuntimeError('Feature disabled') |
| if op in ('hex_encode','base64_encode','base64_decode'): |
| if op=='base64_decode':data=self.text(payload['text'],s['bytes']*2) |
| else: |
| values=payload['bytes'] |
| if len(values)>s['bytes']:raise ValueError('Byte capacity exceeded') |
| data=self.tensor([self.integer(v,0,255) for v in values],torch.long) |
| self.calls[FEATURES.index(feature)].add_(1) |
| if op=='hex_encode':chars=module.hex_encode(data) |
| elif op=='base64_encode':chars=module.base64_encode(data) |
| else: |
| if data.numel()==0:return {'bytes':[]} |
| if data.numel()%4:raise ValueError('Base64 length must be a multiple of four') |
| decoded,valid,pad=module.base64_decode(data) |
| if not bool(valid):raise ValueError('Invalid or noncanonical Base64') |
| if decoded.numel()-int(pad)>s['bytes']:raise ValueError('Decoded byte capacity exceeded') |
| return {'bytes':decoded[:decoded.numel()-int(pad)].cpu().tolist()} |
| return {'text':bytes(chars.cpu().tolist()).decode('ascii')} |
| left=payload.get('left',[]);right=payload.get('right',[0]*len(left)) |
| if not 1<=len(left)<=s['bytes'] or len(left)!=len(right):raise ValueError('Invalid SIMD vector dimensions') |
| for value in left+right:self.integer(value,0,2**32-1) |
| rotation=self.integer(payload.get('rotation',0),0,31) |
| out=self(feature,left=self.tensor(left,torch.long),right=self.tensor(right,torch.long),operation=op,rotation=rotation) |
| return {'values':out.cpu().tolist()} |
| raise AssertionError(feature) |
|
|
| def schema(self,payload): |
| s=self.spec;schema=payload['schema'];value=payload['value'] |
| if not isinstance(schema,dict) or not isinstance(value,dict):raise ValueError('Flat object schemas only') |
| unsupported=set(schema)-{'type','properties','required','additionalProperties','description','title','$schema'} |
| if unsupported or schema.get('type','object')!='object':raise ValueError('Unsupported schema keyword/type') |
| properties=schema.get('properties',{}) |
| if not isinstance(properties,dict) or type(schema.get('additionalProperties',True)) is not bool: |
| raise ValueError('Only boolean additionalProperties supported') |
| if not isinstance(schema.get('required',[]),list) or any(not isinstance(k,str) for k in schema.get('required',[])): |
| raise ValueError('required must be a string list') |
| if len(properties)>s['fields'] or len(value)>s['fields']:raise ValueError('Field capacity exceeded') |
| F,K,V,E=s['fields'],s['key_bytes'],s['value_bytes'],s['enums'] |
| keys_s=torch.zeros(F,K,device=self.device);keys_a=torch.zeros_like(keys_s) |
| values_a=torch.zeros(F,V,device=self.device);enums=torch.zeros(F,E,V,device=self.device) |
| ps=torch.zeros(F,device=self.device);pa=ps.clone();required=ps.clone();ep=torch.zeros(F,E,device=self.device) |
| types=torch.zeros(F,7,device=self.device);actual=torch.zeros_like(types) |
| numeric=torch.zeros(F,device=self.device,dtype=torch.float64) |
| enum_numeric=torch.zeros(F,E,device=self.device,dtype=torch.float64) |
| enum_is_numeric=torch.zeros(F,E,device=self.device) |
| type_names=['null','boolean','integer','number','string','array','object'] |
| def encoded(data,limit): |
| if not isinstance(data,str):raise ValueError('Schema keys must be strings') |
| b=data.encode('utf-8') |
| if not b or len(b)>limit or b'\0' in b:raise ValueError('Schema string capacity exceeded or NUL') |
| return self.tensor(list(b)+[0]*(limit-len(b))) |
| def canonical(v):return json.dumps(v,ensure_ascii=False,sort_keys=True,separators=(',',':'),allow_nan=False) |
| def numeric_value(v): |
| if type(v) in (int,float): |
| if not -2**40<=v<=2**40:raise ValueError('Schema numeric domain exceeds +/-2^40') |
| return float(v),True |
| return 0.,False |
| for i,(name,prop) in enumerate(properties.items()): |
| if set(prop)-{'type','enum','description','title','default','examples'}:raise ValueError('Unsupported property keyword') |
| keys_s[i]=encoded(name,K);ps[i]=1;required[i]=name in schema.get('required',[]) |
| allowed=prop.get('type',type_names) |
| allowed=[allowed] if isinstance(allowed,str) else allowed |
| for t in allowed: |
| if t not in type_names:raise ValueError('Unsupported type') |
| types[i,type_names.index(t)]=1 |
| if t=='number':types[i,2]=1 |
| choices=prop.get('enum',[]) |
| if len(choices)>E or ('enum' in prop and not choices):raise ValueError('Invalid enum capacity') |
| for j,choice in enumerate(choices): |
| if isinstance(choice,(list,dict)):raise ValueError('Only scalar enum members supported') |
| enums[i,j]=encoded(canonical(choice),V);ep[i,j]=1 |
| number,is_number=numeric_value(choice);enum_numeric[i,j]=number;enum_is_numeric[i,j]=is_number |
| if any(name not in properties for name in schema.get('required',[])):raise ValueError('Required property absent in schema') |
| for i,(name,v) in enumerate(value.items()): |
| keys_a[i]=encoded(name,K);pa[i]=1 |
| wire=canonical(v) |
| |
| |
| if len(wire.encode('utf-8'))<=V:values_a[i]=encoded(wire,V) |
| typ={type(None):0,bool:1,int:2,float:3,str:4,list:5,dict:6}.get(type(v)) |
| if typ is None:raise ValueError('Non-JSON value') |
| actual[i,typ]=1 |
| number,is_number=numeric_value(v) |
| if is_number:numeric[i]=number |
| out=self('schema',key_schema=keys_s,key_actual=keys_a,present_schema=ps,present_actual=pa, |
| required=required,allowed_types=types,actual_types=actual,enum_values=enums,enum_present=ep, |
| actual_values=values_a,numeric=numeric,enum_numeric=enum_numeric,enum_is_numeric=enum_is_numeric, |
| additional=bool(schema.get('additionalProperties',True))) |
| return {k:(bool(v) if k=='valid' else int(v)) for k,v in out.items()} |
|
|
