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