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 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,
+    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=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
+            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)
         inputs_embeds = torch.cat((tok_embs[:, 0:1, :], img_embs, tok_embs[:, 1:, :]), dim=1)
-        
-        # 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)
+        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
+        return inputs_embeds, attention_mask, input_ids
 
     def generate(self, question, image, max_new_tokens=30, **kwargs):
-        inputs_embeds, attention_mask = self.preprocess_inputs(question, image)
-        
+        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)
         output_sequences = self.language_model.generate(
             inputs_embeds=inputs_embeds,
             attention_mask=attention_mask,
@@ -100,7 +100,5 @@ class GPT2Vision(PreTrainedModel):
             max_new_tokens=max_new_tokens,
             **kwargs
         )
-        
-        # Decode the output, skipping the input tokens
-        output = self.tokenizer.decode(output_sequences[0][inputs_embeds.size(1):], skip_special_tokens=True)
+        output = self.tokenizer.decode(output_sequences[0], skip_special_tokens=True)
         return output