Update modeling_minicpm.py

modeling_minicpm.py

@@ -48,7 +48,7 @@ from transformers.utils import (
     replace_return_docstrings,
 )
 from transformers.utils.import_utils import is_torch_fx_available
-from .configuration_minicpm import MiniCPMConfig
+from .configuration_minicpm import MiniCPM3Config
 import re
 
 try:
@@ -69,7 +69,7 @@ if is_torch_fx_available():
 
 logger = logging.get_logger(__name__)
 
-_CONFIG_FOR_DOC = "MiniCPMConfig"
+_CONFIG_FOR_DOC = "MiniCPM3Config"
 
 
 def _get_unpad_data(attention_mask):
@@ -220,7 +220,6 @@ class MiniCPMLongRoPE(MiniCPMRotaryEmbedding):
         self.scaling_factor = math.sqrt(
                 1 + math.log(scale) /
                 math.log(self.original_max_position_embeddings))
-        print("using longrope!!!")
         super().__init__(dim, max_position_embeddings, base, device)
 
     def _set_cos_sin_cache(self, seq_len, device, dtype):
@@ -235,8 +234,6 @@ class MiniCPMLongRoPE(MiniCPMRotaryEmbedding):
             torch.outer(t, 1.0 / ext_factors).to(device=device),
             self.inv_freq.to(device=device).to(dtype)
         )
-        # print("??? rope freqs: ", freqs.shape)
-        # print("??? rope freqs: ", freqs)
         # Different from paper, but it uses a different permutation in order to obtain the same calculation
         emb = torch.cat((freqs, freqs), dim=-1)
         self.register_buffer("cos_cached", emb.cos().to(dtype) * self.scaling_factor, persistent=False)
@@ -334,7 +331,7 @@ def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
 class MiniCPMAttention(nn.Module):
     """Multi-headed attention from 'Attention Is All You Need' paper"""
 
-    def __init__(self, config: MiniCPMConfig, layer_idx: Optional[int] = None):
+    def __init__(self, config: MiniCPM3Config, layer_idx: Optional[int] = None):
         super().__init__()
         self.config = config
         self.layer_idx = layer_idx
@@ -601,17 +598,8 @@ class MiniCPMFlashAttention2(MiniCPMAttention):
             kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx)
 
         cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
-        # print("before rope q shape: ", q_pe.shape)
-        # print("before rope k shape: ", k_pe.shape)
-        # print("before rope q: ", q_pe.transpose(1, 2))
-        # print("before rope k: ", k_pe)
         q_pe, k_pe = apply_rotary_pos_emb(q_pe, k_pe, cos, sin, position_ids)
-        # print("after rope q shape: ", q_pe.shape)
-        # print("after rope k shape: ", k_pe.shape)
-        # print("after rope q: ", q_pe.transpose(1, 2))
-        # print("after rope k: ", k_pe)
-        # exit(1)
-
+        
         query_states = k_pe.new_empty(bsz, self.num_heads, q_len, self.q_head_dim)
         query_states[:, :, :, : self.qk_nope_head_dim] = q_nope
         query_states[:, :, :, self.qk_nope_head_dim :] = q_pe
@@ -637,12 +625,6 @@ class MiniCPMFlashAttention2(MiniCPMAttention):
 
         dropout_rate = self.attention_dropout if self.training else 0.0
 
-        # In PEFT, usually we cast the layer norms in float32 for training stability reasons
-        # therefore the input hidden states gets silently casted in float32. Hence, we need
-        # cast them back in the correct dtype just to be sure everything works as expected.
-        # This might slowdown training & inference so it is recommended to not cast the LayerNorms
-        # in fp32. (DeepseekV2RMSNorm handles it correctly)
-
         input_dtype = query_states.dtype
         if input_dtype == torch.float32:
             # Handle the case where the model is quantized
@@ -802,7 +784,7 @@ class MiniCPMSdpaAttention(MiniCPMAttention):
         if output_attentions:
             # TODO: Improve this warning with e.g. `model.config.attn_implementation = "manual"` once this is implemented.
             logger.warning_once(
-                "MiniCPMModel is using MiniCPMSdpaAttention, but `torch.nn.functional.scaled_dot_product_attention` does not support `output_attentions=True`. Falling back to the manual attention implementation, "
+                "MiniCPM3Model is using MiniCPMSdpaAttention, but `torch.nn.functional.scaled_dot_product_attention` does not support `output_attentions=True`. Falling back to the manual attention implementation, "
                 'but specifying the manual implementation will be required from Transformers version v5.0.0 onwards. This warning can be removed using the argument `attn_implementation="eager"` when loading the model.'
             )
             return super().forward(
@@ -902,7 +884,7 @@ MINICPM_ATTENTION_CLASSES = {
 
 
 class MiniCPMDecoderLayer(nn.Module):
-    def __init__(self, config: MiniCPMConfig, layer_idx: int):
+    def __init__(self, config: MiniCPM3Config, layer_idx: int):
         super().__init__()
         self.hidden_size = config.hidden_size
         self.self_attn = MINICPM_ATTENTION_CLASSES[config._attn_implementation](config=config, layer_idx=layer_idx)
@@ -986,7 +968,7 @@ MINICPM_START_DOCSTRING = r"""
     and behavior.
 
     Parameters:
-        config ([`MiniCPMConfig`]):
+        config ([`MiniCPM3Config`]):
             Model configuration class with all the parameters of the model. Initializing with a config file does not
             load the weights associated with the model, only the configuration. Check out the
             [`~PreTrainedModel.from_pretrained`] method to load the model weights.
@@ -998,7 +980,7 @@ MINICPM_START_DOCSTRING = r"""
     MINICPM_START_DOCSTRING,
 )
 class MiniCPM3PreTrainedModel(PreTrainedModel):
-    config_class = MiniCPMConfig
+    config_class = MiniCPM3Config
     base_model_prefix = "model"
     supports_gradient_checkpointing = True
     _no_split_modules = ["MiniCPMDecoderLayer"]
@@ -1098,10 +1080,10 @@ class MiniCPM3Model(MiniCPM3PreTrainedModel):
     Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`MiniCPMDecoderLayer`]
 
     Args:
-        config: MiniCPMConfig
+        config: MiniCPM3Config
     """
 
-    def __init__(self, config: MiniCPMConfig):
+    def __init__(self, config: MiniCPM3Config):
         super().__init__(config)
         self.padding_idx = config.pad_token_id
         self.vocab_size = config.vocab_size
@@ -1239,9 +1221,7 @@ class MiniCPM3Model(MiniCPM3PreTrainedModel):
                 all_self_attns += (layer_outputs[1],)
 
         hidden_states = self.norm(hidden_states)
-        # print("model outputs shape: ", hidden_states.shape)
-        # print("model outputs: ", hidden_states)
-
+        
         # add hidden states from the last decoder layer
         if output_hidden_states:
             all_hidden_states += (hidden_states,)
@@ -1264,7 +1244,7 @@ class MiniCPM3ForCausalLM(MiniCPM3PreTrainedModel):
 
     def __init__(self, config):
         super().__init__(config)
-        self.model = MiniCPMModel(config)
+        self.model = MiniCPM3Model(config)
         self.vocab_size = config.vocab_size
         self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
 
@@ -1461,15 +1441,11 @@ class MiniCPM3ForCausalLM(MiniCPM3PreTrainedModel):
                         "temperature": temperature, "logits_processor": logits_processor, **kwargs}
         
         history.append({"role": role, "content": query})
-        history_str = tokenizer.apply_chat_template(history, tokenize=False, add_generation_prompt=False)
+        history_str = tokenizer.apply_chat_template(history, tokenize=False, add_generation_prompt=True)
         inputs = tokenizer(history_str, return_tensors='pt').to(self.device)
         outputs = self.generate(**inputs, **gen_kwargs)
         outputs = outputs.tolist()[0][len(inputs["input_ids"][0]):-1]
         response = tokenizer.decode(outputs)
-        pattern = re.compile(r".*?(?=<AI>|<用户>)", re.DOTALL)
-        matches = pattern.findall(response)
-        if len(matches) > 0:
-            response = matches[0]
         history.append({"role": "assistant", "content": response})
         return response, history
 
@@ -1493,7 +1469,7 @@ class MiniCPM3ForSequenceClassification(MiniCPM3PreTrainedModel):
     def __init__(self, config):
         super().__init__(config)
         self.num_labels = config.num_labels
-        self.model = MiniCPMModel(config)
+        self.model = MiniCPM3Model(config)
         self.score = nn.Linear(config.hidden_size, self.num_labels, bias=False)
 
         # Initialize weights and apply final processing