Upload model

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)
@@ -81,13 +94,11 @@ class MinistralFlexAttention(Ministral3Attention):
             - **Block Diagonal Mask (M_BD)**: Self-attention within noised blocks
             - **Offset Block Causal Mask (M_OBC)**: Cross-attention for conditional context
             - **Block Causal Mask (M_BC)**: Attention to update x0
-
             Args:
                 b, h: Batch and head indices (ignored for mask logic).
                 q_idx, kv_idx: Query and Key indices.
                 seq_len: Total sequence length.
                 block_size: Defines the block structure.
-
             Returns:
                 A boolean attention mask.
             """
@@ -378,16 +389,13 @@ class MinistralDiffEncoderModel(Ministral3PreTrainedModel, GenerationMixin):
     ):
         """
         Two-stage corruption with optional per-block sampling.
-
         • Stage 1:  m ~ U(eps, 1)   →   k = round(m · len)  (exact budget).
         • Stage 2:  sample exactly k positions with weights
                     w_i(m) = exp[ λ · (1−m) · i ]   (late-heavy when m→0,
                                                      uniform when m→1).
-
           If `block_size` is given, the procedure is run *independently*
           inside each contiguous block of that length (last block may be shorter).
           When block_size is provided, m is sampled per-block and p_mask is per-block.
-
         Args
         ----
         input_ids : (B, L)  LongTensor
@@ -479,6 +487,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:
 
@@ -555,7 +564,8 @@ class MinistralDiffEncoderModel(Ministral3PreTrainedModel, GenerationMixin):
             return BaseModelOutput(last_hidden_state=enc_out.last_hidden_state)
 
         logits = self.diffusion_head(enc_out.last_hidden_state)  # (batch, len_B, vocab)
-
+        causal_logits = None
+        
         if labels is not None and self.config.dlm_paradigm in ['block_diff', 'sbd_block_diff']:
             if self.config.dlm_paradigm == 'sbd_block_diff':
                 causal_logits = logits[:, input_ids_len:]
@@ -565,7 +575,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 +712,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,
@@ -729,5 +740,4 @@ class MinistralDiffEncoderModel(Ministral3PreTrainedModel, GenerationMixin):
 
         return out_ids, nfe
 
-__all__ = ["MinistralDiffEncoderModel", "MinistralFlexAttention"]
-
+__all__ = ["MinistralDiffEncoderModel", "MinistralFlexAttention"]