Update modeling_gpt2vision.py

modeling_gpt2vision.py

@@ -20,102 +20,106 @@ def resize_token_embeds(model_name="openai-community/gpt2"):
 
 tokenizer = resize_token_embeds()
 
+class MLP(nn.Module):
+    def __init__(self, in_features: int, hidden_features: int = None, out_features: int = None):
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        self.fc1 = nn.Linear(in_features, hidden_features)
+        self.act = nn.GELU(approximate="tanh")
+        self.fc2 = nn.Linear(hidden_features, out_features)
+        self.dropout = nn.Dropout(p=0.1)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.dropout(x)
+        x = self.fc2(x)
+        return x
+        
 class GPT2Vision(PreTrainedModel):
     config_class = GPT2VisionConfig
 
     def __init__(self, config):
         super().__init__(config)
         self.vision_encoder = VisionEncoder()
-       
-        if isinstance(config.gpt2_config, dict):
-            gpt2_config = GPT2Config(**config.gpt2_config)
-        else:
-            gpt2_config = config.gpt2_config
-        self.text_model = GPT2LMHeadModel(gpt2_config)
-
-        self.text_model.resize_token_embeddings(len(tokenizer))
+        self.mlp = MLP(in_features=768, hidden_features=768 * 4, out_features=768)
+        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.image_token_id = self.tokenizer.convert_tokens_to_ids(IMAGE_TOKEN)
+        self.img_tokens = 197  # This should match IMG_TOKENS in your training code
 
-        
     @property
     def device(self):
-        return self.text_model.device
+        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",
+        ).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)
+        
+        return {
+            "input_ids": text_inputs.input_ids,
+            "attention_mask": attention_mask,
+            "images": images
+        }
+            
+    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)
+        print("img_embs",img_embs.size())
+        img_embs = self.mlp(img_embs)
+        
+        tok_embs = self.language_model.get_input_embeddings()(input_ids)
 
-    def encode_image(self, image,device):
-        return self.vision_encoder(image,device=device)
+        inputs_embeds = torch.cat((tok_embs[:, 0:1, :], img_embs, tok_embs[:, 1:, :]), dim=1)
 
-    def input_embeds(self, prompt, image_embeds, tokenizer):
-        def _tokenize(txt):
-            return tokenizer(
-                txt, return_tensors="pt", add_special_tokens=False
-            ).input_ids.to(self.device)
+        # 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]})."
 
-        text_emb = self.text_model.get_input_embeddings()
+        return inputs_embeds, attention_mask
 
-        # Add BOS token
-        embeds = []
-        embeds.append(
-            text_emb((torch.tensor([[tokenizer.bos_token_id]], device=self.device)))
+    def forward(self, batch, **kwargs):
+        inputs_embeds, attention_mask = self.preprocess_inputs(batch)
+    
+        outputs = self.language_model(
+            inputs_embeds=inputs_embeds,
+            attention_mask=attention_mask,
+            **kwargs
         )
-
-        if "<image>" not in prompt:
-            embeds.append(text_emb(_tokenize(prompt)))
-        else:
-            assert prompt.count("<image>") == 1
-            before, after = prompt.split("<image>")
-            embeds.append(text_emb(_tokenize(f"{before}<image>")))
-            embeds.append(image_embeds.to(self.device))
-            embeds.append(text_emb(_tokenize(f"</image>{after}")))
-
-        return torch.cat(embeds, dim=1)
-
-    def generate(
-        self,
-        image_embeds,
-        prompt,
-        tokenizer,
-        eos_text="<|endoftext|>",
-        max_new_tokens=128,
-        **kwargs,
-    ):
-        eos_tokens = tokenizer(eos_text, add_special_tokens=False)["input_ids"]
-
-        generate_config = {
-            "eos_token_id": eos_tokens,
-            "bos_token_id": tokenizer.bos_token_id,
-            "pad_token_id": tokenizer.eos_token_id,
-            "max_new_tokens": max_new_tokens,
-            **kwargs,
-        }
-
-        with torch.no_grad():
-            
-            inputs_embeds = self.input_embeds(prompt, image_embeds, tokenizer)
-            print("inputs_embeds",inputs_embeds.size())
-            output_ids = self.text_model.generate(
-                inputs_embeds=inputs_embeds, **generate_config
-            )
-
-        return tokenizer.batch_decode(output_ids, skip_special_tokens=True)
-
-    def answer_question(
-        self,
-        image_embeds,
-        question,
-        tokenizer,
-        chat_history="",
-        result_queue=None,
-        **kwargs,
-    ):
-        prompt = f"<image>\n\n{chat_history}Question: {question}\n\nAnswer: "
-        answer = self.generate(
-            image_embeds,
-            prompt,
-            tokenizer,
-            eos_text="<|endoftext|>",
-            max_new_tokens=256,
-            **kwargs,
-        )[0]
-        return answer 
+        return outputs
+        
+    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 = self.preprocess_inputs(encoded_batch)
+        output_sequences = self.language_model.generate(
+            inputs_embeds=inputs_embeds,
+            attention_mask=attention_mask,
+            pad_token_id=self.tokenizer.eos_token_id,
+            eos_token_id=self.tokenizer.eos_token_id,
+            max_new_tokens=max_new_tokens,
+            **kwargs
+        )
+        output = self.tokenizer.decode(output_sequences[0], skip_special_tokens=True)
+        return output