Update modeling_gpt2vision.py

modeling_gpt2vision.py

@@ -18,6 +18,30 @@ def resize_token_embeds(model_name="openai-community/gpt2"):
 
 tokenizer = resize_token_embeds()
 
+def create_labels(input_ids, tokenizer, attention_mask):
+    labels = input_ids.clone()
+    
+    labels[attention_mask == 0] = -100
+
+    answer_start_tokens = tokenizer.encode("Answer:", add_special_tokens=False)
+
+    for i, seq in enumerate(input_ids):
+        # Find the start of the answer
+        answer_start = (seq == answer_start_tokens[0]).nonzero(as_tuple=True)[0]
+        if len(answer_start) > 0:
+            answer_start = answer_start[0]
+            if seq[answer_start:answer_start+len(answer_start_tokens)].tolist() == answer_start_tokens:
+                # Mask out everything before the answer
+                labels[i, :answer_start] = -100
+                
+                # Find the end of the sequence (last non-padding token)
+                sequence_end = attention_mask[i].nonzero(as_tuple=True)[0][-1]
+                
+                # Keep the last token (EOS) as part of the label
+                labels[i, sequence_end+1:] = -100
+
+    return labels
+    
 class MLP(nn.Module):
     def __init__(self, in_features: int, hidden_features: int = None, out_features: int = None):
         super().__init__()
@@ -62,7 +86,7 @@ class GPT2Vision(PreTrainedModel):
             input_texts,
             padding='max_length',
             truncation=True,
-            max_length=384,
+            max_length=768,
             return_tensors="pt",
         ).to(device)
         pixel_values = self.vision_encoder(images, device)
@@ -72,20 +96,46 @@ class GPT2Vision(PreTrainedModel):
             "pixel_values": pixel_values
         }
 
+    def freeze_model_components(self, freeze_vision=True, freeze_language=True, freeze_mlp=False):
+        for param in self.vision_model.parameters():
+            param.requires_grad = not freeze_vision
+        for param in self.language_model.parameters():
+            param.requires_grad = not freeze_language
+        for param in self.mlp.parameters():
+            param.requires_grad = not freeze_mlp
+            
     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)
+
+        labels = create_labels(input_ids, self.tokenizer, attention_mask)
+        labels = labels.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
 
+        img_labels = torch.full((labels.size(0), img_embs.size(1)), fill_value=-100, dtype=torch.long, device=self.device)
+        labels = torch.cat((labels[:, 0:1], img_labels, labels[:, 1:]), dim=1)
+
+        return inputs_embeds, attention_mask, input_ids, labels
+
+
+    def forward(self, batch, **kwargs):
+        inputs_embeds, attention_mask, input_ids, labels = self.preprocess_inputs(batch)
+    
+        outputs = self.language_model(inputs_embeds=inputs_embeds, attention_mask=attention_mask, labels=labels)
+        return outputs
+        
     def generate(self, question, image, max_new_tokens=30, **kwargs):
         prompt = f"Question: {question}\nAnswer:"
         batch = {"image": [image], "text": prompt}