Upload 3 files

modeling_colgranitevision/config.json

@@ -0,0 +1,167 @@
+{
+  "adapter_path": "",
+  "architectures": [
+    "ColGraniteVision"
+  ],
+    "auto_map": {
+        "AutoModelForVision2Seq": "modeling_colgranitevision.ColGraniteVision",
+        "AutoProcessor": "processing_colgranitevision.ColGraniteVisionProcessor"
+    },
+  "base_model": "ibm-granite/granite-vision-3.1-2b-preview",
+  "emb_dim_doc": 128,
+  "emb_dim_query": 128,
+  "image_grid_pinpoints": [
+    [
+      384,
+      768
+    ],
+    [
+      384,
+      1152
+    ],
+    [
+      384,
+      1536
+    ],
+    [
+      384,
+      1920
+    ],
+    [
+      384,
+      2304
+    ],
+    [
+      384,
+      2688
+    ],
+    [
+      384,
+      3072
+    ],
+    [
+      384,
+      3456
+    ],
+    [
+      384,
+      3840
+    ],
+    [
+      768,
+      384
+    ],
+    [
+      768,
+      768
+    ],
+    [
+      768,
+      1152
+    ],
+    [
+      768,
+      1536
+    ],
+    [
+      768,
+      1920
+    ],
+    [
+      1152,
+      384
+    ],
+    [
+      1152,
+      768
+    ],
+    [
+      1152,
+      1152
+    ],
+    [
+      1536,
+      384
+    ],
+    [
+      1536,
+      768
+    ],
+    [
+      1920,
+      384
+    ],
+    [
+      1920,
+      768
+    ],
+    [
+      2304,
+      384
+    ],
+    [
+      2688,
+      384
+    ],
+    [
+      3072,
+      384
+    ],
+    [
+      3456,
+      384
+    ],
+    [
+      3840,
+      384
+    ]
+  ],
+  "image_seq_length": 576,
+  "image_token_index": 49155,
+  "model_type": "llava_next",
+  "multimodal_projector_bias": true,
+  "projector_hidden_act": "gelu",
+  "text_config": {
+    "architectures": [
+      "GraniteForCausalLM"
+    ],
+    "attention_dropout": 0.1,
+    "attention_multiplier": 0.015625,
+    "bos_token_id": 0,
+    "embedding_multiplier": 12.0,
+    "eos_token_id": 0,
+    "hidden_size": 2048,
+    "intermediate_size": 8192,
+    "logits_scaling": 8.0,
+    "max_position_embeddings": 16384,
+    "model_type": "granite",
+    "num_hidden_layers": 40,
+    "num_key_value_heads": 8,
+    "pad_token_id": 0,
+    "residual_multiplier": 0.22,
+    "rms_norm_eps": 1e-05,
+    "rope_theta": 300000,
+    "tie_word_embeddings": true,
+    "torch_dtype": "bfloat16",
+    "vocab_size": 49156
+  },
+  "torch_dtype": "float32",
+  "transformers_version": "4.50.0.dev0",
+  "use_image_newline_parameter": true,
+  "vision_config": {
+    "hidden_size": 1152,
+    "image_size": 384,
+    "intermediate_size": 4304,
+    "model_type": "siglip_vision_model",
+    "num_attention_heads": 16,
+    "num_hidden_layers": 27,
+    "patch_size": 14
+  },
+  "vision_feature_layer": [
+    -24,
+    -20,
+    -12,
+    -1
+  ],
+  "vision_feature_select_strategy": "full"
+}

modeling_colgranitevision/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ddf363fafb46eb2bda156c288b4023345220f243dfa40cdf8dea7e4e3e2f6ab0
+size 113333568

modeling_colgranitevision/modeling_colgranitevision.py

@@ -0,0 +1,114 @@
+from typing import ClassVar, Optional
+
+import torch
+from torch import nn
+from transformers import LlavaNextConfig, \
+    LlavaNextPreTrainedModel
+
+from custom_llava_next import LlavaNextWithCustomPacking as LlavaNextForConditionalGeneration
+
+
+# from transformers.models.paligemma.modeling_paligemma import (
+#     PaliGemmaConfig,
+#     PaliGemmaForConditionalGeneration,
+#     PaliGemmaPreTrainedModel,
+# )
+
+
+class ColGraniteVision(LlavaNextPreTrainedModel):
+    """
+    ColGraniteVision model implementation.
+    """
+
+    main_input_name: ClassVar[str] = "doc_input_ids"  # transformers-related
+
+    def __init__(self, config: LlavaNextConfig):
+        super().__init__(config=config)
+
+        model = LlavaNextForConditionalGeneration(config=config)
+        if model.language_model._tied_weights_keys is not None:
+            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)
+
+        self.post_init()
+
+    def forward(self, *args, **kwargs) -> torch.Tensor:
+        # Delete output_hidden_states from kwargs
+        kwargs.pop("output_hidden_states", None)
+        if "pixel_values" in kwargs:
+            kwargs["pixel_values"] = kwargs["pixel_values"].to(dtype=self.dtype)
+
+        outputs = self.model(*args, output_hidden_states=True, **kwargs)  # (batch_size, sequence_length, hidden_size)
+        last_hidden_states = outputs.hidden_states[-1]  # (batch_size, sequence_length, hidden_size)
+
+        attention_mask = kwargs["attention_mask"]
+        if "pixel_values" in kwargs:
+            input_ids = kwargs['input_ids']
+            image_mask = (input_ids == self.config.image_token_index)
+            # inputs_embeds = last_hidden_states.masked_scatter(image_mask)
+            N, M = image_mask.shape
+            # Create an index matrix: each row is 0, 1, ..., M-1
+            idx = torch.arange(M, device=image_mask.device).expand(N, M)
+            # Replace False positions with -1 so they are ignored by topk (since all valid indices are >=0)
+            masked_idx = torch.where(image_mask, idx, torch.tensor(-1, device=image_mask.device))
+            topk_values, _ = torch.topk(masked_idx, k=729, dim=1)
+            last_k_indices, _ = torch.sort(topk_values, dim=1)
+            last_k_indices_exp = last_k_indices.unsqueeze(-1).expand(-1, -1, last_hidden_states.size(-1))
+            last_hidden_states = torch.gather(last_hidden_states, 1, last_k_indices_exp)
+            attention_mask = torch.gather(attention_mask, 1, last_k_indices)
+
+        attention_mask = attention_mask.unsqueeze(-1)
+
+        proj = self.custom_text_proj(last_hidden_states)  # (batch_size, sequence_length, dim)
+
+        # L2 normalization
+        proj = proj / (proj.norm(dim=-1, keepdim=True) + 1e-8)
+
+        # proj = proj * kwargs["attention_mask"].unsqueeze(-1)  # (batch_size, sequence_length, dim)
+        proj = proj * attention_mask  # (batch_size, sequence_length, dim)
+
+        return proj
+
+    def get_input_embeddings(self):
+        return self.model.language_model.get_input_embeddings()
+
+    def set_input_embeddings(self, value):
+        self.model.language_model.set_input_embeddings(value)
+
+    def get_output_embeddings(self):
+        return self.model.language_model.get_output_embeddings()
+
+    def set_output_embeddings(self, new_embeddings):
+        self.model.language_model.set_output_embeddings(new_embeddings)
+
+    def set_decoder(self, decoder):
+        self.model.language_model.set_decoder(decoder)
+
+    def get_decoder(self):
+        return self.model.language_model.get_decoder()
+
+    def tie_weights(self):
+        return self.model.language_model.tie_weights()
+
+    def resize_token_embeddings(
+            self,
+            new_num_tokens: Optional[int] = None,
+            pad_to_multiple_of=None,
+    ) -> nn.Embedding:
+        model_embeds = self.model.language_model.resize_token_embeddings(new_num_tokens, pad_to_multiple_of)
+
+        # Update vocab size
+        self.config.text_config.vocab_size = model_embeds.num_embeddings
+        self.config.vocab_size = model_embeds.num_embeddings
+        self.model.vocab_size = model_embeds.num_embeddings
+
+        return model_embeds
+
+    @property
+    def patch_size(self) -> int:
+        return self.model.vision_tower.config.patch_size