File size: 2,521 Bytes
dc2a503 | 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 | # Exports lightonai/mDenseOn to ONNX in float32, with a dynamic batch and length, and checks it against PyTorch.
# pip install torch transformers onnx onnxscript onnxruntime (made with PyTorch 2.15, Transformers 5.22.2)
# python export.py onnx/model.onnx
import sys
import numpy as np
import onnxruntime as ort
import torch
from transformers import AutoModel, AutoTokenizer
MODEL = 'lightonai/mDenseOn'
out = sys.argv[1] if len(sys.argv) > 1 else 'onnx/model.onnx'
tok = AutoTokenizer.from_pretrained(MODEL)
# Eager attention: the sliding-window masks become part of the graph.
model = AutoModel.from_pretrained(MODEL, dtype=torch.float32, attn_implementation='eager').eval()
class Encoder(torch.nn.Module):
def __init__(self, m):
super().__init__()
self.m = m
def forward(self, input_ids, attention_mask):
return self.m(input_ids=input_ids, attention_mask=attention_mask).last_hidden_state
example = tok(['query: hello world', 'document: ' + 'a longer text that pads the first one ' * 40], padding=True, truncation=True, max_length=512, return_tensors='pt')
batch, seq = torch.export.Dim('batch', min=1, max=1024), torch.export.Dim('seq', min=2, max=8192)
with torch.inference_mode():
prog = torch.onnx.export(Encoder(model), (example['input_ids'], example['attention_mask']), dynamo=True, input_names=['input_ids', 'attention_mask'], output_names=['last_hidden_state'],
dynamic_shapes={'input_ids': {0: batch, 1: seq}, 'attention_mask': {0: batch, 1: seq}}, opset_version=20, external_data=False, optimize=True)
prog.save(out, external_data=False)
# The same CLS vectors as PyTorch, at lengths below and above the 128-token window, alone and in batches.
texts = ['query: how do I ship to the live server', 'document: ' + ' '.join(['The quick brown fox jumps over the lazy dog.'] * 60), 'query: Πώς φτιάχνω μια παρουσίαση;', 'document: ok']
sess = ort.InferenceSession(out, providers=['CPUExecutionProvider'])
for batch_texts in [texts, texts[:1], texts[1:3]]:
e = tok(batch_texts, padding=True, truncation=True, max_length=512, return_tensors='pt')
with torch.inference_mode():
ref = torch.nn.functional.normalize(model(**e).last_hidden_state[:, 0], dim=-1).numpy()
h = sess.run(None, {'input_ids': e['input_ids'].numpy(), 'attention_mask': e['attention_mask'].numpy()})[0][:, 0]
h /= np.linalg.norm(h, axis=1, keepdims=True)
print('cosine to PyTorch:', np.round((h * ref).sum(1), 6).tolist())
|