| """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() |
|
|
