Upload model

modeling_ministral.py

@@ -417,15 +417,19 @@ class Ministral3Model(Ministral3PreTrainedModel):
         if position_ids is None:
             position_ids = cache_position.unsqueeze(0)
 
-        mask_function = create_causal_mask if self.config.sliding_window is None else create_sliding_window_causal_mask
-        causal_mask = mask_function(
-            config=self.config,
-            input_embeds=inputs_embeds,
-            attention_mask=attention_mask,
-            cache_position=cache_position,
-            past_key_values=past_key_values,
-            position_ids=position_ids,
-        )
+        if self.training:
+            causal_mask = None
+        
+        else:
+            mask_function = create_causal_mask if self.config.sliding_window is None else create_sliding_window_causal_mask
+            causal_mask = mask_function(
+                config=self.config,
+                input_embeds=inputs_embeds,
+                attention_mask=attention_mask,
+                cache_position=cache_position,
+                past_key_values=past_key_values,
+                position_ids=position_ids,
+            )
 
         hidden_states = inputs_embeds
         position_embeddings = self.rotary_emb(hidden_states, position_ids=position_ids)

modeling_ministral_dlm.py

@@ -518,7 +518,7 @@ class MinistralDiffEncoderModel(Ministral3PreTrainedModel, GenerationMixin):
 
         if labels is not None and self.config.dlm_paradigm != 'autoregressive':
             if masked_indices is not None:
-                #assert p_mask is not None
+                assert p_mask is not None
 
                 if loss_mask is not None:
                     masked_indices[loss_mask == 0] = 0