File size: 3,194 Bytes
0690c04 | 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 | """Text-only D3 Core AI export, adapted from Max Farrell's Apache-2.0 recipe.
Changes by hugging-mac: no vision graph or media inputs; explicit source identity.
"""
import argparse,json,hashlib
from pathlib import Path
import torch
from coreai_torch import TorchConverter,get_decomp_table
from safetensors.torch import load_file
from transformers import Qwen3_5Model
class DecisionText(torch.nn.Module):
def __init__(self,backbone,readout):
super().__init__();self.language_model=backbone.language_model
self.register_buffer('readout',readout.float())
def forward(self,input_ids,attention_mask,position_ids):
hidden=self.language_model(input_ids=input_ids,attention_mask={
'full_attention':attention_mask[:,None,None,:].bool(),
'linear_attention':attention_mask,
},position_ids=position_ids,use_cache=False).last_hidden_state
return hidden[:,-1].float()@self.readout.T
def main():
p=argparse.ArgumentParser();p.add_argument('--source',required=True);p.add_argument('--output',required=True);p.add_argument('--length',type=int,default=512);p.add_argument('--revision',required=True);a=p.parse_args()
if a.length<1:raise ValueError('length must be positive')
torch.set_num_threads(4);source=Path(a.source);out=Path(a.output);out.mkdir(exist_ok=True,parents=True)
config=json.loads((source/'decision_config.json').read_text())
assert config['attention_mode']=='noncausal_full_attention'
manifest=json.loads((source/'MODEL_MANIFEST.json').read_text())
for name,expected in manifest['identity']['files_sha256'].items():
with (source/name).open('rb') as f:actual=hashlib.file_digest(f,'sha256').hexdigest()
if actual!=expected:raise ValueError('Source checksum mismatch: '+name)
print('Source identity verified; loading BF16 backbone',flush=True)
backbone=Qwen3_5Model.from_pretrained(source,dtype=torch.bfloat16,attn_implementation='sdpa').eval()
model=DecisionText(backbone,load_file(str(source/'readout.safetensors'))['weight']).eval()
sample=(torch.ones(1,a.length,dtype=torch.int32),torch.ones(1,a.length,dtype=torch.int32),torch.arange(a.length,dtype=torch.int32).view(1,1,-1).expand(4,1,-1))
print('Exporting text-only graph',flush=True)
with torch.no_grad():
exported=torch.export.export(model,sample,strict=False)
print('Decomposing',flush=True);exported=exported.run_decompositions(get_decomp_table())
print('Converting',flush=True)
program=TorchConverter().add_exported_program(exported).to_coreai()
program.save_asset(out/'decision.aimodel')
(out/'manifest.json').write_text(json.dumps(dict(source='vllm-sr/d3-edge',revision=a.revision,sequence_length=a.length,precision='bfloat16',readout_precision='float32',attention_mode=config['attention_mode'],modality='text',inputs=['input_ids','attention_mask','position_ids'],output=exported.graph_signature.user_outputs[0],function='main',format='Apple Core AI .aimodel',source_weight_sha256=manifest['identity']['files_sha256']['model-00001-of-00001.safetensors']),indent=2)+'\n')
print('Saved text-only decision.aimodel',flush=True)
if __name__=='__main__':main()
|