Added custom MinistralDiffOutputWithPast return type and skip_loss functionality

modeling_ministral_dlm.py

@@ -1,4 +1,5 @@
 import copy
+from dataclasses import dataclass
 from typing import Callable, Optional, Tuple, Union
 import random
 import os
@@ -10,6 +11,7 @@ import torch
 import torch.nn.functional as F
 from torch import nn
 from transformers.modeling_outputs import CausalLMOutputWithPast, BaseModelOutput
+from transformers.utils import ModelOutput
 
 from torch.nn.attention.flex_attention import flex_attention, create_block_mask
 
@@ -29,6 +31,17 @@ from .chat_utils import generate_with_prefix_cache_block_diff
 from .modeling_ministral import Ministral3Model, Ministral3PreTrainedModel, Ministral3Attention, apply_rotary_pos_emb, repeat_kv, _get_llama_4_attn_scale
 from .configuration_ministral_dlm import MinistralDLMConfig
 
+
+@dataclass
+class MinistralDiffOutputWithPast(ModelOutput):
+    loss: torch.FloatTensor | None = None
+    logits: torch.FloatTensor | None = None
+    causal_logits: torch.FloatTensor | None = None
+    past_key_values: Cache | None = None
+    hidden_states: tuple[torch.FloatTensor, ...] | None = None
+    attentions: tuple[torch.FloatTensor, ...] | None = None
+
+
 # @torch.compile(dynamic=True, mode="reduce-overhead")
 # @torch.compile(mode="default")
 # @torch.compile(fullgraph=True, mode="reduce-overhead", dynamic=False)
@@ -479,6 +492,7 @@ class MinistralDiffEncoderModel(Ministral3PreTrainedModel, GenerationMixin):
         loss_mask: Optional[torch.Tensor] = None,
         ce_loss_weight: float = 1.0,
         output_last_hidden_states_only: bool = False,
+        skip_loss: bool = False,
         **kwargs,
     ) -> CausalLMOutputWithPast:
 
@@ -565,7 +579,7 @@ class MinistralDiffEncoderModel(Ministral3PreTrainedModel, GenerationMixin):
             logits = logits[:, :input_ids_len]
 
         loss = None
-        if labels is not None:  
+        if labels is not None and not skip_loss:
             if self.config.dlm_paradigm == 'autoregressive':
                 shift_logits = logits[..., :-1, :].contiguous()
                 shift_labels = labels[..., 1:].contiguous()
@@ -702,9 +716,10 @@ class MinistralDiffEncoderModel(Ministral3PreTrainedModel, GenerationMixin):
                     else: 
                         loss = (loss, num_mask_tokens)
 
-        return CausalLMOutputWithPast(
+        return MinistralDiffOutputWithPast(
             loss=loss if not is_teacher else logits,
             logits=logits,
+            causal_logits=causal_logits,
             past_key_values=enc_out.past_key_values,
             hidden_states=None,
             attentions=None,