Update modeling_minicpm.py

modeling_minicpm.py

@@ -21,12 +21,16 @@
 import math
 import warnings
 from typing import List, Optional, Tuple, Union, Dict
-
+import os
+from tqdm import tqdm
 import torch
 import torch.nn.functional as F
 import torch.utils.checkpoint
 from torch import nn
 from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
+import numpy as np
+from copy import deepcopy
+from transformers import AutoTokenizer
 
 from transformers.activations import ACT2FN
 from transformers.cache_utils import Cache, DynamicCache
@@ -35,6 +39,7 @@ from transformers.modeling_attn_mask_utils import (
     _prepare_4d_attention_mask,
     _prepare_4d_causal_attention_mask,
     _prepare_4d_causal_attention_mask_for_sdpa,
+    _prepare_4d_attention_mask_for_sdpa,
 )
 from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast, SequenceClassifierOutputWithPast
 from transformers.modeling_utils import PreTrainedModel
@@ -320,9 +325,6 @@ class MiniCPMAttention(nn.Module):
         self.rope_theta = config.rope_theta
                 
         self.is_causal = config.is_causal
-        
-        logger.info(f"self.is_causal = {self.is_causal}")
-
 
         if (self.head_dim * self.num_heads) != self.hidden_size:
             raise ValueError(
@@ -979,6 +981,8 @@ class MiniCPMModel(MiniCPMPreTrainedModel):
         self.norm = MiniCPMRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
 
         self.gradient_checkpointing = False
+        self.is_causal = config.is_causal
+        self.adapt_mean_pooling = config.adapt_mean_pooling
         # Initialize weights and apply final processing
         self.post_init()
 
@@ -1000,6 +1004,7 @@ class MiniCPMModel(MiniCPMPreTrainedModel):
         output_attentions: Optional[bool] = None,
         output_hidden_states: Optional[bool] = None,
         return_dict: Optional[bool] = None,
+        adapt_mean_pooling: Optional[bool] = None,
     ) -> Union[Tuple, BaseModelOutputWithPast]:
         output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
         output_hidden_states = (
@@ -1044,24 +1049,35 @@ class MiniCPMModel(MiniCPMPreTrainedModel):
             inputs_embeds = self.embed_tokens(input_ids) * self.config.scale_emb
 
         _attention_mask = attention_mask
-        
         if self._use_flash_attention_2:
             # 2d mask is passed through the layers
             attention_mask = attention_mask if (attention_mask is not None and 0 in attention_mask) else None
         elif self._use_sdpa and not output_attentions:
             # output_attentions=True can not be supported when using SDPA, and we fall back on
             # the manual implementation that requires a 4D causal mask in all cases.
-            attention_mask = _prepare_4d_causal_attention_mask_for_sdpa(
-                attention_mask,
-                (batch_size, seq_length),
-                inputs_embeds,
-                past_key_values_length,
-            )
+            if self.is_causal:
+                attention_mask = _prepare_4d_causal_attention_mask_for_sdpa (
+                    attention_mask,
+                    (batch_size, seq_length),
+                    inputs_embeds,
+                    past_key_values_length,
+                )
+            else:
+                attention_mask = _prepare_4d_attention_mask_for_sdpa(
+                    attention_mask,
+                    inputs_embeds.dtype,
+                )
         else:
             # 4d mask is passed through the layers
-            attention_mask = _prepare_4d_causal_attention_mask(
-                attention_mask, (batch_size, seq_length), inputs_embeds, past_key_values_length
-            )
+            if self.is_causal:
+                attention_mask = _prepare_4d_causal_attention_mask (
+                    attention_mask, (batch_size, seq_length), inputs_embeds, past_key_values_length
+                )
+            else:
+                attention_mask = _prepare_4d_attention_mask(
+                    attention_mask,
+                    inputs_embeds.dtype,
+                )
 
         # embed positions
         hidden_states = inputs_embeds
@@ -1109,14 +1125,18 @@ class MiniCPMModel(MiniCPMPreTrainedModel):
         if output_hidden_states:
             all_hidden_states += (hidden_states,)
 
-        # gen weight before mean pooling
-        attention_mask_ = _attention_mask * _attention_mask.cumsum(dim=1)
-        s = hidden_states * attention_mask_.unsqueeze(-1).float()
-        d = attention_mask_.sum(dim=1, keepdim=True).unsqueeze(1).float() /_attention_mask.sum(dim=1, keepdim=True).unsqueeze(1).float() 
+        next_cache = None
         
-        hidden_states = s / d
+        # gen weight before mean pooling
+        if adapt_mean_pooling is None:
+            adapt_mean_pooling = self.adapt_mean_pooling
+        if adapt_mean_pooling:
+            attention_mask_ = _attention_mask * _attention_mask.cumsum(dim=1)
+            s = hidden_states * attention_mask_.unsqueeze(-1).float()
+            d = attention_mask_.sum(dim=1, keepdim=True).unsqueeze(1).float() /_attention_mask.sum(dim=1, keepdim=True).unsqueeze(1).float()
+
+            hidden_states = s / d
         
-        next_cache = None
         if use_cache:
             next_cache = next_decoder_cache.to_legacy_cache() if use_legacy_cache else next_decoder_cache
         if not return_dict:
@@ -1127,7 +1147,8 @@ class MiniCPMModel(MiniCPMPreTrainedModel):
             hidden_states=all_hidden_states,
             attentions=all_self_attns,
         )
-
+        
+        
 
 class MiniCPMForCausalLM(MiniCPMPreTrainedModel):
     _tied_weights_keys = ["lm_head.weight"]