Hi all,
This is the code I used to successfully export to ONNX format
import torch
from transformers import MBartForConditionalGeneration, MBart50TokenizerFast
model_path = "./mbart_large_50_model"
model = MBartForConditionalGeneration.from_pretrained(model_path)
tokenizer = MBart50TokenizerFast.from_pretrained(model_path)
input_text = "This is just simple text"
inputs = tokenizer(input_text, return_tensors="pt")
onnx_path = "./mbart_large_50.onnx"
# Export to ONNX format:
torch.onnx.export(
model, # PyTorch model
(inputs["input_ids"],),
onnx_path, # the path for the resulting ONNX file
input_names=["input_ids"], # input tensor names
output_names=["logits"], # output tensor names
dynamic_axes={"input_ids": {0: "batch", 1: "sequence"}},
opset_version=14
)
So I got the mbart_large_50.onnx
and bunch of other files
Also, I successfully verified resulting ONNX file, with following code:
import onnx
try:
onnx.checker.check_model("d:/Install/TensorFlow/models/MBART_Base/ONNX/mbart_large_50.onnx")
except onnx.checker.ValidationError as e:
print(f"The model is invalid: {e}")
else:
print("The model is valid!")
So far, so good…
-Now, this is how I am using this ONNX format,for the job of translating between two languages:
from transformers import MBart50TokenizerFast
import numpy as np
tokenizer_path = "d:/Install/TensorFlow/models/MBART_Base/mbart_large_50_model/"
# specify source and target language:
tokenizer = MBart50TokenizerFast.from_pretrained(tokenizer_path, src_lang="en_XX", tgt_lang="hr_HR")
text = "This is just simple text."
inputs = tokenizer(text, return_tensors="np")
input_ids = np.array(inputs["input_ids"], dtype=np.int64)
import onnxruntime as ort
onnx_session = ort.InferenceSession("d:/Install/TensorFlow/models/MBART_Base/ONNX/mbart_large_50.onnx")
onnx_inputs = {"input_ids": input_ids}
onnx_outputs = onnx_session.run(["logits"], onnx_inputs)
logits = onnx_outputs[0]
probabilities = np.exp(logits) / np.sum(np.exp(logits), axis=-1, keepdims=True)
generated_tokens = np.argmax(logits, axis=-1)
translated_text = tokenizer.decode(generated_tokens[0], skip_special_tokens=True)
print("Translated tekst:", translated_text)
But, for the input text: "This is just simple text."
, I am getting following output:
"is just simple text"
as you can see, the text is not translated, it is the same as the input, except the first word is missing. In fact, the same result as I am getting with TorchScript version of model
Wondering, what am I doing wrong?