A script for inference using T5 in Java

Hello, I am having a little trouble for building a script to perform the inference in Java using the T5 model. I am using T5 for the summarization use case.

1️⃣ Model

Luckily the t5-small already has ONNX models exported on the HF: t5-small-onnx, so I just have to use the encoder_model.onnx and decoder_model.onnx files.

2️⃣ The code I am using currently

public class t5_onnx {
    public static void main(String args[]) throws OrtException {
        // Load the model and create InferenceSession
        System.out.println("This is the model loading");
        String encoder = "encoder_model.onnx";
        String decoder = "decoder_model.onnx";
        OrtEnvironment env = OrtEnvironment.getEnvironment();
        OrtSession encoderSession = env.createSession(encoder);
        OrtSession decoderSession = env.createSession(decoder);

        String prompt = "yo!";
        String generatedText = generate(prompt, env, encoderSession, decoderSession);
    }

    static String generate(String prompt, OrtEnvironment env, OrtSession encoderSession, OrtSession decoderSession) throws OrtException {
        // Get the input and output names for the encoder and decoder
        String encoderInputName = encoderSession.getInputNames().iterator().next();
        String encoderOutputName = encoderSession.getOutputNames().iterator().next();
        String decoderInputName = decoderSession.getInputNames().iterator().next();
        String decoderOutputName = decoderSession.getOutputNames().iterator().next();

        // Encoding
        // INPUT IDS
        long[] inputData = new long[10];
        for (int i = 0; i < inputData.length; i++) {
            inputData[i] = (long) Math.random() * 10;
        }
        long[] input_ids_shape = new long[]{1, inputData.length};  // Shape of the input data

        // ATTENTION MASK FOR INPUT IDS
        long[] attention_mask = new long[inputData.length];
        Arrays.fill(attention_mask, 1);
        long[] attention_mask_shape = new long[]{1, inputData.length};  // Shape of the input data

        // OnnxTensor - INPUT IDS
        OnnxTensor inputTensor = OnnxTensor.createTensor(env, LongBuffer.wrap(inputData), input_ids_shape);
        OnnxTensor attentionTensor = OnnxTensor.createTensor(env, LongBuffer.wrap(attention_mask), attention_mask_shape);
        Map<String, OnnxTensor> encoder_inputs = Map.of("input_ids", inputTensor, "attention_mask", attentionTensor);

        OrtSession.Result encoderOutput = encoderSession.run(encoder_inputs);
        OnnxTensor encoderOutputTensor = (OnnxTensor) encoderOutput.get(encoderOutputName).get();

//        OrtSession.Result decoderOutput = decoderSession.run(inputMap);
//        OnnxTensor decoderOutputTensor = (OnnxTensor) decoderOutput.get(decoderOutputName).get();

        return "STATIC RETURN";
    }

As you can see:

• I can successfully forward pass in the encoder, but I am not sure how to create the OnnxTensor for the Decoder as input
• How to sample?

🙏🏻 A small request

Can you please share a correct way to go for this?

Because:

• Currently, I am picking random tokens to forward: How can I tokenize the inputs?
• How can I perform a successful forward pass and get the summary for my input?

Can I get a code guidance to do this? Because in CausalLM models, we have to go for a loop until the max tokens are reached, I am not sure what should I do for this seq to seq model and how the "doSample" and "top K" can be applied here.

Please share a script to do this... I would appreciate your help.
Thanks!

Can you please help @tianleiwu 🙏🏻

Start from the Java snippet in the issue and the referenced encoder_model.onnx and decoder_model.onnx files. Determine how the T5 summarization flow should tokenize input, pass encoder output to the decoder, and generate tokens for a seq2seq model. Done means a documented Java example produces a summary and explains sampling options such as doSample and top K.