File size: 3,607 Bytes
ebef316 f7f4846 ebef316 f7f4846 ebef316 f7f4846 ebef316 | 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 | """Export AMOS Qwen3.5 bidirectional decisions and trained quote/box heads."""
import argparse,json,math,sys
from pathlib import Path
import torch
from coreai_torch import TorchConverter,get_decomp_table
from qwen_export_helpers import DecisionVision
class Text(torch.nn.Module):
def __init__(self,model,letters):
super().__init__();self.language_model=model.model.language_model;self.noncausal=bool(getattr(model.config,'amos_noncausal',False))
self.register_buffer('readout',model.get_output_embeddings().weight[letters].detach().float().clone())
def forward(self,input_ids,attention_mask,position_ids,visual_embeddings,visual_mask):
emb=self.language_model.embed_tokens(input_ids)
emb=torch.where(visual_mask,visual_embeddings.to(emb.dtype),emb)
mask=attention_mask[:,None,None,:].bool()
if not self.noncausal:
width=input_ids.shape[1];idx=torch.arange(width,device=input_ids.device);mask=mask&(idx[:,None]>=idx[None,:])[None,None]
h=self.language_model(inputs_embeds=emb,attention_mask={'full_attention':mask,'linear_attention':None},position_ids=position_ids,use_cache=False).last_hidden_state
return h.float(),h.float()@self.readout.T
class Heads(torch.nn.Module):
def __init__(self,heads):super().__init__();self.heads=heads
def forward(self,question_hidden,state_hidden,image_hidden):
h=self.heads;d=question_hidden.shape[0];dtype=h.span_q.weight.dtype
q=question_hidden.to(dtype);s=state_hidden.to(dtype);im=image_hidden.to(dtype)
span=torch.einsum('cd,lcd->cl',h.span_q(q).view(2,d),h.span_k(s).view(-1,2,d))/math.sqrt(d)
region=h.region_k(im)@h.region_q(q)/math.sqrt(d)
box=h.box(q,im,region,7,7) if h.box is not None else torch.zeros(4)
return span.float(),region.float(),box.float()
def save(m,args,path):
with torch.no_grad():ep=torch.export.export(m,args,strict=False).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('name');p.add_argument('--length',type=int,default=512);p.add_argument('--precision',default='bfloat16');a=p.parse_args()
source=Path('upstream')/a.name;out=Path('bundles')/(a.name+'-CoreAI');out.mkdir(parents=True,exist_ok=True)
sys.path.insert(0,str((source/'serve').resolve()))
from amos.decoder import load_model,load_heads
torch.set_num_threads(4);model,codec=load_model(str(source/'model'),'cpu',dtype=getattr(torch,a.precision));model.eval();heads=load_heads(source/'model',model,codec)
d=model.config.text_config.hidden_size;L=a.length
inputs=(torch.full((1,L),codec.pad,dtype=torch.int32),torch.ones(1,L,dtype=torch.int32),torch.arange(L,dtype=torch.int32).view(1,1,-1).expand(4,1,-1),torch.zeros(1,L,d),torch.zeros(1,L,1,dtype=torch.bool))
print('Exporting text',flush=True);outputs=save(Text(model,codec.letters).eval(),inputs,out/'text.aimodel')
print('Exporting vision',flush=True);visual=model.model.visual;vo=save(DecisionVision(visual).eval(),(torch.zeros(196,1536),),out/'vision.aimodel')
print('Exporting heads',flush=True);ho=save(Heads(heads).eval(),(torch.zeros(d),torch.zeros(L,d),torch.zeros(49,d)),out/'heads.aimodel')
(out/'manifest.json').write_text(json.dumps({'sequence_length':L,'hidden_size':d,'outputs':outputs,'vision_output':vo[0],'head_outputs':ho,'precision':a.precision,'echo':codec.echo,'kinds':heads.kinds,'keep':heads.keep,'has_box':heads.box is not None,'image_side':224,'noncausal':bool(getattr(model.config,'amos_noncausal',False))},indent=2))
print('Saved',flush=True)
if __name__=='__main__':main()
|