Update modeling_gpt2vision.py

modeling_gpt2vision.py

@@ -52,46 +52,46 @@ class GPT2Vision(PreTrainedModel):
     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,
+    def preprocess_inputs(self, question, image):
+        # Convert image to RGB
+        image = image.convert("RGB")
+        
+        # Tokenize the question
+        tokens = [self.tokenizer.bos_token_id]
+        q_tokens = self.tokenizer(
+            f"\n\nQuestion: {question}\n\nAnswer:",
+            add_special_tokens=False,
             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'].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)
+            max_length=384,
+        ).input_ids
+        tokens.extend(q_tokens)
+        
+        # Convert tokens to tensor
+        tokens = torch.tensor([tokens], dtype=torch.long).to(self.device)
+        
+        # Create attention mask
+        attention_mask = torch.ones_like(tokens, dtype=torch.bool)
+        
+        # Process image
+        with torch.no_grad():
+            img_embs = self.vision_encoder([image], device=self.device)
+        
+        # Get token embeddings
+        tok_embs = self.language_model.get_input_embeddings()(tokens)
+        
+        # Concatenate image embeddings and token embeddings
         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)
+        
+        # Update attention mask to include image tokens
+        img_attention = torch.ones((img_embs.size(0), img_embs.size(1)), dtype=torch.bool, 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
+        
+        return inputs_embeds, attention_mask
 
     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)
+        inputs_embeds, attention_mask = self.preprocess_inputs(question, image)
+        
         output_sequences = self.language_model.generate(
             inputs_embeds=inputs_embeds,
             attention_mask=attention_mask,
@@ -100,5 +100,7 @@ class GPT2Vision(PreTrainedModel):
             max_new_tokens=max_new_tokens,
             **kwargs
         )
-        output = self.tokenizer.decode(output_sequences[0], skip_special_tokens=True)
+        
+        # Decode the output, skipping the input tokens
+        output = self.tokenizer.decode(output_sequences[0][inputs_embeds.size(1):], skip_special_tokens=True)
         return output