Edcastro/MiniCPM-1B-sft-bf16-edcastr_JavaScript-v1
029
1# coding=utf-82# Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved.3#4# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX5# and OPT implementations in this library. It has been modified from its6# original forms to accommodate minor architectural differences compared7# to GPT-NeoX and OPT used by the Meta AI team that trained the model.8#9# Licensed under the Apache License, Version 2.0 (the "License");10# you may not use this file except in compliance with the License.11# You may obtain a copy of the License at12#13# http://www.apache.org/licenses/LICENSE-2.014#15# Unless required by applicable law or agreed to in writing, software16# distributed under the License is distributed on an "AS IS" BASIS,17# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.18# See the License for the specific language governing permissions and19# limitations under the License.20""" PyTorch MiniCPM model."""21import math22import warnings23from typing import List, Optional, Tuple, Union, Dict24 25import torch26import torch.nn.functional as F27import torch.utils.checkpoint28from torch import nn29from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss30 31from transformers.activations import ACT2FN32from transformers.cache_utils import Cache, DynamicCache33from transformers.modeling_attn_mask_utils import (34 AttentionMaskConverter,35 _prepare_4d_attention_mask,36 _prepare_4d_causal_attention_mask,37 _prepare_4d_causal_attention_mask_for_sdpa,38)39from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast, SequenceClassifierOutputWithPast40from transformers.modeling_utils import PreTrainedModel41from transformers.pytorch_utils import ALL_LAYERNORM_LAYERS, is_torch_greater_or_equal_than_1_1342from transformers.utils import (43 add_start_docstrings,44 add_start_docstrings_to_model_forward,45 is_flash_attn_2_available,46 is_flash_attn_greater_or_equal_2_10,47 logging,48 replace_return_docstrings,49)50from transformers.utils.import_utils import is_torch_fx_available51from .configuration_minicpm import MiniCPMConfig52import re53 54try:55 from flash_attn import flash_attn_func, flash_attn_varlen_func56 from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input # noqa57except:58 pass59 60 61# This makes `_prepare_4d_causal_attention_mask` a leaf function in the FX graph.62# It means that the function will not be traced through and simply appear as a node in the graph.63if is_torch_fx_available():64 if not is_torch_greater_or_equal_than_1_13:65 import torch.fx66 67 _prepare_4d_causal_attention_mask = torch.fx.wrap(_prepare_4d_causal_attention_mask)68 69 70logger = logging.get_logger(__name__)71 72_CONFIG_FOR_DOC = "MiniCPMConfig"73 74 75def _get_unpad_data(attention_mask):76 seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32)77 indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()78 max_seqlen_in_batch = seqlens_in_batch.max().item()79 cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.torch.int32), (1, 0))80 return (81 indices,82 cu_seqlens,83 max_seqlen_in_batch,84 )85 86 87def _expand_mask(mask: torch.Tensor, dtype: torch.dtype, tgt_len: Optional[int] = None):88 warnings.warn(89 "Calling `transformers.models.minicpm.modeling_minicpm._prepare_4d_attention_mask` is deprecated and will be removed in v4.37. Use `transformers.modeling_attn_mask_utils._prepare_4d_attention_mask"90 )91 return _prepare_4d_attention_mask(mask=mask, dtype=dtype, tgt_len=tgt_len)92 93 94def _make_causal_mask(95 input_ids_shape: torch.Size, dtype: torch.dtype, device: torch.device, past_key_values_length: int = 096):97 warnings.warn(98 "Calling `transformers.models.minicpm.modeling_minicpm._make_causal_mask` is deprecated and will be removed in v4.37. Use `transformers.models.minicpm.modeling_minicpm.AttentionMaskConverter._make_causal_mask"99 )100 return AttentionMaskConverter._make_causal_mask(101 input_ids_shape=input_ids_shape, dtype=dtype, device=device, past_key_values_length=past_key_values_length102 )103 104# @torch.jit.script # type: ignore105def rms_layernorm(hidden: torch.Tensor, weight: torch.Tensor, eps: float):106 old_dtype = hidden.dtype107 variance = hidden.to(torch.float32).pow(2).mean(dim=-1, keepdim=True)108 hidden = (hidden * torch.rsqrt(variance + eps)).to(old_dtype)109 return hidden * weight110 111 112class MiniCPMRMSNorm(nn.Module):113 def __init__(self, hidden_size, eps=1e-6):114 """115 MiniCPMRMSNorm is equivalent to T5LayerNorm116 """117 super().__init__()118 self.weight = nn.Parameter(torch.ones(hidden_size))119 self.variance_epsilon = eps120 121 def forward(self, hidden_states):122 return rms_layernorm(hidden_states, self.weight, self.variance_epsilon)123 124 125ALL_LAYERNORM_LAYERS.append(MiniCPMRMSNorm)126 127 128class MiniCPMRotaryEmbedding(nn.Module):129 def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None):130 super().__init__()131 132 self.dim = dim133 self.max_position_embeddings = max_position_embeddings134 self.base = base135 inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2).float().to(device) / self.dim))136 self.register_buffer("inv_freq", inv_freq, persistent=False)137 138 # Build here to make `torch.jit.trace` work.139 self._set_cos_sin_cache(140 # seq_len=max_position_embeddings, device=self.inv_freq.device, dtype=torch.get_default_dtype()141 seq_len=max_position_embeddings, device=self.inv_freq.device, dtype=torch.float32142 )143 144 def _set_cos_sin_cache(self, seq_len, device, dtype):145 self.max_seq_len_cached = seq_len146 t = torch.arange(self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype)147 freqs = torch.outer(t, self.inv_freq)148 # Different from paper, but it uses a different permutation in order to obtain the same calculation149 emb = torch.cat((freqs, freqs), dim=-1)150 151 self.register_buffer("cos_cached", emb.cos().to(dtype), persistent=False)152 self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False)153 154 def forward(self, x, seq_len=None):155 # x: [bs, num_attention_heads, seq_len, head_size]156 if seq_len > self.max_seq_len_cached:157 self._set_cos_sin_cache(seq_len=seq_len, device=x.device, dtype=x.dtype)158 159 return (160 self.cos_cached[:seq_len].to(dtype=x.dtype),161 self.sin_cached[:seq_len].to(dtype=x.dtype),162 )163 164 165class MiniCPMLinearScalingRotaryEmbedding(MiniCPMRotaryEmbedding):166 """MiniCPMRotaryEmbedding extended with linear scaling. Credits to the Reddit user /u/kaiokendev"""167 168 def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None, scaling_factor=1.0):169 self.scaling_factor = scaling_factor170 super().__init__(dim, max_position_embeddings, base, device)171 172 def _set_cos_sin_cache(self, seq_len, device, dtype):173 self.max_seq_len_cached = seq_len174 t = torch.arange(self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype)175 t = t / self.scaling_factor176 177 freqs = torch.outer(t, self.inv_freq)178 # Different from paper, but it uses a different permutation in order to obtain the same calculation179 emb = torch.cat((freqs, freqs), dim=-1)180 self.register_buffer("cos_cached", emb.cos().to(dtype), persistent=False)181 self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False)182 183 184class MiniCPMDynamicNTKScalingRotaryEmbedding(MiniCPMRotaryEmbedding):185 """MiniCPMRotaryEmbedding extended with Dynamic NTK scaling. Credits to the Reddit users /u/bloc97 and /u/emozilla"""186 187 def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None, scaling_factor=1.0):188 self.scaling_factor = scaling_factor189 super().__init__(dim, max_position_embeddings, base, device)190 191 def _set_cos_sin_cache(self, seq_len, device, dtype):192 self.max_seq_len_cached = seq_len193 194 if seq_len > self.max_position_embeddings:195 base = self.base * (196 (self.scaling_factor * seq_len / self.max_position_embeddings) - (self.scaling_factor - 1)197 ) ** (self.dim / (self.dim - 2))198 inv_freq = 1.0 / (base ** (torch.arange(0, self.dim, 2).float().to(device) / self.dim))199 self.register_buffer("inv_freq", inv_freq, persistent=False)200 201 t = torch.arange(self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype)202 203 freqs = torch.outer(t, self.inv_freq)204 # Different from paper, but it uses a different permutation in order to obtain the same calculation205 emb = torch.cat((freqs, freqs), dim=-1)206 207 self.register_buffer("cos_cached", emb.cos().to(dtype), persistent=False)208 self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False)209 210 211def rotate_half(x):212 """Rotates half the hidden dims of the input."""213 x1 = x[..., : x.shape[-1] // 2]214 x2 = x[..., x.shape[-1] // 2 :]215 return torch.cat((-x2, x1), dim=-1)216 217 218def apply_rotary_pos_emb(q, k, cos, sin, position_ids, unsqueeze_dim=1):219 """Applies Rotary Position Embedding to the query and key tensors.220 221 Args:222 q (`torch.Tensor`): The query tensor.223 k (`torch.Tensor`): The key tensor.224 cos (`torch.Tensor`): The cosine part of the rotary embedding.225 sin (`torch.Tensor`): The sine part of the rotary embedding.226 position_ids (`torch.Tensor`):227 The position indices of the tokens corresponding to the query and key tensors. For example, this can be228 used to pass offsetted position ids when working with a KV-cache.229 unsqueeze_dim (`int`, *optional*, defaults to 1):230 The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and231 sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note232 that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and233 k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes234 cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have235 the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.236 Returns:237 `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.238 """239 # cos = cos[position_ids].unsqueeze(unsqueeze_dim)240 # sin = sin[position_ids].unsqueeze(unsqueeze_dim)241 # q_embed = (q * cos) + (rotate_half(q) * sin)242 # k_embed = (k * cos) + (rotate_half(k) * sin)243 orig_dtype = k.dtype244 cos = cos[position_ids].unsqueeze(unsqueeze_dim) # [bs, 1, seq_len, dim]245 sin = sin[position_ids].unsqueeze(unsqueeze_dim) # [bs, 1, seq_len, dim]246 q_fp32 = q.to(dtype=torch.float32, device=q.device)247 k_fp32 = k.to(dtype=torch.float32, device=k.device)248 q_embed = (q_fp32 * cos) + (rotate_half(q_fp32) * sin)249 k_embed = (k_fp32 * cos) + (rotate_half(k_fp32) * sin)250 return q_embed.to(dtype=orig_dtype), k_embed.to(dtype=orig_dtype)251 252class MiniCPMMLP(nn.Module):253 def __init__(self, config):254 super().__init__()255 self.config = config256 self.hidden_size = config.hidden_size257 self.intermediate_size = config.intermediate_size258 self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)259 self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)260 self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)261 self.act_fn = ACT2FN[config.hidden_act]262 263 def forward(self, x):264 if self.config.pretraining_tp > 1:265 slice = self.intermediate_size // self.config.pretraining_tp266 gate_proj_slices = self.gate_proj.weight.split(slice, dim=0)267 up_proj_slices = self.up_proj.weight.split(slice, dim=0)268 down_proj_slices = self.down_proj.weight.split(slice, dim=1)269 270 gate_proj = torch.cat(271 [F.linear(x, gate_proj_slices[i]) for i in range(self.config.pretraining_tp)], dim=-1272 )273 up_proj = torch.cat([F.linear(x, up_proj_slices[i]) for i in range(self.config.pretraining_tp)], dim=-1)274 275 intermediate_states = (self.act_fn(gate_proj) * up_proj).split(slice, dim=2)276 down_proj = [277 F.linear(intermediate_states[i], down_proj_slices[i]) for i in range(self.config.pretraining_tp)278 ]279 down_proj = sum(down_proj)280 else:281 down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))282 283 return down_proj284 285 286def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:287 """288 This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,289 num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)290 """291 batch, num_key_value_heads, slen, head_dim = hidden_states.shape292 if n_rep == 1:293 return hidden_states294 hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)295 return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)296 297 298 299class MiniCPMAttention(nn.Module):300 """Multi-headed attention from 'Attention Is All You Need' paper"""301 302 def __init__(self, config: MiniCPMConfig, layer_idx: Optional[int] = None):303 super().__init__()304 self.config = config305 self.layer_idx = layer_idx306 if layer_idx is None:307 logger.warning_once(308 f"Instantiating {self.__class__.__name__} without passing `layer_idx` is not recommended and will "309 "to errors during the forward call, if caching is used. Please make sure to provide a `layer_idx` "310 "when creating this class."311 )312 313 self.attention_dropout = config.attention_dropout314 self.hidden_size = config.hidden_size315 self.num_heads = config.num_attention_heads316 self.head_dim = self.hidden_size // self.num_heads317 self.num_key_value_heads = config.num_key_value_heads318 self.num_key_value_groups = self.num_heads // self.num_key_value_heads319 self.max_position_embeddings = config.max_position_embeddings320 self.rope_theta = config.rope_theta321 self.is_causal = True322 323 if (self.head_dim * self.num_heads) != self.hidden_size:324 raise ValueError(325 f"hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size}"326 f" and `num_heads`: {self.num_heads})."327 )328 329 self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=config.attention_bias)330 self.k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=config.attention_bias)331 self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=config.attention_bias)332 self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=config.attention_bias)333 self._init_rope()334 335 def _init_rope(self):336 if self.config.rope_scaling is None:337 self.rotary_emb = MiniCPMRotaryEmbedding(338 self.head_dim,339 max_position_embeddings=self.max_position_embeddings,340 base=self.rope_theta,341 )342 else:343 scaling_type = self.config.rope_scaling["type"]344 scaling_factor = self.config.rope_scaling["factor"]345 if scaling_type == "linear":346 self.rotary_emb = MiniCPMLinearScalingRotaryEmbedding(347 self.head_dim,348 max_position_embeddings=self.max_position_embeddings,349 scaling_factor=scaling_factor,350 base=self.rope_theta,351 )352 elif scaling_type == "dynamic":353 self.rotary_emb = MiniCPMDynamicNTKScalingRotaryEmbedding(354 self.head_dim,355 max_position_embeddings=self.max_position_embeddings,356 scaling_factor=scaling_factor,357 base=self.rope_theta,358 )359 else:360 raise ValueError(f"Unknown RoPE scaling type {scaling_type}")361 362 def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int):363 return tensor.view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2).contiguous()364 365 def forward(366 self,367 hidden_states: torch.Tensor,368 attention_mask: Optional[torch.Tensor] = None,369 position_ids: Optional[torch.LongTensor] = None,370 past_key_value: Optional[Cache] = None,371 output_attentions: bool = False,372 use_cache: bool = False,373 **kwargs,374 ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:375 if "padding_mask" in kwargs:376 warnings.warn(377 "Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use `attention_mask` instead.`"378 )379 380 bsz, q_len, _ = hidden_states.size()381 382 if self.config.pretraining_tp > 1:383 key_value_slicing = (self.num_key_value_heads * self.head_dim) // self.config.pretraining_tp384 query_slices = self.q_proj.weight.split(385 (self.num_heads * self.head_dim) // self.config.pretraining_tp, dim=0386 )387 key_slices = self.k_proj.weight.split(key_value_slicing, dim=0)388 value_slices = self.v_proj.weight.split(key_value_slicing, dim=0)389 390 query_states = [F.linear(hidden_states, query_slices[i]) for i in range(self.config.pretraining_tp)]391 query_states = torch.cat(query_states, dim=-1)392 393 key_states = [F.linear(hidden_states, key_slices[i]) for i in range(self.config.pretraining_tp)]394 key_states = torch.cat(key_states, dim=-1)395 396 value_states = [F.linear(hidden_states, value_slices[i]) for i in range(self.config.pretraining_tp)]397 value_states = torch.cat(value_states, dim=-1)398 399 else:400 query_states = self.q_proj(hidden_states)401 key_states = self.k_proj(hidden_states)402 value_states = self.v_proj(hidden_states)403 404 query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)405 key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)406 value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)407 408 kv_seq_len = key_states.shape[-2]409 if past_key_value is not None:410 if self.layer_idx is None:411 raise ValueError(412 f"The cache structure has changed since version v4.36. If you are using {self.__class__.__name__} "413 "for auto-regressive decoding with k/v caching, please make sure to initialize the attention class "414 "with a layer index."415 )416 kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx)417 cos, sin = self.rotary_emb(value_states.to(torch.float32), seq_len=kv_seq_len)418 419 query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)420 421 if past_key_value is not None:422 cache_kwargs = {"sin": sin, "cos": cos} # Specific to RoPE models423 key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)424 425 key_states = repeat_kv(key_states, self.num_key_value_groups)426 value_states = repeat_kv(value_states, self.num_key_value_groups)427 428 attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim)429 if attn_weights.size() != (bsz, self.num_heads, q_len, kv_seq_len):430 raise ValueError(431 f"Attention weights should be of size {(bsz, self.num_heads, q_len, kv_seq_len)}, but is"432 f" {attn_weights.size()}"433 )434 435 if attention_mask is not None:436 if attention_mask.size() != (bsz, 1, q_len, kv_seq_len):437 raise ValueError(438 f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}"439 )440 attn_weights = attn_weights + attention_mask441 442 # upcast attention to fp32443 attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype)444 attn_weights = nn.functional.dropout(attn_weights, p=self.attention_dropout, training=self.training)445 attn_output = torch.matmul(attn_weights, value_states)446 447 if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim):448 raise ValueError(449 f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.head_dim)}, but is"450 f" {attn_output.size()}"451 )452 453 attn_output = attn_output.transpose(1, 2).contiguous()454 455 attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)456 457 if self.config.pretraining_tp > 1:458 attn_output = attn_output.split(self.hidden_size // self.config.pretraining_tp, dim=2)459 o_proj_slices = self.o_proj.weight.split(self.hidden_size // self.config.pretraining_tp, dim=1)460 attn_output = sum([F.linear(attn_output[i], o_proj_slices[i]) for i in range(self.config.pretraining_tp)])461 else:462 attn_output = self.o_proj(attn_output)463 464 if not output_attentions:465 attn_weights = None466 467 return attn_output, attn_weights, past_key_value468 469 470class MiniCPMFlashAttention2(MiniCPMAttention):471 """472 MiniCPM flash attention module. This module inherits from `MiniCPMAttention` as the weights of the module stays473 untouched. The only required change would be on the forward pass where it needs to correctly call the public API of474 flash attention and deal with padding tokens in case the input contains any of them.475 """476 477 def __init__(self, *args, **kwargs):478 super().__init__(*args, **kwargs)479 480 # TODO: Should be removed once Flash Attention for RoCm is bumped to 2.1.481 # flash_attn<2.1 generates top-left aligned causal mask, while what is needed here is bottom-right alignement, that was made default for flash_attn>=2.1. This attribute is used to handle this difference. Reference: https://github.com/Dao-AILab/flash-attention/releases/tag/v2.1.0.482 # Beware that with flash_attn<2.1, using q_seqlen != k_seqlen (except for the case q_seqlen == 1) produces a wrong mask (top-left).483 self._flash_attn_uses_top_left_mask = not is_flash_attn_greater_or_equal_2_10()484 485 def forward(486 self,487 hidden_states: torch.Tensor,488 attention_mask: Optional[torch.LongTensor] = None,489 position_ids: Optional[torch.LongTensor] = None,490 past_key_value: Optional[Cache] = None,491 output_attentions: bool = False,492 use_cache: bool = False,493 **kwargs,494 ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:495 # MiniCPMFlashAttention2 attention does not support output_attentions496 if "padding_mask" in kwargs:497 warnings.warn(498 "Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use `attention_mask` instead.`"499 )500 501 # overwrite attention_mask with padding_mask502 attention_mask = kwargs.pop("padding_mask")503 504 output_attentions = False505 506 bsz, q_len, _ = hidden_states.size()507 508 query_states = self.q_proj(hidden_states)509 key_states = self.k_proj(hidden_states)510 value_states = self.v_proj(hidden_states)511 512 # Flash attention requires the input to have the shape513 # batch_size x seq_length x head_dim x hidden_dim514 # therefore we just need to keep the original shape515 query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)516 key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)517 value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)518 519 kv_seq_len = key_states.shape[-2]520 if past_key_value is not None:521 kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx)522 cos, sin = self.rotary_emb(value_states.to(torch.float32), seq_len=kv_seq_len)523 query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)524 525 if past_key_value is not None:526 cache_kwargs = {"sin": sin, "cos": cos} # Specific to RoPE models527 key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)528 529 # TODO: These transpose are quite inefficient but Flash Attention requires the layout [batch_size, sequence_length, num_heads, head_dim]. We would need to refactor the KV cache530 # to be able to avoid many of these transpose/reshape/view.531 query_states = query_states.transpose(1, 2)532 key_states = key_states.transpose(1, 2)533 value_states = value_states.transpose(1, 2)534 535 dropout_rate = self.attention_dropout if self.training else 0.0536 537 # In PEFT, usually we cast the layer norms in float32 for training stability reasons538 # therefore the input hidden states gets silently casted in float32. Hence, we need539 # cast them back in the correct dtype just to be sure everything works as expected.540 # This might slowdown training & inference so it is recommended to not cast the LayerNorms541 # in fp32. (MiniCPMRMSNorm handles it correctly)542 543 input_dtype = query_states.dtype544 if input_dtype == torch.float32:545 # Handle the case where the model is quantized546 if hasattr(self.config, "_pre_quantization_dtype"):547 target_dtype = self.config._pre_quantization_dtype548 else:549 target_dtype = self.q_proj.weight.dtype550 551 logger.warning_once(552 f"The input hidden states seems to be silently casted in float32, this might be related to"553 f" the fact you have upcasted embedding or layer norm layers in float32. We will cast back the input in"554 f" {target_dtype}."555 )556 557 query_states = query_states.to(target_dtype)558 key_states = key_states.to(target_dtype)559 value_states = value_states.to(target_dtype)560 561 attn_output = self._flash_attention_forward(562 query_states, key_states, value_states, attention_mask, q_len, dropout=dropout_rate563 )564 565 attn_output = attn_output.reshape(bsz, q_len, self.hidden_size).contiguous()566 attn_output = self.o_proj(attn_output)567 568 if not output_attentions:569 attn_weights = None570 571 return attn_output, attn_weights, past_key_value572 573 def _flash_attention_forward(574 self, query_states, key_states, value_states, attention_mask, query_length, dropout=0.0, softmax_scale=None575 ):576 """577 Calls the forward method of Flash Attention - if the input hidden states contain at least one padding token578 first unpad the input, then computes the attention scores and pad the final attention scores.579 580 Args:581 query_states (`torch.Tensor`):582 Input query states to be passed to Flash Attention API583 key_states (`torch.Tensor`):584 Input key states to be passed to Flash Attention API585 value_states (`torch.Tensor`):586 Input value states to be passed to Flash Attention API587 attention_mask (`torch.Tensor`):588 The padding mask - corresponds to a tensor of size `(batch_size, seq_len)` where 0 stands for the589 position of padding tokens and 1 for the position of non-padding tokens.590 dropout (`int`, *optional*):591 Attention dropout592 softmax_scale (`float`, *optional*):593 The scaling of QK^T before applying softmax. Default to 1 / sqrt(head_dim)594 """595 if not self._flash_attn_uses_top_left_mask:596 causal = self.is_causal597 else:598 # TODO: Remove the `query_length != 1` check once Flash Attention for RoCm is bumped to 2.1. For details, please see the comment in MiniCPMFlashAttention2 __init__.599 causal = self.is_causal and query_length != 1600 # Contains at least one padding token in the sequence601 if attention_mask is not None:602 batch_size = query_states.shape[0]603 query_states, key_states, value_states, indices_q, cu_seq_lens, max_seq_lens = self._upad_input(604 query_states, key_states, value_states, attention_mask, query_length605 )606 607 cu_seqlens_q, cu_seqlens_k = cu_seq_lens608 max_seqlen_in_batch_q, max_seqlen_in_batch_k = max_seq_lens609 attn_output_unpad = flash_attn_varlen_func(610 query_states,611 key_states,612 value_states,613 cu_seqlens_q=cu_seqlens_q,614 cu_seqlens_k=cu_seqlens_k,615 max_seqlen_q=max_seqlen_in_batch_q,616 max_seqlen_k=max_seqlen_in_batch_k,617 dropout_p=dropout,618 softmax_scale=softmax_scale,619 causal=causal,620 )621 622 attn_output = pad_input(attn_output_unpad, indices_q, batch_size, query_length)623 else:624 attn_output = flash_attn_func(625 query_states, key_states, value_states, dropout, softmax_scale=softmax_scale, causal=causal626 )627 628 return attn_output629 630 def _upad_input(self, query_layer, key_layer, value_layer, attention_mask, query_length):631 indices_k, cu_seqlens_k, max_seqlen_in_batch_k = _get_unpad_data(attention_mask)632 batch_size, kv_seq_len, num_key_value_heads, head_dim = key_layer.shape633 634 key_layer = index_first_axis(635 key_layer.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim), indices_k636 )637 value_layer = index_first_axis(638 value_layer.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim), indices_k639 )640 if query_length == kv_seq_len:641 query_layer = index_first_axis(642 query_layer.reshape(batch_size * kv_seq_len, self.num_heads, head_dim), indices_k643 )644 cu_seqlens_q = cu_seqlens_k645 max_seqlen_in_batch_q = max_seqlen_in_batch_k646 indices_q = indices_k647 elif query_length == 1:648 max_seqlen_in_batch_q = 1649 cu_seqlens_q = torch.arange(650 batch_size + 1, dtype=torch.int32, device=query_layer.device651 ) # There is a memcpy here, that is very bad.652 indices_q = cu_seqlens_q[:-1]653 query_layer = query_layer.squeeze(1)654 else:655 # The -q_len: slice assumes left padding.656 attention_mask = attention_mask[:, -query_length:]657 query_layer, indices_q, cu_seqlens_q, max_seqlen_in_batch_q = unpad_input(query_layer, attention_mask)658 659 return (660 query_layer,661 key_layer,662 value_layer,663 indices_q,664 (cu_seqlens_q, cu_seqlens_k),665 (max_seqlen_in_batch_q, max_seqlen_in_batch_k),666 )667 668 669class MiniCPMSdpaAttention(MiniCPMAttention):670 """671 MiniCPM attention module using torch.nn.functional.scaled_dot_product_attention. This module inherits from672 `MiniCPMAttention` as the weights of the module stays untouched. The only changes are on the forward pass to adapt to673 SDPA API.674 """675 676 # Adapted from MiniCPMAttention.forward677 def forward(678 self,679 hidden_states: torch.Tensor,680 attention_mask: Optional[torch.Tensor] = None,681 position_ids: Optional[torch.LongTensor] = None,682 past_key_value: Optional[Cache] = None,683 output_attentions: bool = False,684 use_cache: bool = False,685 ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:686 if output_attentions:687 # TODO: Improve this warning with e.g. `model.config.attn_implementation = "manual"` once this is implemented.688 logger.warning_once(689 "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, "690 '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.'691 )692 return super().forward(693 hidden_states=hidden_states,694 attention_mask=attention_mask,695 position_ids=position_ids,696 past_key_value=past_key_value,697 output_attentions=output_attentions,698 use_cache=use_cache,699 )700 701 bsz, q_len, _ = hidden_states.size()702 703 query_states = self.q_proj(hidden_states)704 key_states = self.k_proj(hidden_states)705 value_states = self.v_proj(hidden_states)706 707 query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)708 key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)709 value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)710 711 kv_seq_len = key_states.shape[-2]712 if past_key_value is not None:713 kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx)714 cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)715 716 query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)717 718 if past_key_value is not None:719 cache_kwargs = {"sin": sin, "cos": cos} # Specific to RoPE models720 key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)721 722 key_states = repeat_kv(key_states, self.num_key_value_groups)723 value_states = repeat_kv(value_states, self.num_key_value_groups)724 725 if attention_mask is not None:726 if attention_mask.size() != (bsz, 1, q_len, kv_seq_len):727 raise ValueError(728 f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}"729 )730 731 # SDPA with memory-efficient backend is currently (torch==2.1.2) bugged with non-contiguous inputs with custom attn_mask,732 # Reference: https://github.com/pytorch/pytorch/issues/112577.733 if query_states.device.type == "cuda" and attention_mask is not None:734 query_states = query_states.contiguous()735 key_states = key_states.contiguous()736 value_states = value_states.contiguous()737 738 attn_output = torch.nn.functional.scaled_dot_product_attention(739 query_states,740 key_states,741 value_states,742 attn_mask=attention_mask,743 dropout_p=self.attention_dropout if self.training else 0.0,744 # The q_len > 1 is necessary to match with AttentionMaskConverter.to_causal_4d that does not create a causal mask in case q_len == 1.745 is_causal=self.is_causal and attention_mask is None and q_len > 1,746 )747 748 attn_output = attn_output.transpose(1, 2).contiguous()749 attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)750 751 attn_output = self.o_proj(attn_output)752 753 return attn_output, None, past_key_value754 755 756MINICPM_ATTENTION_CLASSES = {757 "eager": MiniCPMAttention,758 "flash_attention_2": MiniCPMFlashAttention2,759 "sdpa": MiniCPMSdpaAttention,760}761 762 763class MiniCPMDecoderLayer(nn.Module):764 def __init__(self, config: MiniCPMConfig, layer_idx: int):765 super().__init__()766 self.hidden_size = config.hidden_size767 self.self_attn = MINICPM_ATTENTION_CLASSES[config._attn_implementation](config=config, layer_idx=layer_idx)768 769 self.mlp = MiniCPMMLP(config)770 self.input_layernorm = MiniCPMRMSNorm(config.hidden_size, eps=config.rms_norm_eps)771 self.post_attention_layernorm = MiniCPMRMSNorm(config.hidden_size, eps=config.rms_norm_eps)772 773 self.scale_depth = config.scale_depth774 self.num_hidden_layers = config.num_hidden_layers775 776 def forward(777 self,778 hidden_states: torch.Tensor,779 attention_mask: Optional[torch.Tensor] = None,780 position_ids: Optional[torch.LongTensor] = None,781 past_key_value: Optional[Tuple[torch.Tensor]] = None,782 output_attentions: Optional[bool] = False,783 use_cache: Optional[bool] = False,784 **kwargs,785 ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]:786 """787 Args:788 hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`789 attention_mask (`torch.FloatTensor`, *optional*):790 attention mask of size `(batch_size, sequence_length)` if flash attention is used or `(batch_size, 1,791 query_sequence_length, key_sequence_length)` if default attention is used.792 output_attentions (`bool`, *optional*):793 Whether or not to return the attentions tensors of all attention layers. See `attentions` under794 returned tensors for more detail.795 use_cache (`bool`, *optional*):796 If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding797 (see `past_key_values`).798 past_key_value (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states799 """800 if "padding_mask" in kwargs:801 warnings.warn(802 "Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use `attention_mask` instead.`"803 )804 805 residual = hidden_states806 hidden_states = self.input_layernorm(hidden_states)807 # Self Attention808 hidden_states, self_attn_weights, present_key_value = self.self_attn(809 hidden_states=hidden_states,810 attention_mask=attention_mask,811 position_ids=position_ids,812 past_key_value=past_key_value,813 output_attentions=output_attentions,814 use_cache=use_cache,815 **kwargs,816 )817 818 hidden_states = residual + hidden_states * (self.scale_depth / math.sqrt(self.num_hidden_layers))819 820 # Fully Connected821 residual = hidden_states822 hidden_states = self.post_attention_layernorm(hidden_states)823 824 hidden_states = self.mlp(hidden_states)825 hidden_states = residual + hidden_states * (self.scale_depth / math.sqrt(self.num_hidden_layers))826 827 outputs = (hidden_states,)828 829 if output_attentions:830 outputs += (self_attn_weights,)831 832 if use_cache:833 outputs += (present_key_value,)834 835 return outputs836 837 838MINICPM_START_DOCSTRING = r"""839 This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the840 library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads841 etc.)842 843 This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.844 Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage845 and behavior.846 847 Parameters:848 config ([`MiniCPMConfig`]):849 Model configuration class with all the parameters of the model. Initializing with a config file does not850 load the weights associated with the model, only the configuration. Check out the851 [`~PreTrainedModel.from_pretrained`] method to load the model weights.852"""853 854 855@add_start_docstrings(856 "The bare MiniCPM Model outputting raw hidden-states without any specific head on top.",857 MINICPM_START_DOCSTRING,858)859class MiniCPMPreTrainedModel(PreTrainedModel):860 config_class = MiniCPMConfig861 base_model_prefix = "model"862 supports_gradient_checkpointing = True863 _no_split_modules = ["MiniCPMDecoderLayer"]864 _skip_keys_device_placement = "past_key_values"865 _supports_flash_attn_2 = True866 _supports_sdpa = True867 _supports_cache_class = True868 869 def _init_weights(self, module):870 std = self.config.initializer_range871 if isinstance(module, nn.Linear):872 module.weight.data.normal_(mean=0.0, std=std)873 if module.bias is not None:874 module.bias.data.zero_()875 elif isinstance(module, nn.Embedding):876 module.weight.data.normal_(mean=0.0, std=std)877 if module.padding_idx is not None:878 module.weight.data[module.padding_idx].zero_()879 880 881MINICPM_INPUTS_DOCSTRING = r"""882 Args:883 input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):884 Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide885 it.886 887 Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and888 [`PreTrainedTokenizer.__call__`] for details.889 890 [What are input IDs?](../glossary#input-ids)891 attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):892 Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:893 894 - 1 for tokens that are **not masked**,895 - 0 for tokens that are **masked**.896 897 [What are attention masks?](../glossary#attention-mask)898 899 Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and900 [`PreTrainedTokenizer.__call__`] for details.901 902 If `past_key_values` is used, optionally only the last `input_ids` have to be input (see903 `past_key_values`).904 905 If you want to change padding behavior, you should read [`modeling_opt._prepare_decoder_attention_mask`]906 and modify to your needs. See diagram 1 in [the paper](https://arxiv.org/abs/1910.13461) for more907 information on the default strategy.908 909 - 1 indicates the head is **not masked**,910 - 0 indicates the head is **masked**.911 position_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):912 Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,913 config.n_positions - 1]`.914 915 [What are position IDs?](../glossary#position-ids)916 past_key_values (`Cache` or `tuple(tuple(torch.FloatTensor))`, *optional*):917 Pre-computed hidden-states (key and values in the self-attention blocks and in the cross-attention918 blocks) that can be used to speed up sequential decoding. This typically consists in the `past_key_values`919 returned by the model at a previous stage of decoding, when `use_cache=True` or `config.use_cache=True`.920 921 Two formats are allowed:922 - a [`~cache_utils.Cache`] instance;923 - Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of924 shape `(batch_size, num_heads, sequence_length, embed_size_per_head)`). This is also known as the legacy925 cache format.926 927 The model will output the same cache format that is fed as input. If no `past_key_values` are passed, the928 legacy cache format will be returned.929 930 If `past_key_values` are used, the user can optionally input only the last `input_ids` (those that don't931 have their past key value states given to this model) of shape `(batch_size, 1)` instead of all `input_ids`932 of shape `(batch_size, sequence_length)`.933 inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):934 Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This935 is useful if you want more control over how to convert `input_ids` indices into associated vectors than the936 model's internal embedding lookup matrix.937 use_cache (`bool`, *optional*):938 If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see939 `past_key_values`).940 output_attentions (`bool`, *optional*):941 Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned942 tensors for more detail.943 output_hidden_states (`bool`, *optional*):944 Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for945 more detail.946 return_dict (`bool`, *optional*):947 Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.948"""949 950 951@add_start_docstrings(952 "The bare MiniCPM Model outputting raw hidden-states without any specific head on top.",953 MINICPM_START_DOCSTRING,954)955class MiniCPMModel(MiniCPMPreTrainedModel):956 """957 Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`MiniCPMDecoderLayer`]958 959 Args:960 config: MiniCPMConfig961 """962 963 def __init__(self, config: MiniCPMConfig):964 super().__init__(config)965 self.padding_idx = config.pad_token_id966 self.vocab_size = config.vocab_size967 968 self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)969 self.layers = nn.ModuleList(970 [MiniCPMDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]971 )972 self._use_sdpa = config._attn_implementation == "sdpa"973 self._use_flash_attention_2 = config._attn_implementation == "flash_attention_2"974 975 self.norm = MiniCPMRMSNorm(config.hidden_size, eps=config.rms_norm_eps)976 977 self.gradient_checkpointing = False978 # Initialize weights and apply final processing979 self.post_init()980 981 def get_input_embeddings(self):982 return self.embed_tokens983 984 def set_input_embeddings(self, value):985 self.embed_tokens = value986 987 @add_start_docstrings_to_model_forward(MINICPM_INPUTS_DOCSTRING)988 def forward(989 self,990 input_ids: torch.LongTensor = None,991 attention_mask: Optional[torch.Tensor] = None,992 position_ids: Optional[torch.LongTensor] = None,993 past_key_values: Optional[List[torch.FloatTensor]] = None,994 inputs_embeds: Optional[torch.FloatTensor] = None,995 use_cache: Optional[bool] = None,996 output_attentions: Optional[bool] = None,997 output_hidden_states: Optional[bool] = None,998 return_dict: Optional[bool] = None,999 ) -> Union[Tuple, BaseModelOutputWithPast]:1000 output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions1001 output_hidden_states = (1002 output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states1003 )1004 use_cache = use_cache if use_cache is not None else self.config.use_cache1005 1006 return_dict = return_dict if return_dict is not None else self.config.use_return_dict1007 1008 # retrieve input_ids and inputs_embeds1009 if input_ids is not None and inputs_embeds is not None:1010 raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")1011 elif input_ids is not None:1012 batch_size, seq_length = input_ids.shape[:2]1013 elif inputs_embeds is not None:1014 batch_size, seq_length = inputs_embeds.shape[:2]1015 else:1016 raise ValueError("You have to specify either input_ids or inputs_embeds")1017 1018 if self.gradient_checkpointing and self.training:1019 if use_cache:1020 logger.warning_once(1021 "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."1022 )1023 use_cache = False1024 1025 past_key_values_length = 01026 if use_cache:1027 use_legacy_cache = not isinstance(past_key_values, Cache)1028 if use_legacy_cache:1029 past_key_values = DynamicCache.from_legacy_cache(past_key_values)1030 past_key_values_length = past_key_values.get_usable_length(seq_length)1031 1032 if position_ids is None:1033 device = input_ids.device if input_ids is not None else inputs_embeds.device1034 position_ids = torch.arange(1035 past_key_values_length, seq_length + past_key_values_length, dtype=torch.long, device=device1036 )1037 position_ids = position_ids.unsqueeze(0)1038 1039 if inputs_embeds is None:1040 inputs_embeds = self.embed_tokens(input_ids) * self.config.scale_emb1041 1042 if self._use_flash_attention_2:1043 # 2d mask is passed through the layers1044 attention_mask = attention_mask if (attention_mask is not None and 0 in attention_mask) else None1045 elif self._use_sdpa and not output_attentions:1046 # output_attentions=True can not be supported when using SDPA, and we fall back on1047 # the manual implementation that requires a 4D causal mask in all cases.1048 attention_mask = _prepare_4d_causal_attention_mask_for_sdpa(1049 attention_mask,1050 (batch_size, seq_length),1051 inputs_embeds,1052 past_key_values_length,1053 )1054 else:1055 # 4d mask is passed through the layers1056 attention_mask = _prepare_4d_causal_attention_mask(1057 attention_mask, (batch_size, seq_length), inputs_embeds, past_key_values_length1058 )1059 1060 # embed positions1061 hidden_states = inputs_embeds1062 1063 # decoder layers1064 all_hidden_states = () if output_hidden_states else None1065 all_self_attns = () if output_attentions else None1066 next_decoder_cache = None1067 1068 for decoder_layer in self.layers:1069 if output_hidden_states:1070 all_hidden_states += (hidden_states,)1071 1072 if self.gradient_checkpointing and self.training:1073 layer_outputs = self._gradient_checkpointing_func(1074 decoder_layer.__call__,1075 hidden_states,1076 attention_mask,1077 position_ids,1078 past_key_values,1079 output_attentions,1080 use_cache,1081 )1082 else:1083 layer_outputs = decoder_layer(1084 hidden_states,1085 attention_mask=attention_mask,1086 position_ids=position_ids,1087 past_key_value=past_key_values,1088 output_attentions=output_attentions,1089 use_cache=use_cache,1090 )1091 1092 hidden_states = layer_outputs[0]1093 1094 if use_cache:1095 next_decoder_cache = layer_outputs[2 if output_attentions else 1]1096 1097 if output_attentions:1098 all_self_attns += (layer_outputs[1],)1099 1100 hidden_states = self.norm(hidden_states)1101 1102 # add hidden states from the last decoder layer1103 if output_hidden_states:1104 all_hidden_states += (hidden_states,)1105 1106 next_cache = None1107 if use_cache:1108 next_cache = next_decoder_cache.to_legacy_cache() if use_legacy_cache else next_decoder_cache1109 if not return_dict:1110 return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None)1111 return BaseModelOutputWithPast(1112 last_hidden_state=hidden_states,1113 past_key_values=next_cache,1114 hidden_states=all_hidden_states,1115 attentions=all_self_attns,1116 )1117 1118 1119class MiniCPMForCausalLM(MiniCPMPreTrainedModel):1120 _tied_weights_keys = ["lm_head.weight"]1121 1122 def __init__(self, config):1123 super().__init__(config)1124 self.model = MiniCPMModel(config)1125 self.vocab_size = config.vocab_size1126 self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)1127 1128 # Initialize weights and apply final processing1129 self.post_init()1130 1131 def get_input_embeddings(self):1132 return self.model.embed_tokens1133 1134 def set_input_embeddings(self, value):1135 self.model.embed_tokens = value1136 1137 def get_output_embeddings(self):1138 return self.lm_head1139 1140 def set_output_embeddings(self, new_embeddings):1141 self.lm_head = new_embeddings1142 1143 def set_decoder(self, decoder):1144 self.model = decoder1145 1146 def get_decoder(self):1147 return self.model1148 1149 @add_start_docstrings_to_model_forward(MINICPM_INPUTS_DOCSTRING)1150 @replace_return_docstrings(output_type=CausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC)1151 def forward(1152 self,1153 input_ids: torch.LongTensor = None,1154 attention_mask: Optional[torch.Tensor] = None,1155 position_ids: Optional[torch.LongTensor] = None,1156 past_key_values: Optional[List[torch.FloatTensor]] = None,1157 inputs_embeds: Optional[torch.FloatTensor] = None,1158 labels: Optional[torch.LongTensor] = None,1159 use_cache: Optional[bool] = None,1160 output_attentions: Optional[bool] = None,1161 output_hidden_states: Optional[bool] = None,1162 return_dict: Optional[bool] = None,1163 ) -> Union[Tuple, CausalLMOutputWithPast]:1164 r"""1165 Args:1166 labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):1167 Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,1168 config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored1169 (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.1170 1171 Returns:1172 1173 Example:1174 1175 ```python1176 >>> from transformers import AutoTokenizer, MiniCPMForCausalLM1177 1178 >>> model = MiniCPMForCausalLM.from_pretrained(PATH_TO_CONVERTED_WEIGHTS)1179 >>> tokenizer = AutoTokenizer.from_pretrained(PATH_TO_CONVERTED_TOKENIZER)1180 1181 >>> prompt = "Hey, are you conscious? Can you talk to me?"1182 >>> inputs = tokenizer(prompt, return_tensors="pt")1183 1184 >>> # Generate1185 >>> generate_ids = model.generate(inputs.input_ids, max_length=30)1186 >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]1187 "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you."1188 ```"""1189 output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions1190 output_hidden_states = (1191 output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states1192 )1193 return_dict = return_dict if return_dict is not None else self.config.use_return_dict1194 1195 # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)1196 outputs = self.model(1197 input_ids=input_ids,1198 attention_mask=attention_mask,1199 position_ids=position_ids,1200 past_key_values=past_key_values,