Update modeling_arlow_gpt.py

modeling_arlow_gpt.py

@@ -1,42 +1,34 @@
+# modeling_arlow_gpt.py
 import torch
 import torch.nn as nn
 from transformers import PreTrainedModel, CLIPModel, GPT2Model
-from typing import Optional, Union, Dict
+from typing import Optional, Union, Dict, Tuple
 from .configuration_arlow_gpt import ArlowGPTConfig
 
 class ArlowGPTPreTrainedModel(PreTrainedModel):
-    """
-    An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained models.
-    """
+    """Base class for ArlowGPT model."""
     config_class = ArlowGPTConfig
     base_model_prefix = "arlow_gpt"
     supports_gradient_checkpointing = True
 
     def _init_weights(self, module):
         if isinstance(module, nn.Linear):
-            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range if hasattr(self.config, "initializer_range") else 0.02)
+            module.weight.data.normal_(mean=0.0, std=0.02)
             if module.bias is not None:
                 module.bias.data.zero_()
 
-class ArlowGPT(ArlowGPTPreTrainedModel):
+class ArlowGPTModel(ArlowGPTPreTrainedModel):
     def __init__(self, config: ArlowGPTConfig):
         super().__init__(config)
         
-        # Load the models
         self.clip = CLIPModel.from_pretrained(config.clip_model_name)
         self.gpt2 = GPT2Model.from_pretrained(config.gpt2_model_name)
 
-        # Projection layers
         self.feature_projection = nn.Linear(
             self.clip.vision_model.config.hidden_size + self.gpt2.config.hidden_size,
             config.projection_dim
         )
 
-        self.output_projection = nn.Linear(
-            config.projection_dim,
-            config.vocab_size
-        )
-
         # Initialize weights and apply final processing
         self.post_init()
 
@@ -45,7 +37,6 @@ class ArlowGPT(ArlowGPTPreTrainedModel):
         input_ids: torch.Tensor,
         attention_mask: torch.Tensor,
         pixel_values: torch.Tensor,
-        labels: Optional[torch.Tensor] = None,
         return_dict: bool = True,
     ) -> Union[torch.Tensor, Dict[str, torch.Tensor]]:
         vision_outputs = self.clip.get_image_features(pixel_values=pixel_values)
@@ -66,8 +57,38 @@ class ArlowGPT(ArlowGPTPreTrainedModel):
             dim=-1
         )
 
-        projected_features = self.feature_projection(combined_features)
-        logits = self.output_projection(projected_features)
+        hidden_states = self.feature_projection(combined_features)
+
+        if return_dict:
+            return {"hidden_states": hidden_states}
+        return hidden_states
+
+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()
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: torch.Tensor,
+        pixel_values: torch.Tensor,
+        labels: Optional[torch.Tensor] = None,
+        return_dict: bool = True,
+    ) -> Union[torch.Tensor, Dict[str, torch.Tensor]]:
+        outputs = self.arlow_gpt(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            pixel_values=pixel_values,
+            return_dict=True
+        )
+
+        hidden_states = outputs["hidden_states"]
+        logits = self.output_projection(hidden_states)
 
         loss = None
         if labels is not None:
@@ -79,4 +100,20 @@ class ArlowGPT(ArlowGPTPreTrainedModel):
                 "loss": loss,
                 "logits": logits
             }
-        return logits
+        return (loss, logits) if loss is not None else logits
+
+    def prepare_inputs_for_generation(
+        self, input_ids, past=None, attention_mask=None, **kwargs
+    ):
+        # only last token for inputs_ids if past is defined in kwargs
+        if past:
+            input_ids = input_ids[:, -1].unsqueeze(-1)
+            if attention_mask is not None:
+                attention_mask = attention_mask[:, -1].unsqueeze(-1)
+
+        return {
+            "input_ids": input_ids,
+            "attention_mask": attention_mask,
+            "pixel_values": kwargs.get("pixel_values", None),
+            "past_key_values": past,
+        }