Update modeling_gpt2vision.py

modeling_gpt2vision.py

@@ -45,42 +45,21 @@ class GPT2Vision(PreTrainedModel):
         self.language_model = GPT2LMHeadModel(config.gpt2_config)
         self.language_model.resize_token_embeddings(len(tokenizer))
         self.tokenizer = tokenizer
-        tokenizer.pad_token = tokenizer.eos_token
+        self.tokenizer.pad_token = self.tokenizer.eos_token
         self.image_token_id = self.tokenizer.convert_tokens_to_ids(IMAGE_TOKEN)
 
     @property
     def device(self):
         return next(self.language_model.parameters()).device
 
-    def tokenize_encode(self, batch, device):
-        text = batch['text']
-        images = batch['image']
-        if isinstance(text, str):
-            text = [text]
-        input_texts = [f"{IMAGE_TOKEN}{t}" for t in text]
-        text_inputs = self.tokenizer(
-            input_texts,
-            padding='max_length',
-            truncation=True,
-            max_length=768,
-            return_tensors="pt",
-            pad_to_multiple_of=8,
-        ).to(device)
-        pixel_values = self.vision_encoder(images, device)
-        return {
-            "input_ids": text_inputs.input_ids,
-            "attention_mask": text_inputs.attention_mask,
-            "pixel_values": pixel_values
-        }
-
     def preprocess_inputs(self, batch):
-        pixel_values = batch['pixel_values']
+        img_embs = batch['pixel_values']
         input_ids = batch['input_ids']
         attention_mask = batch['attention_mask']
         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)
+        img_embs = img_embs.to(self.device)
+        
         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)
@@ -88,10 +67,22 @@ class GPT2Vision(PreTrainedModel):
         return inputs_embeds, attention_mask, input_ids
 
     def generate(self, question, image, max_new_tokens=30, **kwargs):
-        prompt = f"Question: {question}\nAnswer:"
-        batch = {"image": [image], "text": prompt}
-        encoded_batch = self.tokenize_encode(batch, self.device)
-        inputs_embeds, attention_mask, input_ids = self.preprocess_inputs(encoded_batch)
+        # Process the image
+        img_embs = self.vision_encoder(image.unsqueeze(0), device=self.device)
+        img_embs = self.mlp(img_embs)
+        
+        # Tokenize the question
+        prompt = f"{IMAGE_TOKEN}Question: {question}\nAnswer:"
+        encoded_input = self.tokenizer(prompt, return_tensors="pt", padding=True, truncation=True)
+        
+        batch = {
+            "pixel_values": img_embs,
+            "input_ids": encoded_input.input_ids,
+            "attention_mask": encoded_input.attention_mask
+        }
+        
+        inputs_embeds, attention_mask, input_ids = self.preprocess_inputs(batch)
+        
         output_sequences = self.language_model.generate(
             inputs_embeds=inputs_embeds,
             attention_mask=attention_mask,
@@ -100,5 +91,6 @@ class GPT2Vision(PreTrainedModel):
             max_new_tokens=max_new_tokens,
             **kwargs
         )
+        
         output = self.tokenizer.decode(output_sequences[0], skip_special_tokens=True)
         return output