Update transformers to 4.51.3 (#18)

- Update transformers to 4.51.3 (43e59dc346665e6de9e4edfd1f3a4807c629804e)
- Raise error if vision tower receives kwargs (443f857075c53b4af6ed4c00af8036c9cc9a11ce)

README.md

@@ -84,10 +84,10 @@ pillow
 protobuf
 sentencepiece
 torch
-transformers>=4.48.0,<4.49
+transformers>=4.48.0,<4.52
 ```
 
-Note: MAIRA-2 has last been tested with transformers v4.48.0.
+Note: MAIRA-2 has last been tested with transformers v4.51.3.
 
 First, initialise the model and put it in eval mode.
 ```python

config.json

@@ -8,7 +8,6 @@
     "AutoModelForVision2Seq": "modeling_maira2.Maira2ForConditionalGeneration"
   },
   "hidden_size": 4096,
-  "ignore_index": -100,
   "image_seq_length": 576,
   "image_token_index": 32204,
   "model_type": "maira2",
@@ -21,19 +20,33 @@
     "architectures": [
       "LlamaForCausalLM"
     ],
+    "attention_bias": false,
+    "attention_dropout": 0.0,
+    "head_dim": 128,
+    "hidden_act": "silu",
+    "hidden_size": 4096,
+    "initializer_range": 0.02,
+    "intermediate_size": 11008,
     "max_position_embeddings": 4096,
+    "mlp_bias": false,
     "model_type": "llama",
+    "num_attention_heads": 32,
+    "num_hidden_layers": 32,
+    "num_key_value_heads": 32,
     "pad_token_id": 0,
+    "pretraining_tp": 1,
     "rms_norm_eps": 1e-05,
     "rope_scaling": {
       "factor": 1.5,
       "rope_type": "linear"
     },
+    "rope_theta": 10000.0,
     "torch_dtype": "bfloat16",
+    "use_cache": true,
     "vocab_size": 32207
   },
   "torch_dtype": "float32",
-  "transformers_version": "4.48.0",
+  "transformers_version": "4.51.3",
   "vision_config": {
     "apply_layernorm": true,
     "architectures": [
@@ -45,11 +58,13 @@
     "hidden_dropout_prob": 0.0,
     "hidden_size": 768,
     "image_size": 518,
+    "initializer_range": 0.02,
     "layer_norm_eps": 1e-06,
     "layerscale_value": 1.0,
     "mlp_ratio": 4,
     "model_type": "dinov2",
     "num_attention_heads": 12,
+    "num_channels": 3,
     "num_hidden_layers": 12,
     "out_features": [
       "stage12"
@@ -57,6 +72,7 @@
     "out_indices": [
       12
     ],
+    "patch_size": 14,
     "qkv_bias": true,
     "reshape_hidden_states": false,
     "stage_names": [
@@ -75,6 +91,7 @@
       "stage12"
     ],
     "torch_dtype": "float32",
+    "use_mask_token": true,
     "use_swiglu_ffn": false
   },
   "vision_feature_layer": -1,

generation_config.json

@@ -5,5 +5,5 @@
   "max_length": 4096,
   "max_new_tokens": 450,
   "pad_token_id": 0,
-  "transformers_version": "4.48.0"
+  "transformers_version": "4.51.3"
 }

modeling_maira2.py

@@ -2,9 +2,16 @@
 #  Licensed under the MSRLA License. See LICENSE in the repo root for license information.
 
 
+from typing import Any
+
 import torch
 from torch.nn import Linear, Module, Sequential
-from transformers import AutoBackbone, AutoModelForCausalLM, LlavaForConditionalGeneration, LlavaPreTrainedModel
+from transformers import (
+    AutoBackbone,
+    AutoModelForCausalLM,
+    LlavaForConditionalGeneration,
+    LlavaPreTrainedModel,
+)
 from transformers.activations import ACT2FN
 from transformers.utils import check_min_version
 
@@ -48,7 +55,6 @@ class Maira2ForConditionalGeneration(LlavaForConditionalGeneration):
     config_class = Maira2Config
 
     def __init__(self, config: Maira2Config) -> None:
-
         # Check transformers version is at least 4.46.0.dev0  otherwise the model fails
         # silently since get_image_features is not called in the forward pass
         check_min_version("4.46.0.dev0")
@@ -62,11 +68,17 @@ class Maira2ForConditionalGeneration(LlavaForConditionalGeneration):
             config.text_config,
             attn_implementation=config._attn_implementation,
         )
-        self.pad_token_id = self.config.pad_token_id if self.config.pad_token_id is not None else -1
+        self.pad_token_id = (
+            self.config.pad_token_id if self.config.pad_token_id is not None else -1
+        )
         self.post_init()
 
     def get_image_features(
-        self, pixel_values: torch.FloatTensor, vision_feature_layer: int, vision_feature_select_strategy: str
+        self,
+        pixel_values: torch.FloatTensor,
+        vision_feature_layer: int | list[int],
+        vision_feature_select_strategy: str,
+        **kwargs: Any,
     ) -> torch.Tensor:
         """
         This method extracts the image features from the vision backbone using the specified feature layer and
@@ -74,15 +86,27 @@ class Maira2ForConditionalGeneration(LlavaForConditionalGeneration):
         class instead of the `hidden_states` which are used in the default implementation of `get_image_features` in LlavaForConditionalGeneration.
         The feature_maps returned by Dinov2Backbone are the hideen_states with a layernorm applied to them.
         """
+        if isinstance(vision_feature_layer, list):
+            raise ValueError(
+                "MAIRA-2 does not support list values for vision_feature_layer."
+            )
+
+        if vision_feature_select_strategy not in ["default", "full"]:
+            raise ValueError(
+                f"Unexpected select feature strategy: {self.config.vision_feature_select_strategy}"
+            )
+
+        extra_kwargs = {k: v for k, v in kwargs.items() if v is not None}
+        if extra_kwargs:
+            raise ValueError(
+                f"MAIRA-2 does not support passing extra kwargs to the vision tower, received: {extra_kwargs}"
+            )
         image_outputs = self.vision_tower(pixel_values, output_hidden_states=True)
+
         selected_image_feature = image_outputs.feature_maps[vision_feature_layer]
 
         if vision_feature_select_strategy == "default":
             selected_image_feature = selected_image_feature[:, 1:]
-        elif vision_feature_select_strategy == "full":
-            selected_image_feature = selected_image_feature
-        else:
-            raise ValueError(f"Unexpected select feature strategy: {self.config.vision_feature_select_strategy}")
 
         image_features = self.multi_modal_projector(selected_image_feature)
         return image_features  # type: ignore[no-any-return]