Update modeling_auristream.py

modeling_auristream.py

@@ -76,7 +76,7 @@ class AuriStream(PreTrainedModel):
         elif isinstance(module, nn.Embedding):
             torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
    
-    def forward(self, seq, tgt=None, output_hidden_states=False, return_dict=False, up_until_layer=None):
+    def forward(self, seq, tgt=None, output_logits=False, output_hidden_states=False, return_dict=False, up_until_layer=None):
         """
         Input: coch: torch.Tensor of shape (b, t)
                tgt_coch: torch.Tensor of shape (b, t) or None
@@ -106,13 +106,12 @@ class AuriStream(PreTrainedModel):
             x = block(x)
             if self.dwa is not None:
                 x = self.dwa(x)
-            
 
         # append the last hidden state if we did not exit early
         if up_until_layer is None or block_idx == len(self.transformer.h) - 1:
             all_hidden_states.append(x)
 
-        if output_hidden_states:
+        if output_hidden_states and not output_logits:
             model_output = BaseModelOutput(
                 last_hidden_state=x,
                 hidden_states=all_hidden_states,
@@ -123,6 +122,10 @@ class AuriStream(PreTrainedModel):
         logits = self.coch_head(x)
 
         if tgt is not None:
+
+            if output_logits:
+                all_logits = [logits]
+            
             loss = F.cross_entropy(
                 logits.reshape(-1, self.config.vocab_size), tgt.reshape(-1),
             )
@@ -134,14 +137,22 @@ class AuriStream(PreTrainedModel):
                     loss += F.cross_entropy(
                         future_logits.reshape(-1, self.config.vocab_size), tgt[:, (i+1):].reshape(-1),
                     )
+                    if output_logits:
+                        all_logits.append(future_logits)
                 # divide loss by number of future heads
                 loss = loss / (len(self.future_heads) + 1)
-
+            
             if return_dict:
-                model_output = CausalLMOutput(
-                    loss=loss,
-                    logits=logits,
-                )
+                if output_logits:
+                    model_output = CausalLMOutput(
+                        loss=loss,
+                        logits=all_logits,
+                    )
+                else:
+                    model_output = CausalLMOutput(
+                        loss=loss,
+                        logits=logits,
+                    )
                 return model_output
             
             return logits, loss