refactor code

config.json

@@ -1,18 +1,18 @@
 {
-  "_name_or_path": "ibm-granite/granite-vision-3.3-2b",
+  "_name_or_path": "ibm_granite/granite-vision-3.3-2b",
   "adapter_path": null,
-  "auto_map": {
-        "AutoModel": "modeling_colgranitevision.ColGraniteVision",
-        "AutoProcessor": "processing_colgranitevision.ColGraniteVisionProcessor",
-        "AutoConfig":    "colgranitevision_config.ColGraniteVisionConfig"
+    "auto_map": {
+        "AutoModel": "modeling_granite_vision_embedding.GraniteVisionEmb",
+        "AutoProcessor": "processing_granite_vision_embedding.GraniteVisionEmbProcessor",
+        "AutoConfig":    "granite_vision_embedding_config.GraniteVisionEmbConfig"
     },
   "architectures": [
-    "ColGraniteVision"
+    "GraniteVisionEmb"
   ],
+  "base_image_feature_location": "last",
   "base_model": null,
   "emb_dim_doc": 128,
   "emb_dim_query": 128,
-  "base_image_feature_location": "last",
   "image_grid_pinpoints": [
     [
       384,
@@ -121,7 +121,7 @@
   ],
   "image_seq_length": 576,
   "image_token_index": 49155,
-  "model_type": "colgranitevision",
+  "model_type": "granitevisionemb",
   "multimodal_projector_bias": true,
   "pretrained_language_model": "",
   "pretrained_vision_tower": "",
@@ -149,12 +149,12 @@
     "rms_norm_eps": 1e-05,
     "rope_theta": 300000,
     "tie_word_embeddings": true,
-    "torch_dtype": "float32",
+    "torch_dtype": "bfloat16",
     "vocab_size": 49156
   },
   "tie_word_embeddings": true,
   "torch_dtype": "float32",
-  "transformers_version": "4.50.0.dev0",
+  "transformers_version": "4.49.0",
   "use_image_newline_parameter": true,
   "vision_config": {
     "_attn_implementation_autoset": true,
@@ -167,7 +167,7 @@
     "num_attention_heads": 16,
     "num_hidden_layers": 27,
     "patch_size": 14,
-    "torch_dtype": "float32"
+    "torch_dtype": "bfloat16"
   },
   "vision_feature_layer": [
     -24,

colgranitevision_config.py → granite_vision_embedding_config.py

@@ -1,8 +1,8 @@
 from transformers import LlavaNextConfig
 
 
-class ColGraniteVisionConfig(LlavaNextConfig):
-    model_type = "colgranitevision"
+class GraniteVisionEmbConfig(LlavaNextConfig):
+    model_type = "granitevisionemb"
 
     def __init__(self, **kwargs):
         self.base_model = kwargs.get("base_model", None)
@@ -11,3 +11,5 @@ class ColGraniteVisionConfig(LlavaNextConfig):
         self.base_image_feature_location = kwargs.get("base_image_feature_location", "last")
         self.adapter_path = kwargs.get("adapter_path", None)
         super().__init__(**kwargs)
+
+

modeling_colgranitevision.py → modeling_granite_vision_embedding.py

@@ -7,11 +7,14 @@ from transformers import LlavaNextPreTrainedModel
 from transformers.models.llava_next.modeling_llava_next import LlavaNextForConditionalGeneration
 from transformers.models.llava_next.modeling_llava_next import unpad_image, get_anyres_image_grid_shape
 
-from .colgranitevision_config import ColGraniteVisionConfig
+try:
+    from .granite_vision_embedding_config import GraniteVisionEmbConfig
+except:
+    from granite_vision_embedding_config import GraniteVisionEmbConfig
 
 
 class LlavaNextWithCustomPacking(LlavaNextForConditionalGeneration):
-    
+
     def pack_image_features(
             self,
             image_features,
@@ -93,15 +96,15 @@ class LlavaNextWithCustomPacking(LlavaNextForConditionalGeneration):
         return image_features, feature_lens
 
 
-class ColGraniteVision(LlavaNextPreTrainedModel):
+class GraniteVisionEmb(LlavaNextPreTrainedModel):
     """
-    ColGraniteVision model implementation.
+    GraniteVisionEmb model implementation.
     """
 
     main_input_name: ClassVar[str] = "doc_input_ids"  # transformers-related
-    config_class = ColGraniteVisionConfig
+    config_class = GraniteVisionEmbConfig
 
-    def __init__(self, config: ColGraniteVisionConfig):
+    def __init__(self, config: GraniteVisionEmbConfig):
         super().__init__(config=config)
 
         model = LlavaNextWithCustomPacking(config=config)
@@ -109,8 +112,6 @@ class ColGraniteVision(LlavaNextPreTrainedModel):
             self._tied_weights_keys = [f"model.language_model.{k}" for k in model.language_model._tied_weights_keys]
         self.model = model
 
-        # TODO: Wait for ColPali2 to create a ColPaliConfig to allow specifying the embedding dimension.
-        # We could do it now but it would break all the models trying to load the model from the checkpoint.
         self.dim = 128
         self.custom_text_proj = nn.Linear(self.model.config.text_config.hidden_size, self.dim)
 

preprocessor_config.json

@@ -127,7 +127,7 @@
     0.5,
     0.5
   ],
-  "processor_class": "ColGraniteVisionProcessor",
+  "processor_class": "GraniteVisionEmbProcessor",
   "resample": 3,
   "rescale_factor": 0.00392156862745098,
   "size": {

processing_colgranitevision.py → processing_granite_vision_embedding.py

@@ -21,7 +21,7 @@ def floor_by_factor(number: float, factor: int) -> int:
     return math.floor(number / factor) * factor
 
 
-class ColGraniteVisionProcessor(LlavaNextProcessor):
+class GraniteVisionEmbProcessor(LlavaNextProcessor):
     """
     Processor for ColPali.
     """
@@ -140,14 +140,14 @@ class ColGraniteVisionProcessor(LlavaNextProcessor):
             max_size=self.max_size,
             fill_color=0
         )
-        
+
     def resize_and_pad_centered_to_long_side(
-        self,
-        image: Image.Image,
-        factor: int,
-        min_size: int,
-        max_size: int,
-        fill_color=0
+            self,
+            image: Image.Image,
+            factor: int,
+            min_size: int,
+            max_size: int,
+            fill_color=0
     ) -> Image.Image:
         """
         Resizes and pads an image such that:
@@ -183,10 +183,10 @@ class ColGraniteVisionProcessor(LlavaNextProcessor):
 
         # Resize the image
         resized_image = image.resize((target_width, target_height), Image.LANCZOS)
-        final_image =resized_image.convert("RGB")
+        final_image = resized_image.convert("RGB")
 
         return final_image
-        
+
     def resize_and_pad_centered(self,
                                 image: Image.Image,
                                 factor: int,
@@ -439,4 +439,4 @@ class ColGraniteVisionProcessor(LlavaNextProcessor):
         assert scores.shape[0] == len(qs), f"Expected {len(qs)} scores, got {scores.shape[0]}"
 
         scores = scores.to(torch.float32)
-        return scores
+        return scores

processor_config.json

@@ -1,6 +1,6 @@
 {
-  "processor_class": "ColGraniteVisionProcessor",
+  "processor_class": "GraniteVisionEmbProcessor",
   "auto_map": {
-    "AutoProcessor": "processing_colgranitevision.ColGraniteVisionProcessor"
+    "AutoProcessor": "processing_granite_vision_embedding.GraniteVisionEmbProcessor"
   }
 }