Update modeling_quiet.py

modeling_quiet.py

@@ -22,6 +22,7 @@ import inspect
 import math
 import warnings
 from typing import List, Optional, Tuple, Union
+from dataclasses import dataclass
 
 import torch
 import torch.nn.functional as F
@@ -56,6 +57,31 @@ logger = logging.get_logger(__name__)
 
 _CONFIG_FOR_DOC = "QuietConfig"
 
+@dataclass
+class ModelOutput:
+    """
+    Base class for model's outputs, with potential hidden states and attentions.
+    """
+
+    def to_tuple(self):
+        """
+        Convert the output to a tuple.
+        """
+        return tuple(self[k] for k in self.keys())
+
+@dataclass
+class BaseModelOutput(ModelOutput):
+    last_hidden_state: torch.FloatTensor = None
+    hidden_states: Optional[Tuple[torch.FloatTensor]] = None
+    attentions: Optional[Tuple[torch.FloatTensor]] = None
+
+@dataclass
+class QuietModelOutputWithPast(BaseModelOutput):
+    last_hidden_state: torch.FloatTensor = None
+    past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None
+    hidden_states: Optional[Tuple[torch.FloatTensor]] = None
+    attentions: Optional[Tuple[torch.FloatTensor]] = None
+    logits: torch.FloatTensor = None
 
 # Copied from transformers.models.llama.modeling_llama._get_unpad_data
 def _get_unpad_data(attention_mask):
@@ -1097,7 +1123,7 @@ class QuietModel(QuietPreTrainedModel):
 
         if not return_dict:
             return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None)
-        return BaseModelOutputWithPast(
+        return QuietModelOutputWithPast(
             last_hidden_state=hidden_states,
             past_key_values=next_cache,
             hidden_states=all_hidden_states,
@@ -1216,7 +1242,7 @@ class QuietForCausalLM(QuietPreTrainedModel):
         )
 
         hidden_states = outputs.last_hidden_state
-        base_logits = outputs.logits  # Use the logits from the model output
+        base_logits = outputs.logits
 
         thought_ids, thought_embeddings = self.model._generate_thoughts(hidden_states, max_length=self.thought_length)
         thought_hidden_states = self.model(inputs_embeds=thought_embeddings).last_hidden_state