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);
    }
}