File size: 4,467 Bytes
5cb93ca | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 | """Core AI neural execution; upstream packing/decoding uses PyTorch tensors only."""
import json,math,sys
from pathlib import Path
import numpy as np
import torch
from coreai.runtime import AIModel,NDArray,SpecializationOptions
class AmosMini:
@classmethod
async def load(cls,bundle):
self=cls();self.root=Path(bundle).resolve();sys.path.insert(0,str(self.root/'serve'))
from amos.packing import Codec
cfg=json.loads((self.root/'config.json').read_text())
# Codec's loader accepts a locally pinned tokenizer through this scoped patch.
import amos.packing as packing
old=packing.hf_hub_download
packing.hf_hub_download=lambda repo,filename: str(self.root/filename)
try:self.codec=Codec(cfg['repo'],max_state=cfg['max_state'],digits=cfg.get('digits',False),score_rank=cfg.get('score_rank',False))
finally:packing.hf_hub_download=old
self.manifest=json.loads((self.root/'manifest.json').read_text())
self.text=await AIModel.load(self.root/'text.aimodel',SpecializationOptions.default());self.tf=self.text.load_function('main')
self.vision=await AIModel.load(self.root/'vision.aimodel',SpecializationOptions.default());self.vf=self.vision.load_function('main')
return self
async def raw(self,state,questions,images=()):
from amos.packing import pack,collate
from amos.vision import preprocess
feats=[]
async def feature(im):
pix=preprocess(im)[None].numpy()
out=await self.vf({'pixel_values':NDArray(pix)})
return torch.from_numpy(out[self.manifest['vision_output']].numpy()[0].copy())
# Pack expects pre-projection 196-cell features. Instead use placeholders;
# inject the Core AI projected 49-cell vectors directly into embedding slots.
for im in images:feats.append(await feature(im))
qs=[]
for q in questions:
q=dict(q)
if 'options' in q:
opts=[]
for o in q['options']:
if isinstance(o,str):opts.append(o)
else:opts.append(torch.zeros(196,768));feats.append(await feature(o))
q['options']=opts
qs.append(q)
p=pack(self.codec,state,qs,[torch.zeros(196,768) for _ in images],strict=True)
n=len(p);L=self.manifest['sequence_length'];d=self.manifest['hidden_size']
if n>L:raise ValueError(f'Request has {n} tokens; Core AI capacity is {L}.')
b=collate([p],self.codec.pad)
args={}
for src,dst,pad in [('ids','input_ids',self.codec.pad),('pos','position_ids',0),('seg','segment_ids',-1),('opt','option_ids',0)]:
arr=np.full((1,L),pad,dtype=np.int32);arr[0,:n]=getattr(p,src);args[dst]=NDArray(arr)
ve=np.zeros((1,L,d),np.float32);vm=np.zeros((1,L,1),bool)
for at,f in zip(p.image_at,feats):ve[0,at:at+49]=f.numpy();vm[0,at:at+49]=True
args['visual_embeddings']=NDArray(ve);args['visual_mask']=NDArray(vm)
result=await self.tf(args)
pq,po,sq,sk,rq,rk,bias=[torch.from_numpy(result[k].numpy().copy()) for k in self.manifest['outputs']]
Q=len(p.decide);O=len(b['o_q']);out={'opt':torch.full((Q,max(1,max(p.n_options or [0]))), -torch.inf)}
for i in range(O):
valid=b['option_of'][0]==i
pooled=po[0,:valid.shape[0]][valid].mean(0)
qi=int(b['o_q'][i]);rank=int(b['o_rank'][i]);at=p.decide[qi]
out['opt'][qi,rank]=(pq[0,at]*pooled).sum()/math.sqrt(d)+(bias if p.types[qi]=='multi' else 0)
spans=[i for i,t in enumerate(p.types) if t=='span'];regions=[i for i,t in enumerate(p.types) if t=='region']
out['span_q']=torch.tensor(spans);out['region_q']=torch.tensor(regions)
if spans:
q=sq[0,p.decide].view(Q,2,d)[spans];k=sk[0,:b['ids'].shape[1]].view(-1,2,d)
out['span']=torch.einsum('qcd,lcd->qcl',q,k)/math.sqrt(d)
out['span'].masked_fill_(~b['state_mask'][:,None,:],-torch.inf)
if regions:
out['region']=rq[0,[p.decide[i] for i in regions]]@rk[0,:b['ids'].shape[1]].T/math.sqrt(d)
out['region'].masked_fill_(~b['region_mask'],-torch.inf)
return out,[p]
async def predict(self,state,questions,images=()):
from amos.model import decode
out,ps=await self.raw(state,questions,images)
return decode(out,ps)[0]
async def request(self,req):
from amos.serve import parse,answer
from PIL import Image
import base64,io
state,ids,qs=parse(req)
def image(s):return Image.open(io.BytesIO(base64.b64decode(s.split(',')[-1])))
ims=[image(s) for s in req.get('images',[])]
qs=[{**q,'options':[image(o['image']) if isinstance(o,dict) else o for o in q['options']]} if 'options' in q else q for q in qs]
results=await self.predict(state,qs,ims)
return {'answers':{i:answer(q,r) for i,q,r in zip(ids,qs,results)}}
|