Update modeling_quiet.py

modeling_quiet.py

@@ -57,31 +57,6 @@ 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):
@@ -1123,12 +1098,11 @@ 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 QuietModelOutputWithPast(
+        return BaseModelOutputWithPast(
             last_hidden_state=hidden_states,
             past_key_values=next_cache,
             hidden_states=all_hidden_states,
             attentions=all_self_attns,
-            logits=self.lm_head(hidden_states),
         )
 
 
@@ -1274,8 +1248,8 @@ class QuietForCausalLM(QuietPreTrainedModel):
             return ((loss,) + output) if loss is not None else output
 
         return CausalLMOutputWithPast(
-            loss=loss,
-            logits=mixed_logits,
+            loss=loss if loss is not None else None,
+            logits=(rm_logits if self.n_ahead > 1 else logits) if not self.output_logits_at_the_end else logits,
             past_key_values=outputs.past_key_values,
             hidden_states=outputs.hidden_states,
             attentions=outputs.attentions,