Add token_type_ids to prepare_inputs_for_generation for gpt/gpt2

src/transformers/modeling_gpt2.py

@@ -680,11 +680,15 @@ def get_output_embeddings(self):
 
     def prepare_inputs_for_generation(self, input_ids, past=None, **kwargs):
         # only last token for inputs_ids if past is defined in kwargs
+        token_type_ids = kwargs.get("token_type_ids")
         if past:
             input_ids = input_ids[:, -1].unsqueeze(-1)
+            if token_type_ids is not None:
+                token_type_ids = token_type_ids[:, -1].unsqueeze(-1)
 
         return {
             "input_ids": input_ids,
+            "token_type_ids": token_type_ids,
             "past_key_values": past,
             "use_cache": kwargs.get("use_cache"),
         }

src/transformers/modeling_openai.py

@@ -541,6 +541,21 @@ def __init__(self, config):
     def get_output_embeddings(self):
         return self.lm_head
 
+    def prepare_inputs_for_generation(self, input_ids, past=None, **kwargs):
+        # only last token for inputs_ids if past is defined in kwargs
+        token_type_ids = kwargs.get("token_type_ids")
+        if past:
+            input_ids = input_ids[:, -1].unsqueeze(-1)
+            if token_type_ids is not None:
+                token_type_ids = token_type_ids[:, -1].unsqueeze(-1)
+
+        return {
+            "input_ids": input_ids,
+            "token_type_ids": token_type_ids,
+            "past_key_values": past,
+            "use_cache": kwargs.get("use_cache"),
+        }
+
     @add_start_docstrings_to_callable(OPENAI_GPT_INPUTS_DOCSTRING)
     @add_code_sample_docstrings(
         tokenizer_class=_TOKENIZER_FOR_DOC,