Update modeling_arlow_gpt.py

modeling_arlow_gpt.py

@@ -1,15 +1,14 @@
 # modeling_arlow_gpt.py
-import torch
-import torch.nn as nn
-from transformers import PreTrainedModel, CLIPModel, GPT2Model
+from transformers import PreTrainedModel, PreTrainedModel, CLIPModel, GPT2Model
+from transformers.modeling_outputs import Seq2SeqLMOutput
 from typing import Optional, Union, Dict, Tuple
 from .configuration_arlow_gpt import ArlowGPTConfig
 
 class ArlowGPTPreTrainedModel(PreTrainedModel):
-    """Base class for ArlowGPT model."""
     config_class = ArlowGPTConfig
     base_model_prefix = "arlow_gpt"
     supports_gradient_checkpointing = True
+    _keys_to_ignore_on_load_missing = [r"clip", r"gpt2"]
 
     def _init_weights(self, module):
         if isinstance(module, nn.Linear):
@@ -18,59 +17,32 @@ class ArlowGPTPreTrainedModel(PreTrainedModel):
                 module.bias.data.zero_()
 
 class ArlowGPTModel(ArlowGPTPreTrainedModel):
+    # Same as before
+
+class ArlowGPTForImageTextToText(ArlowGPTPreTrainedModel):
     def __init__(self, config: ArlowGPTConfig):
         super().__init__(config)
-        
-        self.clip = CLIPModel.from_pretrained(config.clip_model_name)
-        self.gpt2 = GPT2Model.from_pretrained(config.gpt2_model_name)
-
-        self.feature_projection = nn.Linear(
-            self.clip.vision_model.config.hidden_size + self.gpt2.config.hidden_size,
-            config.projection_dim
-        )
+        self.arlow_gpt = ArlowGPTModel(config)
+        self.lm_head = nn.Linear(config.projection_dim, config.vocab_size)
 
         # Initialize weights and apply final processing
         self.post_init()
 
-    def forward(
-        self,
-        input_ids: torch.Tensor,
-        attention_mask: torch.Tensor,
-        pixel_values: torch.Tensor,
-        return_dict: bool = True,
-    ) -> Union[torch.Tensor, Dict[str, torch.Tensor]]:
-        vision_outputs = self.clip.get_image_features(pixel_values=pixel_values)
-        text_outputs = self.gpt2(
-            input_ids=input_ids,
-            attention_mask=attention_mask
-        ).last_hidden_state
-
-        batch_size = text_outputs.shape[0]
-        seq_length = text_outputs.shape[1]
-
-        vision_features = vision_outputs.unsqueeze(1).expand(
-            batch_size, seq_length, -1
-        )
-
-        combined_features = torch.cat(
-            [vision_features, text_outputs],
-            dim=-1
-        )
-
-        hidden_states = self.feature_projection(combined_features)
+    def save_pretrained(self, save_directory, **kwargs):
+        """Override save_pretrained to save all components"""
+        super().save_pretrained(save_directory, **kwargs)
+        self.arlow_gpt.save_pretrained(save_directory)
 
-        if return_dict:
-            return {"hidden_states": hidden_states}
-        return hidden_states
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
+        """Override from_pretrained to handle custom loading logic"""
+        config = kwargs.get("config", None)
+        if config is None:
+            config = ArlowGPTConfig.from_pretrained(pretrained_model_name_or_path)
+            kwargs["config"] = config
 
-class ArlowGPTForCausalLM(ArlowGPTPreTrainedModel):
-    def __init__(self, config: ArlowGPTConfig):
-        super().__init__(config)
-        self.arlow_gpt = ArlowGPTModel(config)
-        self.output_projection = nn.Linear(config.projection_dim, config.vocab_size)
-        
-        # Initialize weights and apply final processing
-        self.post_init()
+        config._name_or_path = pretrained_model_name_or_path
+        return super().from_pretrained(pretrained_model_name_or_path, *model_args, **kwargs)
 
     def forward(
         self,
@@ -88,7 +60,7 @@ class ArlowGPTForCausalLM(ArlowGPTPreTrainedModel):
         )
 
         hidden_states = outputs["hidden_states"]
-        logits = self.output_projection(hidden_states)
+        logits = self.lm_head(hidden_states)
 
         loss = None
         if labels is not None:
@@ -96,10 +68,10 @@ class ArlowGPTForCausalLM(ArlowGPTPreTrainedModel):
             loss = loss_fct(logits.view(-1, self.config.vocab_size), labels.view(-1))
 
         if return_dict:
-            return {
-                "loss": loss,
-                "logits": logits
-            }
+            return Seq2SeqLMOutput(
+                loss=loss,
+                logits=logits,
+            )
         return (loss, logits) if loss is not None else logits
 
     def prepare_inputs_for_generation(