Update modeling_gpt2vision.py

modeling_gpt2vision.py

@@ -1,10 +1,8 @@
 import torch
-from torch import nn
-from transformers import PreTrainedModel,AutoTokenizer
-import re
-
-from .vision_encoder import VisionEncoder
+import torch.nn as nn
+from transformers import PreTrainedModel, AutoTokenizer
 from .configuration_gpt2vision import GPT2VisionConfig
+from .vision_encoder import VisionEncoder
 from .modeling_gpt2 import GPT2LMHeadModel
 
 IMAGE_TOKEN = "<image>"
@@ -36,7 +34,7 @@ class MLP(nn.Module):
         x = self.dropout(x)
         x = self.fc2(x)
         return x
-        
+
 class GPT2Vision(PreTrainedModel):
     config_class = GPT2VisionConfig
 
@@ -49,7 +47,6 @@ class GPT2Vision(PreTrainedModel):
         self.tokenizer = tokenizer
         tokenizer.pad_token = tokenizer.eos_token
         self.image_token_id = self.tokenizer.convert_tokens_to_ids(IMAGE_TOKEN)
-        self.img_tokens = 197  
 
     @property
     def device(self):
@@ -60,48 +57,41 @@ class GPT2Vision(PreTrainedModel):
         images = batch['image']
         if isinstance(text, str):
             text = [text]
-
+        input_texts = [f"{IMAGE_TOKEN}{t}" for t in text]
         text_inputs = self.tokenizer(
-            text,
+            input_texts,
             padding='max_length',
             truncation=True,
             max_length=768,
             return_tensors="pt",
+            pad_to_multiple_of=8,
         ).to(device)
-        
-        # Adjust attention mask to account for image tokens and the extra <image> token
-        batch_size = text_inputs.input_ids.shape[0]
-        img_attention = torch.ones((batch_size, self.img_tokens + 1), dtype=torch.long, device=device)
-        attention_mask = torch.cat([img_attention, text_inputs.attention_mask[:, 1:]], dim=1)
-        
+        pixel_values = self.vision_encoder(images, device)
         return {
             "input_ids": text_inputs.input_ids,
-            "attention_mask": attention_mask,
-            "images": images
+            "attention_mask": text_inputs.attention_mask,
+            "pixel_values": pixel_values
         }
-            
-    def preprocess_inputs(self, batch):
-        images = batch['images']
-        input_ids = batch['input_ids'].to(self.device)
-        attention_mask = batch['attention_mask'].to(self.device)
 
-        img_embs = self.vision_encoder(images, device=self.device)
-        img_embs = self.mlp(img_embs)
+    def preprocess_inputs(self, batch):
+        pixel_values = batch['pixel_values'].squeeze(1)
+        input_ids = batch['input_ids'].squeeze(1)
+        attention_mask = batch['attention_mask'].squeeze(1)
+        input_ids = input_ids.to(self.device)
+        attention_mask = attention_mask.to(self.device)
+        pixel_values = pixel_values.to(self.device)
+        img_embs = self.mlp(pixel_values)
         tok_embs = self.language_model.get_input_embeddings()(input_ids)
-
         inputs_embeds = torch.cat((tok_embs[:, 0:1, :], img_embs, tok_embs[:, 1:, :]), dim=1)
+        img_attention = torch.ones((img_embs.size(0), img_embs.size(1)), dtype=torch.long, device=self.device)
+        attention_mask = torch.cat((attention_mask[:, 0:1], img_attention, attention_mask[:, 1:]), dim=1)
+        return inputs_embeds, attention_mask, input_ids
 
-        # Ensure the attention mask aligns with the inputs_embeds
-        assert inputs_embeds.shape[1] == attention_mask.shape[1], f"Mismatch between embeddings ({inputs_embeds.shape[1]}) and attention mask length ({attention_mask.shape[1]})."
-
-        return inputs_embeds, attention_mask
-
-        
     def generate(self, question, image, max_new_tokens=30, **kwargs):
-        prompt = f"\n\nQuestion: {question}\n\nAnswer:"
+        prompt = f"Question: {question}\nAnswer:"
         batch = {"image": [image], "text": prompt}
         encoded_batch = self.tokenize_encode(batch, self.device)
-        inputs_embeds, attention_mask = self.preprocess_inputs(encoded_batch)
+        inputs_embeds, attention_mask, input_ids = self.preprocess_inputs(encoded_batch)
         output_sequences = self.language_model.generate(
             inputs_embeds=inputs_embeds,
             attention_mask=attention_mask,