Update modeling_auristream.py

modeling_auristream.py

@@ -72,7 +72,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
@@ -112,6 +112,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),
             )
@@ -123,14 +127,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