Fix create_causal_mask compatibility for split branches

config.json

@@ -35,5 +35,6 @@
   "bundled_segmentation_model_name": "bundled_backbones/segmenter_encoder",
   "bundled_text_model_name": "bundled_backbones/text_decoder",
   "bundled_tokenizer_name": ".",
-  "segmenter_weights_in_model_state": true
+  "segmenter_weights_in_model_state": true,
+  "visual_projection_type": "mlp4"
 }

gpt2_modified.py

@@ -1,4 +1,5 @@
 from typing import Optional, Union
+import inspect
 
 import torch
 import torch.nn.functional as F
@@ -11,6 +12,8 @@ from transformers.modeling_outputs import BaseModelOutputWithPastAndCrossAttenti
 from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
 from transformers.models.gpt2.modeling_gpt2 import GPT2Attention, GPT2Block, eager_attention_forward
 
+_CREATE_CAUSAL_MASK_EMBEDS_ARG = "inputs_embeds" if "inputs_embeds" in inspect.signature(create_causal_mask).parameters else "input_embeds"
+
 
 class GPT2AttentionModified(GPT2Attention):
     def forward(
@@ -169,14 +172,15 @@ class GPT2ModelModified(GPT2Model):
         if attention_mask is not None and attention_mask.ndim < 4:
             attention_mask = attention_mask.view(batch_size, -1)
 
-        causal_mask = create_causal_mask(
-            config=self.config_causal,
-            inputs_embeds=inputs_embeds,
-            attention_mask=attention_mask,
-            cache_position=cache_position,
-            past_key_values=past_key_values,
-            position_ids=position_ids,
-        )
+        causal_mask_kwargs = {
+            "config": self.config_causal,
+            _CREATE_CAUSAL_MASK_EMBEDS_ARG: inputs_embeds,
+            "attention_mask": attention_mask,
+            "cache_position": cache_position,
+            "past_key_values": past_key_values,
+            "position_ids": position_ids,
+        }
+        causal_mask = create_causal_mask(**causal_mask_kwargs)
 
         _use_sdpa = self._attn_implementation == "sdpa" and output_attentions is False and head_mask is None
         if self.config.add_cross_attention and encoder_hidden_states is not None: