How can I execute decoder of ONNX Export from Seq2Seq model

I made an export of the Helsinki model using python optimum and i am trying to run the model with only the onnx environment and implement beam search from scratch because I have to later port this to a system not running python. So I want to prototype in Python a version without optimum/pytorch.

However I don't find a way to execute/feed the decoder with the outputs of the encoder, due to rank issues.

#Export the model
from transformers import AutoTokenizer
from optimum.onnxruntime import ORTModelForSeq2SeqLM
from optimum.pipelines import pipeline

tokenizer = AutoTokenizer.from_pretrained("Helsinki-NLP/opus-mt-de-en")

model = ORTModelForSeq2SeqLM.from_pretrained("Helsinki-NLP/opus-mt-de-en", from_transformers=True)

onnx_translation = pipeline("translation_de_to_en", model=model, tokenizer=tokenizer)
onnx_translation.save_pretrained("DE_TO_EN_TRANSLATION_HELSINKI")

code to run the encoder works:

import numpy as np
from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("Helsinki-NLP/opus-mt-de-en")
inputs = tokenizer("Das ist ein Test", return_tensors="pt")

encoder_file = "DE_TO_EN_TRANSLATION_HELSINKI/encoder_model.onnx"
decoder_file = "DE_TO_EN_TRANSLATION_HELSINKI/decoder_model.onnx"

sess = rt.InferenceSession(encoder_file)
input_ids = sess.get_inputs()[0].name
attention_mask =  sess.get_inputs()[1].name
ort_inputs = {input_ids: inputs['input_ids'].numpy() ,attention_mask: inputs['attention_mask'].numpy()}
output_encoder = sess.run([label_name], ort_inputs)
print(output_encoder)

However if I then try the decoder:

sess2 = rt.InferenceSession(decoder_file)
input_name = sess.get_inputs()[0].name
input_name2 = sess.get_inputs()[1].name
pred_onx2 = sess.run([label_name], {input_name:inputs['input_ids'].numpy() ,input_name2: output_encoder})
print(output_encoder)

Output:

InvalidArgument: \[ONNXRuntimeError\] : 2 : INVALID_ARGUMENT : Invalid rank for input: attention_mask Got: 4 Expected: 2 Please fix either the inputs or the model.

I don't understand where this goes wrong?

The inputs should be:

input: input_ids tensor(int64) \['batch_size', 'decoder_sequence_length'\]

input: encoder_hidden_states tensor(float) \['batch_size', 'encoder_sequence_length', 512\]

according to

for t in sess.get_inputs():
    print("input:", t.name, t.type, t.shape)

for t in sess.get_outputs():
    print("output:", t.name, t.type, t.shape)

So which one is the attention mask?

I also tried:

sess2 = rt.InferenceSession(decoder_file)
input_name = sess.get_inputs()[0].name
input_name2 = sess.get_inputs()[1].name
pred_onx2 = sess.run([label_name], {input_name:inputs['attention_mask'].numpy() ,input_name2: output_encoder})
print(output_encoder)

Also a side question. If I understand it right, after the first time executing the decoder, I use the Decoder_with_past_model file? Or how is the relation?

appreciate any help