scripts/conversion/export.py

"""Reproducible full-graph Crivo export. Conversion-only Python; app runtime stays TypeScript.
Run with the pinned environment in requirements.txt. Source snapshots must match sources.json.
"""
import argparse, gc, importlib.util, json, shutil, sys
from pathlib import Path
import numpy as np
import onnx
import onnxruntime as ort
from quantize import quantize_weights
import torch
from transformers import AutoTokenizer, ModernBertConfig, ModernBertModel
from safetensors.torch import load_file

parser = argparse.ArgumentParser()
parser.add_argument('size', choices=['17m','32m','68m','150m','400m','1b'])
args = parser.parse_args()
root = Path(__file__).resolve().parents[2]
source = root / '.cache/conversion' / args.size
out = root / '.cache/variants' / args.size
out.mkdir(parents=True, exist_ok=True)
(out/'onnx').mkdir(exist_ok=True)
meta = json.loads((source/'source.json').read_text())
assert meta == json.loads(Path(__file__).with_name('sources.json').read_text())[args.size], 'Source revision mismatch'
torch.set_num_threads(4)
spec = importlib.util.spec_from_file_location('upstream_model', source/'model.py')
upstream = importlib.util.module_from_spec(spec)
spec.loader.exec_module(upstream)
original = json.loads((source/'config.json').read_text())
config = ModernBertConfig(**{k:v for k,v in original.items() if k not in ('architectures','decision')})
config._attn_implementation = 'eager'
config.reference_compile = False
model = upstream.DecisionModel(ModernBertModel(config), config.hidden_size).eval()
model.load_state_dict(load_file(str(source/'model.safetensors')), strict=True)
tok = AutoTokenizer.from_pretrained(source)
tok.save_pretrained(out)
tc = json.loads((out/'tokenizer_config.json').read_text())
tc['tokenizer_class'] = 'PreTrainedTokenizer'
(out/'tokenizer_config.json').write_text(json.dumps(tc,indent=2)+'\n')
exported = {**original, 'model_type':'custom','architectures':['PreTrainedModel'], 'source':meta}
# Browser release initially retains the tested 7,999-token cap, not upstream's 32k claim.
exported['web_max_length'] = 7999
names = ['input_ids','attention_mask','opt_positions','opt_mask','qtype']
encoder = upstream.Encoder(tok)
cases = [
 ('Please refund my duplicate payment.', [('choice','Identify the intent.',{'refund':'Return money','cancel':'Cancel order','other':'Another request'}),('noul','The customer was charged twice.',None),('score','Rate urgency.',['Low','Medium','High'])]),
 ('The order is fine, I only want to know the delivery date.', [('choice','Identify intent.',{'question':'Ask a question','refund':'Return money'} )]),
 ('A long support request. '*30,[('score','Rate urgency.',['Low','High']),('noul','The service is unavailable.',{'true':'Unavailable','false':'Available'})])
]
def inputs(state, questions):
 rows=[]
 for kind,instructions,criteria in questions:
  options,_=upstream.option_texts(kind,criteria,instructions)
  ids,pos=encoder.encode(instructions,state,options)
  rows.append((ids,pos,upstream.QUESTION_TYPES[kind]))
 return upstream.collate(rows,tok.pad_token_id)
samples=[inputs(*case) for case in cases]
references=[]
with torch.inference_mode():
 for feeds in samples: references.append(model(**feeds).numpy())
graph=out/'onnx/model.onnx'
print('EXPORT',args.size,flush=True)
with torch.no_grad():
 torch.onnx.export(model,tuple(samples[0][name] for name in names),str(graph),input_names=names,output_names=['logits'],opset_version=17,dynamo=False,external_data=True,dynamic_axes={**{n:{0:'batch',1:'sequence'} for n in names[:2]},**{n:{0:'batch',1:'options'} for n in names[2:4]},'qtype':{0:'batch'},'logits':{0:'batch',1:'options'}})
del model;gc.collect()
# Keep each external buffer below browser ArrayBuffer limits.
def shard_external(path):
 g=onnx.load(str(path))
 index=0;offset=0;handle=None
 try:
  for tensor in g.graph.initializer:
   if not tensor.HasField('raw_data'): continue
   data=tensor.raw_data
   if handle is None or offset+len(data)>1024**3:
    if handle: handle.close();index+=1
    location=path.name+'_data'+('' if index==0 else '_'+str(index))
    handle=(path.parent/location).open('wb');offset=0
   onnx.external_data_helper.set_external_data(tensor,location,offset,len(data))
   handle.write(data);offset+=len(data)
   tensor.ClearField('raw_data');tensor.data_location=onnx.TensorProto.EXTERNAL
 finally:
  if handle: handle.close()
 onnx.save_model(g,str(path))
 del g;gc.collect()
 return index+1
if args.size=='1b':
 exported['transformers.js_config']={'use_external_data_format':{'model.onnx':shard_external(graph)}}
(out/'config.json').write_text(json.dumps(exported,indent=2)+'\n')
onnx.checker.check_model(str(graph))
opts=ort.SessionOptions();opts.intra_op_num_threads=4
session=ort.InferenceSession(str(graph),sess_options=opts,providers=['CPUExecutionProvider'])
fp_delta=0
for feeds,expected in zip(samples,references):
 actual=session.run(['logits'],{n:feeds[n].numpy() for n in names})[0]
 np.testing.assert_allclose(actual,expected,rtol=1e-3,atol=1e-4)
 fp_delta=max(fp_delta,float(np.max(np.abs(actual-expected))))
del session;gc.collect()
print('QUANTIZE',args.size,flush=True)
quant=out/'onnx/model_quantized.onnx'
quantize_weights(graph, quant)
if args.size=='1b':
 exported['transformers.js_config']['use_external_data_format']['model_quantized.onnx']=shard_external(quant)
 (out/'config.json').write_text(json.dumps(exported,indent=2)+'\n')
onnx.checker.check_model(str(quant))
session=ort.InferenceSession(str(quant),sess_options=opts,providers=['CPUExecutionProvider'])
quant_delta=0
fixtures=[]
for (state,questions),feeds,expected in zip(cases,samples,references):
 actual=session.run(['logits'],{n:feeds[n].numpy() for n in names})[0]
 assert np.isfinite(actual).all()
 p=torch.from_numpy(expected).double().softmax(-1).numpy()
 q=torch.from_numpy(actual).double().softmax(-1).numpy()
 quant_delta=max(quant_delta,float(np.max(np.abs(p-q))))
 fixtures.append({'state':state,'questions':[{'type':k,'instructions':i,**({'criteria':c} if c is not None else {})} for k,i,c in questions],'fp32':p.tolist(),'int8':q.tolist()})
report={'quantization':{'storage':'int8','arithmetic':'fp32','scheme':'symmetric-per-channel-weight-only'},'source':meta,'fp32_max_logit_error':fp_delta,'int8_max_probability_error':quant_delta,'cases':fixtures,'versions':{'torch':torch.__version__,'onnx':onnx.__version__,'onnxruntime':ort.__version__}}
# Retain only the final graphs and referenced external buffers.
keep = {'model.onnx', 'model_quantized.onnx'}
for path in [graph, quant]:
 g=onnx.load(str(path),load_external_data=False)
 for tensor in g.graph.initializer:
  keep.update(item.value for item in tensor.external_data if item.key=='location')
for path in (out/'onnx').iterdir():
 if path.is_file() and path.name not in keep: path.unlink()
(out/'parity.json').write_text(json.dumps(report,indent=2)+'\n')
print('PASS',args.size,fp_delta,quant_delta,flush=True)