Sync modeling_auristream.py from TuKoResearch/AuriStream100M_40Pred_BigAudioDataset_500k

modeling_auristream.py

@@ -134,15 +134,29 @@ class AuriStream(PreTrainedModel):
             
             if return_dict:
                 if output_logits:
-                    model_output = CausalLMOutput(
-                        loss=loss,
-                        logits=all_logits,
-                    )
+                    if output_hidden_states:
+                        model_output = CausalLMOutput(
+                            loss=loss,
+                            logits=all_logits,
+                            hidden_states=all_hidden_states,
+                        )
+                    else:
+                        model_output = CausalLMOutput(
+                            loss=loss,
+                            logits=all_logits,
+                        )
                 else:
-                    model_output = CausalLMOutput(
-                        loss=loss,
-                        logits=logits,
-                    )
+                    if output_hidden_states:
+                        model_output = CausalLMOutput(
+                            loss=loss,
+                            logits=logits,
+                            hidden_states=all_hidden_states,
+                        )
+                    else:
+                        model_output = CausalLMOutput(
+                            loss=loss,
+                            logits=logits,
+                        )
                 return model_output
             
             return logits, loss
@@ -577,4 +591,4 @@ class RMSNorm(nn.Module):
         output = self._norm(x.float()).type_as(x)
         if self.weight is not None:
             return output * self.weight
-        return output
+        return output