support generation on backpacks by overloading prepare_inputs_for_generation.

modeling_backpack_gpt2.py

@@ -153,7 +153,7 @@ class BackpackGPT2Model(BackpackGPT2PreTrainedModel):
     def get_sense_network(self):
         return self.sense_network
 
-    def forward(self, input_ids, position_ids):
+    def forward(self, input_ids, position_ids, **kwargs):
         # Compute senses
         sense_input_embeds = self.word_embeddings(input_ids)
         senses = self.sense_network(sense_input_embeds) # (bs, nv, s, d)
@@ -205,8 +205,48 @@ class BackpackGPT2LMHeadModel(BackpackGPT2PreTrainedModel):
 
   def get_lm_head(self):
       return self.lm_head
+  
+  def prepare_inputs_for_generation(self, input_ids, past_key_values=None, inputs_embeds=None, attention_mask=None, **kwargs):
+    # prepare_inputs_for_generation needs to be overwritten to support generation
+    # this is inspired from the one in GPT2LMHeadModel: https://github.com/huggingface/transformers/blob/d533465150532b0c5de167b574e59f64c68b1154/src/transformers/models/gpt2/modeling_gpt2.py#L1007C4-L1007C4
+    
+    token_type_ids = kwargs.get("token_type_ids", None)
+    # only last token for inputs_ids if past is defined in kwargs
+    if past_key_values:
+        input_ids = input_ids[:, -1].unsqueeze(-1)
+        if token_type_ids is not None:
+            token_type_ids = token_type_ids[:, -1].unsqueeze(-1)
+
+    attention_mask = kwargs.get("attention_mask", None)
+    position_ids = kwargs.get("position_ids", None)
+
+    if attention_mask is not None and position_ids is None:
+        # create position_ids on the fly for batch generation
+        position_ids = attention_mask.long().cumsum(-1) - 1
+        position_ids.masked_fill_(attention_mask == 0, 1)
+        if past_key_values:
+            position_ids = position_ids[:, -1].unsqueeze(-1)
+    else:
+        position_ids = None
+
+    # if `inputs_embeds` are passed, we only want to use them in the 1st generation step
+    if inputs_embeds is not None and past_key_values is None:
+        model_inputs = {"inputs_embeds": inputs_embeds}
+    else:
+        model_inputs = {"input_ids": input_ids}
+
+    model_inputs.update(
+        {
+            "past_key_values": past_key_values,
+            "use_cache": kwargs.get("use_cache"),
+            "position_ids": position_ids,
+            "attention_mask": attention_mask,
+            "token_type_ids": token_type_ids,
+        }
+    )
+    return model_inputs
 
-  def forward(self, input_ids, position_ids=None):
+  def forward(self, input_ids, position_ids=None, **kwargs):
       outputs = self.backpack(input_ids, position_ids=position_ids)
       hidden_states, contextualization = outputs.hidden_states, outputs.contextualization
       lm_logits = self.lm_head(hidden_states) # (bs, s, V)