mDenseOn exported to ONNX: float32 (exact) and an 8-bit build, with the scripts that made them

dc2a503 verified about 16 hours ago

2.52 kB

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

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