Update modeling_gme_qwen2vl.py

modeling_gme_qwen2vl.py

@@ -12,16 +12,25 @@ import torch
 from PIL import Image
 from torch.utils.data import DataLoader
 from tqdm.autonotebook import tqdm
-from transformers import (
-    AutoProcessor,
-    PreTrainedModel,
+from transformers import AutoProcessor, PreTrainedModel
+from transformers.models.qwen2_vl.modeling_qwen2_vl import (
+    Qwen2VisionTransformerPretrainedModel,
     Qwen2VLConfig,
     Qwen2VLForConditionalGeneration,
+    Qwen2VLModel,
+)
+from transformers.utils.versions import require_version
+
+
+require_version(
+    "transformers<4.52.0",
+    "This code has some issues with transformers>=4.52.0, please downgrade: pip install transformers==4.51.3"
 )
-import os
 
 
 class GmeQwen2VLConfig(Qwen2VLConfig):
+    # model_type = ''
+
     def __init__(
         self,
         min_image_tokens: int = 256,
@@ -35,14 +44,25 @@ class GmeQwen2VLConfig(Qwen2VLConfig):
         self.max_length = max_length
 
 
-class GmeQwen2VLForVision2Seq(PreTrainedModel):
+class GmeQwen2VL(PreTrainedModel):
     config_class = GmeQwen2VLConfig
-    base_model_prefix: str = "base"
+    base_model_prefix = "model"
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["Qwen2VLDecoderLayer", "Qwen2VLVisionBlock"]
+    # _skip_keys_device_placement = "past_key_values"
+    _supports_flash_attn_2 = True
+    _supports_sdpa = True
+    # _supports_cache_class = True
+    _supports_static_cache = False  # TODO (joao): fix. torch.compile failing probably due to `cache_positions`
+    # _tied_weights_keys = ["lm_head.weight"]
 
     def __init__(self, config: GmeQwen2VLConfig, **kwargs: Any) -> None:
         super().__init__(config)
-        self.base = Qwen2VLForConditionalGeneration.from_pretrained(config._name_or_path)
-        self.base.tie_weights()  # It's important to produce same outputs.
+        self.visual = Qwen2VisionTransformerPretrainedModel._from_config(config.vision_config)
+        self.model = Qwen2VLModel(config)
+        self.vocab_size = config.vocab_size
+        # self.lm_head = torch.nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        self.rope_deltas = None  # cache rope_deltas here
 
         min_pixels: int = config.min_image_tokens * 28 * 28
         max_pixels: int = config.max_image_tokens * 28 * 28
@@ -55,6 +75,9 @@ class GmeQwen2VLForVision2Seq(PreTrainedModel):
         self.default_instruction: str = "You are a helpful assistant."
         self.sep: str = " "
 
+        # Initialize weights and apply final processing
+        self.post_init()
+
     def forward(
         self,
         input_ids: Optional[torch.LongTensor] = None,
@@ -70,21 +93,21 @@ class GmeQwen2VLForVision2Seq(PreTrainedModel):
         **kwargs
     ) -> torch.Tensor:
         if inputs_embeds is None:
-            inputs_embeds = self.base.model.embed_tokens(input_ids)
+            inputs_embeds = self.model.get_input_embeddings()(input_ids)
             if pixel_values is not None:
-                pixel_values = pixel_values.type(self.base.visual.get_dtype())
-                image_embeds = self.base.visual(pixel_values, grid_thw=image_grid_thw).to(inputs_embeds.device)
-                image_mask = input_ids == self.base.config.image_token_id
+                pixel_values = pixel_values.type(self.visual.get_dtype())
+                image_embeds = self.visual(pixel_values, grid_thw=image_grid_thw).to(inputs_embeds.device)
+                image_mask = input_ids == self.config.image_token_id
                 inputs_embeds[image_mask] = image_embeds
             # if pixel_values_videos is not None:
-            #     pixel_values_videos = pixel_values_videos.type(self.base.visual.get_dtype())
-            #     video_embeds = self.base.visual(pixel_values_videos, grid_thw=video_grid_thw).to(inputs_embeds.device)
-            #     video_mask = input_ids == self.base.config.video_token_id
+            #     pixel_values_videos = pixel_values_videos.type(self.visual.get_dtype())
+            #     video_embeds = self.visual(pixel_values_videos, grid_thw=video_grid_thw).to(inputs_embeds.device)
+            #     video_mask = input_ids == self.config.video_token_id
             #     inputs_embeds[video_mask] = video_embeds
             if attention_mask is not None:
                 attention_mask = attention_mask.to(inputs_embeds.device)
 
-        outputs = self.base.model(
+        outputs = self.model(
             input_ids=None,
             position_ids=position_ids,
             attention_mask=attention_mask,
@@ -311,3 +334,4 @@ def fetch_image(image: str | Image.Image, size_factor: int = IMAGE_FACTOR) -> Im
 
     return image
 ###
+