Upload modeling.py with huggingface_hub

modeling.py

@@ -5,10 +5,15 @@ from typing import Optional, Union
 import numpy as np
 import torch
 from torch import nn
+import transformers
 from transformers import (
     AutoModel,
     LlavaNextForConditionalGeneration,
 )
+
+_V5 = int(transformers.__version__.split(".")[0]) >= 5
+if _V5:
+    from transformers.masking_utils import create_causal_mask
 from transformers.cache_utils import Cache, DynamicCache
 from transformers.modeling_flash_attention_utils import FlashAttentionKwargs
 from transformers.models.granitemoehybrid.modeling_granitemoehybrid import (
@@ -81,10 +86,10 @@ class Granite4VisionForConditionalGeneration(LlavaNextForConditionalGeneration):
         use_cache: Optional[bool] = None,
         output_attentions: Optional[bool] = None,
         output_hidden_states: Optional[bool] = None,
-        cache_position: Optional[torch.LongTensor] = None,
         logits_to_keep: Union[int, torch.Tensor] = 0,
         **kwargs: Unpack[TransformersKwargs],
     ) -> Union[tuple, LlavaNextCausalLMOutputWithPast]:
+        cache_position = kwargs.pop("cache_position", None)
 
         output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
         output_hidden_states = (
@@ -99,8 +104,7 @@ class Granite4VisionForConditionalGeneration(LlavaNextForConditionalGeneration):
             else self.config.vision_feature_select_strategy
         )
 
-        outputs = self.model(
-            input_ids,
+        model_kwargs = dict(
             pixel_values=pixel_values,
             image_sizes=image_sizes,
             vision_feature_layer=vision_feature_layer,
@@ -113,9 +117,10 @@ class Granite4VisionForConditionalGeneration(LlavaNextForConditionalGeneration):
             output_attentions=output_attentions,
             output_hidden_states=output_hidden_states,
             return_dict=True,
-            cache_position=cache_position,
-            **kwargs,
         )
+        if not _V5:
+            model_kwargs["cache_position"] = cache_position
+        outputs = self.model(input_ids, **model_kwargs, **kwargs)
 
         hidden_states = outputs.last_hidden_state
 
@@ -154,17 +159,29 @@ class Granite4VisionForConditionalGeneration(LlavaNextForConditionalGeneration):
         logits_to_keep=None,
         **kwargs,
     ):
-        model_inputs = super().prepare_inputs_for_generation(
-            input_ids,
-            past_key_values=past_key_values,
-            inputs_embeds=inputs_embeds,
-            attention_mask=attention_mask,
-            cache_position=cache_position,
-            logits_to_keep=logits_to_keep,
-            **kwargs,
-        )
+        if _V5:
+            is_first = kwargs.get("is_first_iteration", False)
+            model_inputs = super().prepare_inputs_for_generation(
+                input_ids,
+                past_key_values=past_key_values,
+                inputs_embeds=inputs_embeds,
+                attention_mask=attention_mask,
+                logits_to_keep=logits_to_keep,
+                **kwargs,
+            )
+        else:
+            is_first = cache_position[0] == 0 if cache_position is not None else True
+            model_inputs = super().prepare_inputs_for_generation(
+                input_ids,
+                past_key_values=past_key_values,
+                inputs_embeds=inputs_embeds,
+                attention_mask=attention_mask,
+                cache_position=cache_position,
+                logits_to_keep=logits_to_keep,
+                **kwargs,
+            )
         model_inputs = self._init_hybrid_cache(**model_inputs)
-        if cache_position[0] == 0:
+        if is_first:
             model_inputs["pixel_values"] = pixel_values
             model_inputs["image_sizes"] = image_sizes
 
@@ -182,9 +199,9 @@ class Granite4VisionForConditionalGeneration(LlavaNextForConditionalGeneration):
         **kwargs,
     ):
         """Handle HybridMambaAttentionDynamicCache for GraniteMoeHybrid language model."""
-        empty_past_kv = past_key_values is None or (isinstance(past_key_values, DynamicCache) and past_key_values[0][0] is None)
+        empty_past_kv = past_key_values is None or (isinstance(past_key_values, DynamicCache) and past_key_values.get_seq_length() == 0)
 
-        if not empty_past_kv:
+        if not empty_past_kv and not _V5:
             if (
                 inputs_embeds is not None
                 or cache_position[-1] >= input_ids.shape[1]
@@ -192,7 +209,7 @@ class Granite4VisionForConditionalGeneration(LlavaNextForConditionalGeneration):
                 input_ids = input_ids[:, -cache_position.shape[0] :]
             elif input_ids.shape[1] != cache_position.shape[0]:
                 input_ids = input_ids[:, cache_position]
-        elif use_cache:
+        elif use_cache and empty_past_kv:
             past_key_values = HybridMambaAttentionDynamicCache(
                 self.model.language_model.config, input_ids.shape[0], self.dtype, device=self.device
             )
@@ -214,9 +231,10 @@ class Granite4VisionForConditionalGeneration(LlavaNextForConditionalGeneration):
                 "past_key_values": past_key_values,
                 "use_cache": use_cache,
                 "attention_mask": attention_mask,
-                "cache_position": cache_position,
             }
         )
+        if not _V5:
+            model_inputs["cache_position"] = cache_position
 
         for key, value in kwargs.items():
             if key not in model_inputs:
@@ -258,7 +276,7 @@ class Granite4VisionModel(LlavaNextPreTrainedModel):
 
         self.vocab_size = config.text_config.vocab_size
         self.language_model = AutoModel.from_config(config.text_config)
-        self.pad_token_id = self.config.pad_token_id if self.config.pad_token_id is not None else -1
+        self.pad_token_id = getattr(self.config, "pad_token_id", None) or -1
         self.post_init()
 
     def get_input_embeddings(self):
@@ -473,14 +491,14 @@ class Granite4VisionModel(LlavaNextPreTrainedModel):
         output_attentions: Optional[bool] = None,
         output_hidden_states: Optional[bool] = None,
         return_dict: Optional[bool] = None,
-        cache_position: Optional[torch.LongTensor] = None,
         **kwargs: Unpack[FlashAttentionKwargs],
     ) -> Union[tuple, LlavaNextModelOutputWithPast]:
+        cache_position = kwargs.pop("cache_position", None)
         output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
         output_hidden_states = (
             output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
         )
-        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        return_dict = return_dict if return_dict is not None else self.config.return_dict
         vision_feature_layer = (
             vision_feature_layer if vision_feature_layer is not None else self.config.vision_feature_layer
         )
@@ -522,19 +540,31 @@ class Granite4VisionModel(LlavaNextPreTrainedModel):
         # Custom forward pass with vision injection at specific LLM layers
         hidden_states = inputs_embeds * self.language_model.embedding_multiplier
 
-        if cache_position is None:
-            past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
-            cache_position = torch.arange(
-                past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device
+        if _V5:
+            if position_ids is None:
+                past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
+                position_ids = torch.arange(
+                    past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device
+                ).unsqueeze(0)
+            causal_mask = create_causal_mask(
+                config=self.language_model.config,
+                inputs_embeds=inputs_embeds,
+                attention_mask=attention_mask,
+                past_key_values=past_key_values,
             )
-
-        if position_ids is None:
-            position_ids = cache_position.unsqueeze(0)
-
-        causal_mask = self.language_model._update_causal_mask(
-            attention_mask, inputs_embeds, cache_position, past_key_values, output_attentions
-        )
-        mamba_mask = self.language_model._update_mamba_mask(attention_mask, cache_position)
+            mamba_mask = self.language_model._update_mamba_mask(attention_mask, past_key_values)
+        else:
+            if cache_position is None:
+                past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
+                cache_position = torch.arange(
+                    past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device
+                )
+            if position_ids is None:
+                position_ids = cache_position.unsqueeze(0)
+            causal_mask = self.language_model._update_causal_mask(
+                attention_mask, inputs_embeds, cache_position, past_key_values, output_attentions
+            )
+            mamba_mask = self.language_model._update_mamba_mask(attention_mask, cache_position)
 
         position_embeddings = None
         if self.language_model.rotary_emb is not None:
@@ -558,21 +588,24 @@ class Granite4VisionModel(LlavaNextPreTrainedModel):
             if output_hidden_states:
                 all_hidden_states += (hidden_states,)
 
-            layer_outputs = decoder_layer(
-                hidden_states,
+            layer_kwargs = dict(
                 attention_mask=layer_mask,
                 past_key_values=past_key_values,
-                output_attentions=output_attentions,
                 use_cache=use_cache,
-                cache_position=cache_position,
                 position_embeddings=position_embeddings,
-                **kwargs,
             )
-
-            hidden_states = layer_outputs[0]
-
-            if output_attentions and layer_outputs[1] is not None:
-                all_self_attns += (layer_outputs[1],)
+            if not _V5:
+                layer_kwargs["output_attentions"] = output_attentions
+                layer_kwargs["cache_position"] = cache_position
+            layer_outputs = decoder_layer(hidden_states, **layer_kwargs, **kwargs)
+
+            # v5 decoder layers return a bare tensor; v4 returns a tuple
+            if isinstance(layer_outputs, torch.Tensor):
+                hidden_states = layer_outputs
+            else:
+                hidden_states = layer_outputs[0]
+                if output_attentions and layer_outputs[1] is not None:
+                    all_self_attns += (layer_outputs[1],)
 
         hidden_states = self.language_model.norm(hidden_states)