Merge pull request #56 from SamSamhuns/fix_for_new_transformers_lib_ver

Change model_kwargs argument to encoder_outputs to support transformers>=4.22.1

Author: gwkrsrch

donut/model.py

@@ -206,7 +206,7 @@ def add_special_tokens(self, list_of_tokens: List[str]):
         if newly_added_num > 0:
             self.model.resize_token_embeddings(len(self.tokenizer))
 
-    def prepare_inputs_for_inference(self, input_ids: torch.Tensor, past=None, use_cache: bool = None, **model_kwargs):
+    def prepare_inputs_for_inference(self, input_ids: torch.Tensor, encoder_outputs: torch.Tensor, past=None, use_cache: bool = None, attention_mask: torch.Tensor = None):
         """
         Args:
             input_ids: (batch_size, sequence_lenth)
@@ -223,7 +223,7 @@ def prepare_inputs_for_inference(self, input_ids: torch.Tensor, past=None, use_c
             "attention_mask": attention_mask,
             "past_key_values": past,
             "use_cache": use_cache,
-            "encoder_hidden_states": model_kwargs["encoder_outputs"].last_hidden_state,
+            "encoder_hidden_states": encoder_outputs.last_hidden_state,
         }
         return output