Using ONNX format of the facebook/mbart-large-50-many-to-many-mmt in DJL?

[JeeDevUser]

Hi all,

Used this code to export mentioned model 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                      
)

Got resulting mbart_large_50.onnx but I’m completely stuck on how to use it in DJL to translate between two languages. I just don't know where to start, or how to organize the code.
Any suggestions would be appreciated.

[frankfliu]

@JeeDevUser

facebook/mbart-large-50-many-to-many-mmt is an encoder/decoder model. It contains 2 model inside it. To improve the performance, it use kv_cache at ecoding stage. When you export the model to onnx, you should export 3 model. You can use optimum to export the model:

optimum-cli export onnx --model facebook/mbart-large-50-many-to-many-mmt output_model

You will get the following files:

• decoder_model.onnx
• decoder_with_past_model.onnx
• encoder_model.onnx

You can also use the following code to trace the model:

from typing import Tuple, List

import torch
from torch import nn
from transformers import AutoTokenizer, MBartForConditionalGeneration


class ModelWrapper(nn.Module):

    def __init__(self, model) -> None:
        super().__init__()
        self.model = model

    def encode(
            self,
            input_ids: torch.Tensor,
            attention_mask: torch.Tensor,
    ) -> Tuple[torch.Tensor]:
        return self.model.get_encoder()(
            input_ids=input_ids,
            attention_mask=attention_mask,
            output_attentions=False,
            output_hidden_states=False,
            return_dict=False,
        )

    def forward(
            self,
            attention_mask: torch.Tensor,
            decoder_input_ids: torch.Tensor,
            encoder_outputs: torch.Tensor,
            past_key_values: List[torch.Tensor],
    ) -> Tuple[torch.Tensor]:
        past_kv_list = []
        for i in range(12):
            layers = []
            for j in range(4):
                layers.append(past_key_values[i * 4 + j])
            past_kv_list.append(layers)

        return self.model(
            decoder_input_ids=decoder_input_ids,
            encoder_outputs=(encoder_outputs,),
            attention_mask=attention_mask,
            past_key_values=tuple(past_kv_list),
            use_cache=True,
            output_attentions=False,
            output_hidden_states=False,
            return_dict=False,
        )

    def forward_init(
            self,
            attention_mask: torch.Tensor,
            decoder_input_ids: torch.Tensor,
            encoder_outputs: torch.Tensor,
    ) -> Tuple[torch.Tensor]:
        return self.model(
            decoder_input_ids=decoder_input_ids,
            encoder_outputs=(encoder_outputs,),
            attention_mask=attention_mask,
            use_cache=True,
            output_attentions=False,
            output_hidden_states=False,
            return_dict=False,
        )


def generate_dummy_past_key_values(num_heads=16, num_layers=12, kv_dims=64, batch_size=1):
    past_key_values = []
    for _ in range(num_layers):
        past_key_values.append(torch.zeros(batch_size, num_heads, 1, kv_dims))
        past_key_values.append(torch.zeros(batch_size, num_heads, 1, kv_dims))
        past_key_values.append(torch.zeros(batch_size, num_heads, 21, kv_dims))
        past_key_values.append(torch.zeros(batch_size, num_heads, 21, kv_dims))

    return past_key_values


def jit_trace():
    model_id = "facebook/mbart-large-50-many-to-many-mmt"
    mbart = MBartForConditionalGeneration.from_pretrained(model_id)
    tokenizer = AutoTokenizer.from_pretrained(model_id,
                                              src_lang="fr_XX",
                                              tgt_lang="en_XX")
    # save tokenizer
    tokenizer.save_pretrained("mbart-large")

    model = ModelWrapper(mbart)

    input_text = "Le chef de l 'ONU affirme qu 'il n 'y a pas de solution militai"

    encoding = tokenizer(input_text, return_tensors="pt")
    input_ids = encoding["input_ids"]
    attention_mask = encoding["attention_mask"]

    past_key_values = generate_dummy_past_key_values()
    encoder_outputs = model.encode(input_ids, attention_mask)
    # print(decoder_outputs)

    traced_decoder = torch.jit.trace_module(
        model,
        {
            "encode": [input_ids, attention_mask],
            "forward_init": [attention_mask, torch.tensor([[0]]), encoder_outputs[0]],
            "forward": [attention_mask, torch.tensor([[0, 0]]), encoder_outputs[0], past_key_values],
        }
    )
    torch.jit.save(traced_decoder, "mbart-large/model.pt")


if __name__ == '__main__':
    jit_trace()

In order to execute the model in DJL, you have to implement the same MBartForConditionalGeneration.generate() logic in java. Here is a naive implementation that mimic python logic (everything is hard coded for en_XX -> fr_XX)

import ai.djl.ModelException;
import ai.djl.huggingface.tokenizers.Encoding;
import ai.djl.huggingface.tokenizers.HuggingFaceTokenizer;
import ai.djl.inference.Predictor;
import ai.djl.ndarray.NDArray;
import ai.djl.ndarray.NDList;
import ai.djl.ndarray.NDManager;
import ai.djl.repository.zoo.Criteria;
import ai.djl.repository.zoo.ZooModel;
import ai.djl.translate.TranslateException;

import java.io.IOException;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.util.ArrayList;
import java.util.List;

public class Text2TextGeneration {

    public static void main(String[] args) throws ModelException, IOException, TranslateException {
        String inputText = "Le chef de l 'ONU affirme qu 'il n 'y a pas de solution militai";

        Path path = Paths.get("mbart-large/model.pt");

        Path dir = Paths.get("mbart-large");
        HuggingFaceTokenizer tokenizer = HuggingFaceTokenizer.newInstance(dir);
        Encoding encoding = tokenizer.encode(inputText, true, false);
        List<Long> outputIds = new ArrayList<>();

        Criteria<NDList, NDList> criteria =
                Criteria.builder()
                        .setTypes(NDList.class, NDList.class)
                        .optModelPath(path)
                        .optEngine("PyTorch")
                        .build();
        try (ZooModel<NDList, NDList> model = criteria.loadModel();
             Predictor<NDList, NDList> predictor = model.newPredictor();
             NDManager manager = NDManager.newBaseManager("PyTorch")) {

            long[] ids = encoding.getIds();
            ids[0] = 250008; // fr_XX
            NDArray inputIds = manager.create(ids).expandDims(0);
            NDArray attentionMask = manager.create(encoding.getAttentionMask()).expandDims(0);

            NDArray encodeMethod = manager.create("");
            encodeMethod.setName("module_method:encode");
            NDArray initMethod = manager.create("");
            initMethod.setName("module_method:forward_init");

            NDList encodeInput = new NDList(inputIds, attentionMask, encodeMethod);
            NDList encoderOutputs = predictor.predict(encodeInput);

            NDArray decoderOutputIds = manager.create(new long[][]{{2}});

            NDList initInputs =
                    new NDList(attentionMask, decoderOutputIds, encoderOutputs.get(0), initMethod);
            NDList pastKeyValues = null;
            for (int i = 0; i < 200; ++i) {
                NDList output;
                if (i == 0) {
                    output = predictor.predict(initInputs);
                } else {
                    NDList decoderInputs =
                            new NDList(attentionMask, decoderOutputIds, encoderOutputs.get(0));
                    decoderInputs.addAll(pastKeyValues);
                    output = predictor.predict(decoderInputs);
                }
                NDArray logits = output.get(0);
                pastKeyValues = output.subNDList(1, 49);
                for (NDArray pastKeyValue : pastKeyValues) {
                    pastKeyValue.setName("past_key_values[]");
                }

                logits = logits.get(":, -1, :");
                NDArray nextTokenIds = logits.argMax(-1);
                long tokenId = nextTokenIds.getLong();
                if (tokenId == 2) {
                    // found eos
                    break;
                }

                if (tokenId != 250004) {
                    // en_XX
                    outputIds.add(tokenId);
                }
                decoderOutputIds = decoderOutputIds.concat(nextTokenIds.expandDims(0), -1);
            }
        }

        String text = tokenizer.decode(outputIds.stream().mapToLong(Long::longValue).toArray());
        System.out.println(text);
    }
}

[ImeshaDilshani]

now i have converted my mbart50 trained model to .onnx but when I'm going to test the model it gives input sentences. why ?

import torch
from transformers import MBartForConditionalGeneration, MBart50TokenizerFast

model_path = "./2. translation_model_ta_si"
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
)

---

import onnxruntime as ort
from transformers import MBart50TokenizerFast
import numpy as np
import torch

Load tokenizer

model_path = "./2. translation_model_ta_si" # your model path
tokenizer = MBart50TokenizerFast.from_pretrained(model_path)

Tamil → Sinhala

tokenizer.src_lang = "ta_IN"

Input Tamil sentence

input_text = "நீங்கள் எப்படி இருக்கிறீர்கள்?"

Tokenize input

inputs = tokenizer(input_text, return_tensors="pt")
input_ids = inputs["input_ids"].numpy()

Load ONNX model

onnx_path = "./mbart_large_50.onnx"
ort_session = ort.InferenceSession(onnx_path)

Run inference

ort_inputs = {"input_ids": input_ids}
ort_outs = ort_session.run(["logits"], ort_inputs)

Convert output logits to token ids

logits = ort_outs[0]
predicted_ids = np.argmax(logits, axis=-1)

Decode to get translated text

translated = tokenizer.batch_decode(predicted_ids, skip_special_tokens=True)
print("Translated Output:", translated[0])

---

Translated Output: நீங்கள் எப்படி இருக்கிறீர்கள்?

[frankfliu]

@ImeshaDilshani
As I mentioned, mbart50 is encoder/decoder model, export decoder only won't work. Can you try to use optimum-cli export to export the model to ONNX?

[ImeshaDilshani]

after using this optimum-cli export if generated this files but there is not decoder_with_past_model.onnx file . insted of this there is decoder_model.onnx_data . why?

! optimum-cli export onnx --model "2. translation_model_ta_si" --task translation "my_onnx_model"

Output

d:\CS-2019-063\ONNX\onnx-1\lib\site-packages\transformers\models\mbart\modeling_mbart.py:873: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
if sequence_length != 1:
Could not find ONNX initializer for torch parameter model.encoder.embed_tokens.weight. model.encoder.embed_tokens.weight will not be checked for deduplication.
Found different candidate ONNX initializers (likely duplicate) for the tied weights:
lm_head.weight: {'onnx::MatMul_3921'}
model.decoder.embed_tokens.weight: {'model.shared.weight'}
model.encoder.embed_tokens.weight: set() --> ignored (may be a parameter from a part of the model not exported)
model.shared.weight: {'model.shared.weight'}
-[x] values not close enough, max diff: 4.5299530029296875e-05 (atol: 1e-05)
The ONNX export succeeded with the warning: The maximum absolute difference between the output of the reference model and the ONNX exported model is not within the set tolerance 1e-05:

• logits: max diff = 4.5299530029296875e-05.
  The exported model was saved at: my_onnx_model

then after i went to load the model but it show this type of error

from transformers import MBart50TokenizerFast
from optimum.onnxruntime import ORTModelForSeq2SeqLM

# Path to your ONNX model directory
onnx_model_path = "my_onnx_model"

# Load tokenizer and model
tokenizer = MBart50TokenizerFast.from_pretrained(onnx_model_path)
model = ORTModelForSeq2SeqLM.from_pretrained(onnx_model_path)

# Set source and target languages (adjust as per your training)
tokenizer.src_lang = "ta_IN"
tokenizer.tgt_lang = "si_LK"

output

d:\CS-2019-063\ONNX\onnx-1\lib\site-packages\tqdm\auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html
from .autonotebook import tqdm as notebook_tqdm

---

IndexError Traceback (most recent call last)
Cell In[2], line 9
7 # Load tokenizer and model
8 tokenizer = MBart50TokenizerFast.from_pretrained(onnx_model_path)
----> 9 model = ORTModelForSeq2SeqLM.from_pretrained(onnx_model_path)
11 # Set source and target languages (adjust as per your training)
12 tokenizer.src_lang = "ta_IN"

File d:\CS-2019-063\ONNX\onnx-1\lib\site-packages\optimum\onnxruntime\modeling_ort.py:550, in ORTModel.from_pretrained(cls, model_id, config, export, subfolder, revision, force_download, local_files_only, trust_remote_code, cache_dir, token, provider, providers, provider_options, session_options, use_io_binding, **kwargs)
545 if file_name is not None:
546 logger.warning(
547 f"file_name was set to {file_name} but will be ignored as the model will be converted to ONNX"
548 )
--> 550 return super().from_pretrained(
551 model_id,
552 config=config,
553 export=_export,
554 force_download=force_download,
555 token=token,
556 cache_dir=cache_dir,
557 subfolder=subfolder,
558 local_files_only=local_files_only,
559 trust_remote_code=trust_remote_code,
560 revision=revision,
...
-> 1108 decoder_with_past_path = decoder_with_past_path[0] if decoder_with_past_path else model_files[0]
1109 decoder_path = decoder_with_past_path.parent / decoder_with_past_path.name.replace("_with_past", "")
1110 else:

IndexError: list index out of range