File size: 2,883 Bytes
5cb93ca 8dea295 5cb93ca 8dea295 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 | """Reproduce the AMOS mini Core AI export from pinned source files."""
import argparse,json,sys
from pathlib import Path
import torch
from coreai_torch import TorchConverter,get_decomp_table
PINNED_DEPENDENCIES={'jhu-clsp/ettin-encoder-400m':'7662476d60abb071a5bd319c9f3074f3072c062d','google/siglip2-base-patch16-224':'75de2d55ec2d0b4efc50b3e9ad70dba96a7b2fa2'}
def pin_dependencies():
from huggingface_hub import hf_hub_download
import amos.packing,amos.vision
def pinned(repo,filename,**kwargs):
return hf_hub_download(repo,filename,revision=PINNED_DEPENDENCIES[repo],**kwargs)
amos.packing.hf_hub_download=pinned;amos.vision.hf_hub_download=pinned
class Text(torch.nn.Module):
def __init__(self,m):
super().__init__(); self.m=m
def forward(self,input_ids,position_ids,segment_ids,option_ids,visual_embeddings,visual_mask):
m=self.m
emb=torch.where(visual_mask,visual_embeddings,m.encoder.tok_embeddings(input_ids))
h=m.encoder(input_ids,position_ids,segment_ids,option_ids,inputs_embeds=emb)
return m.ptr_q(h),m.ptr_o(h),m.span_q(h),m.span_k(h),m.region_q(h),m.region_k(h),m.multi_bias
class Vision(torch.nn.Module):
def __init__(self,m):
super().__init__(); self.vision=m.vision; self.projector=m.projector
def forward(self,pixel_values):
v=self.vision
x=v.patch_embedding(pixel_values).flatten(2).transpose(1,2)+v.position_embedding
for b in v.blocks:x=b(x)
return self.projector(v.post_layernorm(x))
def save(m,args,path):
ep=torch.export.export(m,args).run_decompositions(get_decomp_table())
program=TorchConverter().add_exported_program(ep).to_coreai()
program.save_asset(path)
return list(ep.graph_signature.user_outputs)
def main():
p=argparse.ArgumentParser();p.add_argument('--source',default='upstream/amos-1-mini');p.add_argument('--output',default='bundles/amos-1-mini-CoreAI');p.add_argument('--length',type=int,default=512);a=p.parse_args()
source=Path(a.source).resolve();out=Path(a.output);out.mkdir(parents=True,exist_ok=True)
sys.path.insert(0,str(source/'serve'))
from amos.model import load_run
pin_dependencies()
torch.set_num_threads(4)
m,codec=load_run(source);m.eval()
args=(torch.full((1,a.length),codec.pad,dtype=torch.int32),torch.arange(a.length,dtype=torch.int32)[None],torch.zeros(1,a.length,dtype=torch.int32),torch.zeros(1,a.length,dtype=torch.int32),torch.zeros(1,a.length,m.encoder.config.hidden_size),torch.zeros(1,a.length,1,dtype=torch.bool))
print('Exporting text',flush=True)
outputs=save(Text(m).eval(),args,out/'text.aimodel')
print('Exporting vision',flush=True)
vo=save(Vision(m).eval(),(torch.zeros(1,3,224,224),),out/'vision.aimodel')
(out/'manifest.json').write_text(json.dumps({'sequence_length':a.length,'hidden_size':m.encoder.config.hidden_size,'outputs':outputs,'vision_output':vo[0],'precision':'float32'},indent=2))
print('Saved',flush=True)
if __name__=='__main__':main()
|