TorchScript format problem: ai.djl.engine.EngineException: forward() is missing value for argument 'attention_mask'

[JeeDevUser]

Scenario:

-downloaded all files under 'Files and versions' section, to the local drive, from location:
https://huggingface.co/google/pegasus-xsum/tree/main
-converted `pytorch_model.bin` to the TorchScript format (via Python script), 
and got the `pegasus_xsum_torchscript.pt` 

- initialized the tokenizer as follows:

HuggingFaceTokenizer tokenizer = HuggingFaceTokenizer.newInstance(Paths.get('/pathToTheLocation/with_tokenizer_json_etc'));

-loaded the model as follows:

      Model model = Model.newInstance(modelPath, "PyTorch");
      model.load(Paths.get(modelPath));

-Translator creation:

      // . Create Translator:
      Translator<long[], long[]> translator = new Translator<>() {
         @Override
         public NDList processInput(TranslatorContext ctx, long[] input) {
            // Kreiranje LongBuffer-a iz long[]
            LongBuffer buffer = LongBuffer.wrap(input);
            // Kreiranje NDArray objekta iz Buffer-a
            NDArray array = ctx.getNDManager().create(buffer, new Shape(1, input.length), DataType.INT64);

            NDList result =  new NDList(array);
            System.out.println("processInput() returns: " + result);
            return result;
         }


         @Override
         public long[] processOutput(TranslatorContext ctx, NDList list) {
            System.out.println("processOutput(), NDList in processOutput: " + list);
            // Getting NDArray from NDList:
            NDArray outputArray = list.singletonOrThrow();
            // Convert NDArray (output from the Model) to the long[]:
            long[] outputTokens = outputArray.toLongArray();
            System.out.println("processOutput() returns: " + list);
            return outputTokens;
         }
      };

• finally, generating a summary for the input text:

    // 5. Predictor:
     try (Predictor<long[], long[]> predictor = model.newPredictor(translator)) {
        // 6. Generating summary from the tokenized input:
        System.out.println("Tokenized input: " + Arrays.toString(tokenizedInput));
        long[] summaryTokens = predictor.predict(tokenizedInput);

        // 7. decode generated tokens:
        String summaryText = tokenizer.decode(summaryTokens);
        System.out.println("Summary: " + summaryText);
     }catch(Exception ex) {
        System.err.println("error : " + ex.getMessage());
        ex.printStackTrace();
     }

All this causes the following error:

 ai.djl.engine.EngineException: forward() is missing value for argument 'attention_mask'. Declaration: forward(__torch__.ModelWrapper self, Tensor input_ids, Tensor attention_mask, Tensor decoder_input_ids) -> Tensor
ai.djl.translate.TranslateException: ai.djl.engine.EngineException: forward() is missing value for argument 'attention_mask'. Declaration: forward(__torch__.ModelWrapper self, Tensor input_ids, Tensor attention_mask, Tensor decoder_input_ids) -> Tensor
	at ai.djl.inference.Predictor.batchPredict(Predictor.java:197)
	at ai.djl.inference.Predictor.predict(Predictor.java:133)
	at djl.TestDJL.main(TestDJL.java:144)
Caused by: ai.djl.engine.EngineException: forward() is missing value for argument 'attention_mask'. Declaration: forward(__torch__.ModelWrapper self, Tensor input_ids, Tensor attention_mask, Tensor decoder_input_ids) -> Tensor
	at ai.djl.pytorch.jni.PyTorchLibrary.moduleRunMethod(Native Method)
	at ai.djl.pytorch.jni.IValueUtils.forward(IValueUtils.java:57)
	at ai.djl.pytorch.engine.PtSymbolBlock.forwardInternal(PtSymbolBlock.java:146)
	at ai.djl.nn.AbstractBaseBlock.forward(AbstractBaseBlock.java:79)
	at ai.djl.nn.Block.forward(Block.java:127)
	at ai.djl.inference.Predictor.predictInternal(Predictor.java:147)
	at ai.djl.inference.Predictor.batchPredict(Predictor.java:188)

• What could be the problem?

[frankfliu]

@JeeDevUser

Can you provide your script that trace the model?

The model takes multiple inputs, but your processInput only provide in NDArray. your input should be long[][], it at least containds token_ids and attention_mask

[JeeDevUser]

@frankfliu ,
do you mean, python script I used to convert pythorch_model.bin to TorchScript format?
Here is:

import torch
from transformers import PegasusForConditionalGeneration, PegasusTokenizer
model_path = 'd:/IDEJE_ZA_AI_APP/Hugging Face Modeli/pegasus-xsum'
tokenizer_path = 'd:/IDEJE_ZA_AI_APP/Hugging Face Modeli/pegasus-xsum'

model = PegasusForConditionalGeneration.from_pretrained(model_path)
tokenizer = PegasusTokenizer.from_pretrained(tokenizer_path)

# Prepare inputs
input_text = "This is an example input text for summarization."
inputs = tokenizer(input_text, return_tensors="pt", max_length=1024, truncation=True)

# Extract necessary inputs
input_ids = inputs["input_ids"]
attention_mask = inputs["attention_mask"]

# Create dummy decoder inputs (typically starts with the decoder start token)
decoder_input_ids = torch.tensor([[model.config.decoder_start_token_id]])

# Wrap the model for TorchScript trace
class ModelWrapper(torch.nn.Module):
    def __init__(self, model):
        super(ModelWrapper, self).__init__()
        self.model = model

    def forward(self, input_ids, attention_mask, decoder_input_ids):
        outputs = self.model(
            input_ids=input_ids,
            attention_mask=attention_mask,
            decoder_input_ids=decoder_input_ids
        )
        return outputs.logits

wrapped_model = ModelWrapper(model)

# Trace the wrapped model
dummy_inputs = (input_ids, attention_mask, decoder_input_ids)
traced_model = torch.jit.trace(wrapped_model, dummy_inputs)

# Save the TorchScript model
traced_model.save("pegasus_xsum_torchscript.pt")
print("TorchScript model saved successfully!")

• here is complete Java code I am using for text sumarization:

   public static void main(String[] args) throws IOException, ModelException {
      String modelPath = "d:\\IDEJE_ZA_AI_APP\\Hugging Face Modeli\\pegasus-xsum\\pegasus_xsum_torchscript.pt";

      String inputText = "Sample text to be summarized";

      // 1. Tokenizer:
      String tokenizerPath = "d:\\IDEJE_ZA_AI_APP\\Hugging Face Modeli\\pegasus-xsum";
      HuggingFaceTokenizer tokenizer = HuggingFaceTokenizer.newInstance(Paths.get(tokenizerPath));
      Encoding encoding = tokenizer.encode(inputText);
      long[] tokenizedInput = encoding.getIds(); 
      long[] attentionMask = encoding.getAttentionMask(); 
      long[] decoderInputIds = new long[]{1}; // BOS token za Pegasus model ?

      Model model = Model.newInstance(modelPath, "PyTorch");
      model.load(Paths.get(modelPath));

      // 4. Translator
      Translator<long[][], long[]> translator = new Translator<>() {
         @Override
         public NDList processInput(TranslatorContext ctx, long[][] inputs) {
            NDManager manager = ctx.getNDManager();
            NDArray inputIdsArray = manager.create(inputs[0]);
            NDArray attentionMaskArray = manager.create(inputs[1]);
            NDArray decoderInputIdsArray = manager.create(inputs[2]);

            NDList result = new NDList(inputIdsArray, attentionMaskArray, decoderInputIdsArray);
            System.out.println("processInput() returns: " + result);
            return result;
         }

         @Override
         public long[] processOutput(TranslatorContext ctx, NDList list) {
            NDArray outputArray = list.singletonOrThrow();
            NDArray tokenIds = outputArray.argMax(1); 
            long[] outputTokens = tokenIds.toLongArray(); 

            return outputTokens;    
           }
      };



      // 5. Predictor
      try (Predictor<long[][], long[]> predictor = model.newPredictor(translator)) {

         long[][] input = new long[][]{tokenizedInput, attentionMask, decoderInputIds};

         // 6. summary generation:
         long[] summaryTokens = predictor.predict(input);

         String summaryText = tokenizer.decode(summaryTokens);
      } catch (Exception ex) {
         System.err.println("Error: " + ex.getMessage());
         ex.printStackTrace();
      }
   }

Tthis code _does not cause any error,_ but the summary text is not generated correctly. For example, I get the following output:

Summary:

Can you help somehow?