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