Sync modeling_auristream.py from TuKoResearch/AuriStream200M_100Pred_librilight_200k

modeling_auristream.py

@@ -111,30 +111,35 @@ class AuriStream(PreTrainedModel):
         x = self.transformer.ln_f(x)
         logits = self.coch_head(x)
 
-        if tgt is not None:
 
-            if output_logits:
-                all_logits = [logits]
-            
+        if output_logits:
+            all_logits = [logits]
+
+        if tgt is not None:
             loss = F.cross_entropy(
                 logits.reshape(-1, self.config.vocab_size), tgt.reshape(-1),
             )
 
-            # If we have more than one future head, compute the loss for each head
-            if self.future_heads is not None:
-                for i, head in enumerate(self.future_heads):
-                    future_logits = head(x[:, :-(i+1)])
+        # If we have more than one future head, compute the loss for each head
+        if self.future_heads is not None:
+            for i, head in enumerate(self.future_heads):
+                future_logits = head(x[:, :-(i+1)])
+
+                if tgt is not None:
                     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)
+                if output_logits:
+                    all_logits.append(future_logits)
+
+            if tgt is not None:
                 # divide loss by number of future heads
                 loss = loss / (len(self.future_heads) + 1)
             
-            if return_dict:
-                if output_logits:
-                    if output_hidden_states:
+        if return_dict:
+            if output_logits:
+                if output_hidden_states:
+                    if tgt is not None:
                         model_output = CausalLMOutput(
                             loss=loss,
                             logits=all_logits,
@@ -142,23 +147,45 @@ class AuriStream(PreTrainedModel):
                         )
                     else:
                         model_output = CausalLMOutput(
-                            loss=loss,
                             logits=all_logits,
+                            hidden_states=all_hidden_states,
                         )
                 else:
-                    if output_hidden_states:
+                    if tgt is not None:
+                        model_output = CausalLMOutput(
+                            loss=loss,
+                            logits=all_logits,
+                        )
+                    else:
+                        model_output = CausalLMOutput(
+                            logits=all_logits,
+                        )
+            else:
+                if output_hidden_states:
+                    if tgt is not None:
                         model_output = CausalLMOutput(
                             loss=loss,
                             logits=logits,
                             hidden_states=all_hidden_states,
                         )
                     else:
+                        model_output = CausalLMOutput(
+                            logits=logits,
+                            hidden_states=all_hidden_states,
+                        )
+                else:
+                    if tgt is not None:
                         model_output = CausalLMOutput(
                             loss=loss,
                             logits=logits,
                         )
-                return model_output
-            
+                    else:
+                        model_output = CausalLMOutput(
+                            logits=logits,
+                        )
+            return model_output
+
+        if tgt is not None:
             return logits, loss
 
         return logits, None