katuni4ka/tiny-random-glm4
0208
1""" PyTorch ChatGLM model. """2import json3import math4import copy5import warnings6import re7import sys8 9import transformers10import torch11import torch.utils.checkpoint12import torch.nn.functional as F13from torch import nn14from torch.nn import CrossEntropyLoss, LayerNorm, MSELoss, BCEWithLogitsLoss15from torch.nn.utils import skip_init16from typing import Optional, Tuple, Union, List, Callable, Dict, Any17from copy import deepcopy18 19from transformers.modeling_outputs import (20 BaseModelOutputWithPast,21 CausalLMOutputWithPast,22 SequenceClassifierOutputWithPast,23)24from transformers.modeling_utils import PreTrainedModel25from transformers.utils import logging, is_torch_npu_available26from transformers.generation.logits_process import LogitsProcessor27from transformers.generation.utils import LogitsProcessorList, StoppingCriteriaList, GenerationConfig, ModelOutput28 29from .configuration_chatglm import ChatGLMConfig30 31# flags required to enable jit fusion kernels32 33if sys.platform != 'darwin' and not is_torch_npu_available():34 torch._C._jit_set_profiling_mode(False)35 torch._C._jit_set_profiling_executor(False)36 torch._C._jit_override_can_fuse_on_cpu(True)37 torch._C._jit_override_can_fuse_on_gpu(True)38 39logger = logging.get_logger(__name__)40 41_CHECKPOINT_FOR_DOC = "THUDM/ChatGLM"42_CONFIG_FOR_DOC = "ChatGLMConfig"43 44is_transformers_4_42_or_higher = int(transformers.__version__.split(".")[1]) >= 4245is_transformers_4_44_or_higher = int(transformers.__version__.split(".")[1]) >= 4446 47 48def default_init(cls, *args, **kwargs):49 return cls(*args, **kwargs)50 51 52class InvalidScoreLogitsProcessor(LogitsProcessor):53 def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:54 if torch.isnan(scores).any() or torch.isinf(scores).any():55 scores.zero_()56 scores[..., 198] = 5e457 return scores58 59 60def split_tensor_along_last_dim(61 tensor: torch.Tensor,62 num_partitions: int,63 contiguous_split_chunks: bool = False,64) -> List[torch.Tensor]:65 """Split a tensor along its last dimension.66 67 Arguments:68 tensor: input tensor.69 num_partitions: number of partitions to split the tensor70 contiguous_split_chunks: If True, make each chunk contiguous71 in memory.72 73 Returns:74 A list of Tensors75 """76 # Get the size and dimension.77 last_dim = tensor.dim() - 178 last_dim_size = tensor.size()[last_dim] // num_partitions79 # Split.80 tensor_list = torch.split(tensor, last_dim_size, dim=last_dim)81 # Note: torch.split does not create contiguous tensors by default.82 if contiguous_split_chunks:83 return tuple(chunk.contiguous() for chunk in tensor_list)84 85 return tensor_list86 87 88class RotaryEmbedding(nn.Module):89 def __init__(self, dim, rope_ratio=1, original_impl=False, device=None, dtype=None):90 super().__init__()91 inv_freq = 1.0 / (10000 ** (torch.arange(0, dim, 2, device=device).to(dtype=dtype) / dim))92 self.register_buffer("inv_freq", inv_freq)93 self.dim = dim94 self.original_impl = original_impl95 self.rope_ratio = rope_ratio96 97 def forward_impl(98 self, seq_len: int, n_elem: int, dtype: torch.dtype, device: torch.device, base: int = 1000099 ):100 """Enhanced Transformer with Rotary Position Embedding.101 102 Derived from: https://github.com/labmlai/annotated_deep_learning_paper_implementations/blob/master/labml_nn/103 transformers/rope/__init__.py. MIT License:104 https://github.com/labmlai/annotated_deep_learning_paper_implementations/blob/master/license.105 """106 # $\Theta = {\theta_i = 10000^{\frac{2(i-1)}{d}}, i \in [1, 2, ..., \frac{d}{2}]}$107 base = base * self.rope_ratio108 theta = 1.0 / (base ** (torch.arange(0, n_elem, 2, dtype=torch.float, device=device) / n_elem))109 110 # Create position indexes `[0, 1, ..., seq_len - 1]`111 seq_idx = torch.arange(seq_len, dtype=torch.float, device=device)112 113 # Calculate the product of position index and $\theta_i$114 idx_theta = torch.outer(seq_idx, theta).float()115 116 cache = torch.stack([torch.cos(idx_theta), torch.sin(idx_theta)], dim=-1)117 118 # this is to mimic the behaviour of complex32, else we will get different results119 if dtype in (torch.float16, torch.bfloat16, torch.int8):120 cache = cache.bfloat16() if dtype == torch.bfloat16 else cache.half()121 return cache122 123 def forward(self, max_seq_len, offset=0):124 return self.forward_impl(125 max_seq_len, self.dim, dtype=self.inv_freq.dtype, device=self.inv_freq.device126 )127 128 129@torch.jit.script130def apply_rotary_pos_emb(x: torch.Tensor, rope_cache: torch.Tensor) -> torch.Tensor:131 # x: [b, np, sq, hn]132 b, np, sq, hn = x.size(0), x.size(1), x.size(2), x.size(3)133 rot_dim = rope_cache.shape[-2] * 2134 x, x_pass = x[..., :rot_dim], x[..., rot_dim:]135 # truncate to support variable sizes136 rope_cache = rope_cache[:, :sq]137 xshaped = x.reshape(b, np, sq, rot_dim // 2, 2)138 rope_cache = rope_cache.view(-1, 1, sq, xshaped.size(3), 2)139 x_out2 = torch.stack(140 [141 xshaped[..., 0] * rope_cache[..., 0] - xshaped[..., 1] * rope_cache[..., 1],142 xshaped[..., 1] * rope_cache[..., 0] + xshaped[..., 0] * rope_cache[..., 1],143 ],144 -1,145 )146 x_out2 = x_out2.flatten(3)147 return torch.cat((x_out2, x_pass), dim=-1)148 149 150class RMSNorm(torch.nn.Module):151 def __init__(self, normalized_shape, eps=1e-5, device=None, dtype=None, **kwargs):152 super().__init__()153 self.weight = torch.nn.Parameter(torch.empty(normalized_shape, device=device, dtype=dtype))154 self.eps = eps155 156 def forward(self, hidden_states: torch.Tensor):157 input_dtype = hidden_states.dtype158 variance = hidden_states.to(torch.float32).pow(2).mean(-1, keepdim=True)159 hidden_states = hidden_states * torch.rsqrt(variance + self.eps)160 161 return (self.weight * hidden_states).to(input_dtype)162 163 164class CoreAttention(torch.nn.Module):165 def __init__(self, config: ChatGLMConfig, layer_number):166 super(CoreAttention, self).__init__()167 168 self.apply_query_key_layer_scaling = config.apply_query_key_layer_scaling169 self.attention_softmax_in_fp32 = config.attention_softmax_in_fp32170 if self.apply_query_key_layer_scaling:171 self.attention_softmax_in_fp32 = True172 self.layer_number = max(1, layer_number)173 174 projection_size = config.kv_channels * config.num_attention_heads175 176 # Per attention head and per partition values.177 self.hidden_size_per_partition = projection_size178 self.hidden_size_per_attention_head = projection_size // config.num_attention_heads179 self.num_attention_heads_per_partition = config.num_attention_heads180 181 coeff = None182 self.norm_factor = math.sqrt(self.hidden_size_per_attention_head)183 if self.apply_query_key_layer_scaling:184 coeff = self.layer_number185 self.norm_factor *= coeff186 self.coeff = coeff187 188 self.attention_dropout = torch.nn.Dropout(config.attention_dropout)189 190 def forward(self, query_layer, key_layer, value_layer, attention_mask):191 pytorch_major_version = int(torch.__version__.split('.')[0])192 if pytorch_major_version >= 2:193 if attention_mask is None and query_layer.shape[2] == key_layer.shape[2]:194 context_layer = torch.nn.functional.scaled_dot_product_attention(query_layer, key_layer, value_layer,195 is_causal=True)196 else:197 if attention_mask is not None:198 attention_mask = ~attention_mask199 context_layer = torch.nn.functional.scaled_dot_product_attention(query_layer, key_layer, value_layer,200 attention_mask)201 context_layer = context_layer.transpose(1, 2).contiguous()202 new_context_layer_shape = context_layer.size()[:-2] + (self.hidden_size_per_partition,)203 context_layer = context_layer.reshape(*new_context_layer_shape)204 else:205 # Raw attention scores206 207 # [b, np, sq, sk]208 output_size = (query_layer.size(0), query_layer.size(1), query_layer.size(2), key_layer.size(2))209 210 # [b, np, sq, hn] -> [b * np, sq, hn]211 query_layer = query_layer.view(output_size[0] * output_size[1], output_size[2], -1)212 # [b, np, sk, hn] -> [b * np, sk, hn]213 key_layer = key_layer.view(output_size[0] * output_size[1], output_size[3], -1)214 215 # preallocting input tensor: [b * np, sq, sk]216 matmul_input_buffer = torch.empty(217 output_size[0] * output_size[1], output_size[2], output_size[3], dtype=query_layer.dtype,218 device=query_layer.device219 )220 221 # Raw attention scores. [b * np, sq, sk]222 matmul_result = torch.baddbmm(223 matmul_input_buffer,224 query_layer, # [b * np, sq, hn]225 key_layer.transpose(1, 2), # [b * np, hn, sk]226 beta=0.0,227 alpha=(1.0 / self.norm_factor),228 )229 230 # change view to [b, np, sq, sk]231 attention_scores = matmul_result.view(*output_size)232 233 # ===========================234 # Attention probs and dropout235 # ===========================236 237 # attention scores and attention mask [b, np, sq, sk]238 if self.attention_softmax_in_fp32:239 attention_scores = attention_scores.float()240 if self.coeff is not None:241 attention_scores = attention_scores * self.coeff242 if attention_mask is None and attention_scores.shape[2] == attention_scores.shape[3]:243 attention_mask = torch.ones(output_size[0], 1, output_size[2], output_size[3],244 device=attention_scores.device, dtype=torch.bool)245 attention_mask.tril_()246 attention_mask = ~attention_mask247 if attention_mask is not None:248 attention_scores = attention_scores.masked_fill(attention_mask, float("-inf"))249 attention_probs = F.softmax(attention_scores, dim=-1)250 attention_probs = attention_probs.type_as(value_layer)251 252 # This is actually dropping out entire tokens to attend to, which might253 # seem a bit unusual, but is taken from the original Transformer paper.254 attention_probs = self.attention_dropout(attention_probs)255 256 # query layer shape: [b * np, sq, hn]257 # value layer shape: [b, np, sk, hn]258 # attention shape: [b, np, sq, sk]259 # context layer shape: [b, np, sq, hn]260 output_size = (value_layer.size(0), value_layer.size(1), query_layer.size(1), value_layer.size(3))261 # change view [b * np, sk, hn]262 value_layer = value_layer.view(output_size[0] * output_size[1], value_layer.size(2), -1)263 # change view [b * np, sq, sk]264 attention_probs = attention_probs.view(output_size[0] * output_size[1], output_size[2], -1)265 # matmul: [b * np, sq, hn]266 context_layer = torch.bmm(attention_probs, value_layer)267 # change view [b, np, sq, hn]268 context_layer = context_layer.view(*output_size)269 # [b, np, sq, hn] --> [b, sq, np, hn]270 context_layer = context_layer.transpose(1, 2).contiguous()271 # [b, sq, np, hn] --> [b, sq, hp]272 new_context_layer_shape = context_layer.size()[:-2] + (self.hidden_size_per_partition,)273 context_layer = context_layer.reshape(*new_context_layer_shape)274 275 return context_layer276 277 278class SelfAttention(torch.nn.Module):279 """Parallel self-attention layer abstract class.280 281 Self-attention layer takes input with size [s, b, h]282 and returns output of the same size.283 """284 285 def __init__(self, config: ChatGLMConfig, layer_number, device=None):286 super(SelfAttention, self).__init__()287 self.layer_number = max(1, layer_number)288 289 self.projection_size = config.kv_channels * config.num_attention_heads290 291 # Per attention head and per partition values.292 self.hidden_size_per_attention_head = self.projection_size // config.num_attention_heads293 self.num_attention_heads_per_partition = config.num_attention_heads294 295 self.multi_query_attention = config.multi_query_attention296 self.qkv_hidden_size = 3 * self.projection_size297 if self.multi_query_attention:298 self.num_multi_query_groups_per_partition = config.multi_query_group_num299 self.qkv_hidden_size = (300 self.projection_size + 2 * self.hidden_size_per_attention_head * config.multi_query_group_num301 )302 self.query_key_value = nn.Linear(config.hidden_size, self.qkv_hidden_size,303 bias=config.add_bias_linear or config.add_qkv_bias,304 device=device, **_config_to_kwargs(config)305 )306 307 self.core_attention = CoreAttention(config, self.layer_number)308 309 # Output.310 self.dense = nn.Linear(self.projection_size, config.hidden_size, bias=config.add_bias_linear,311 device=device, **_config_to_kwargs(config)312 )313 314 def _allocate_memory(self, inference_max_sequence_len, batch_size, device=None, dtype=None):315 if self.multi_query_attention:316 num_attention_heads = self.num_multi_query_groups_per_partition317 else:318 num_attention_heads = self.num_attention_heads_per_partition319 return torch.empty(320 inference_max_sequence_len,321 batch_size,322 num_attention_heads,323 self.hidden_size_per_attention_head,324 dtype=dtype,325 device=device,326 )327 328 def forward(329 self, hidden_states, attention_mask, rotary_pos_emb, kv_cache=None, use_cache=True330 ):331 # hidden_states: [b, sq, h]332 333 # =================================================334 # Pre-allocate memory for key-values for inference.335 # =================================================336 # =====================337 # Query, Key, and Value338 # =====================339 340 # Attention heads [b, sq, h] --> [b, sq, (np * 3 * hn)]341 mixed_x_layer = self.query_key_value(hidden_states)342 343 if self.multi_query_attention:344 (query_layer, key_layer, value_layer) = mixed_x_layer.split(345 [346 self.num_attention_heads_per_partition * self.hidden_size_per_attention_head,347 self.num_multi_query_groups_per_partition * self.hidden_size_per_attention_head,348 self.num_multi_query_groups_per_partition * self.hidden_size_per_attention_head,349 ],350 dim=-1,351 )352 query_layer = query_layer.view(353 query_layer.size()[:-1] + (self.num_attention_heads_per_partition, self.hidden_size_per_attention_head)354 )355 key_layer = key_layer.view(356 key_layer.size()[:-1] + (self.num_multi_query_groups_per_partition, self.hidden_size_per_attention_head)357 )358 value_layer = value_layer.view(359 value_layer.size()[:-1]360 + (self.num_multi_query_groups_per_partition, self.hidden_size_per_attention_head)361 )362 else:363 new_tensor_shape = mixed_x_layer.size()[:-1] + \364 (self.num_attention_heads_per_partition,365 3 * self.hidden_size_per_attention_head)366 mixed_x_layer = mixed_x_layer.view(*new_tensor_shape)367 368 # [b, sq, np, 3 * hn] --> 3 [b, sq, np, hn]369 (query_layer, key_layer, value_layer) = split_tensor_along_last_dim(mixed_x_layer, 3)370 371 # [b, sq, np, hn] -> [b, np, sq, hn]372 query_layer, key_layer, value_layer = [k.transpose(1, 2) for k in [query_layer, key_layer, value_layer]]373 374 # apply relative positional encoding (rotary embedding)375 if rotary_pos_emb is not None:376 query_layer = apply_rotary_pos_emb(query_layer, rotary_pos_emb)377 key_layer = apply_rotary_pos_emb(key_layer, rotary_pos_emb)378 379 # adjust key and value for inference380 if kv_cache is not None:381 cache_k, cache_v = kv_cache382 key_layer = torch.cat((cache_k, key_layer), dim=2)383 value_layer = torch.cat((cache_v, value_layer), dim=2)384 if use_cache:385 if kv_cache is None:386 kv_cache = torch.cat((key_layer.unsqueeze(0).unsqueeze(0), value_layer.unsqueeze(0).unsqueeze(0)), dim=1)387 else:388 kv_cache = (key_layer, value_layer)389 else:390 kv_cache = None391 392 if self.multi_query_attention:393 key_layer = key_layer.unsqueeze(2)394 key_layer = key_layer.expand(395 -1, -1, self.num_attention_heads_per_partition // self.num_multi_query_groups_per_partition, -1, -1396 )397 key_layer = key_layer.contiguous().view(398 key_layer.size()[:1] + (self.num_attention_heads_per_partition,) + key_layer.size()[3:]399 )400 value_layer = value_layer.unsqueeze(2)401 value_layer = value_layer.expand(402 -1, -1, self.num_attention_heads_per_partition // self.num_multi_query_groups_per_partition, -1, -1403 )404 value_layer = value_layer.contiguous().view(405 value_layer.size()[:1] + (self.num_attention_heads_per_partition,) + value_layer.size()[3:]406 )407 408 # ==================================409 # core attention computation410 # ==================================411 412 context_layer = self.core_attention(query_layer, key_layer, value_layer, attention_mask)413 414 # =================415 # Output. [sq, b, h]416 # =================417 418 output = self.dense(context_layer)419 420 return output, kv_cache421 422 423def _config_to_kwargs(args):424 common_kwargs = {425 "dtype": args.torch_dtype if not isinstance(args.torch_dtype, str) else getattr(torch, args.torch_dtype)426 }427 return common_kwargs428 429 430class MLP(torch.nn.Module):431 """MLP.432 433 MLP will take the input with h hidden state, project it to 4*h434 hidden dimension, perform nonlinear transformation, and project the435 state back into h hidden dimension.436 """437 438 def __init__(self, config: ChatGLMConfig, device=None):439 super(MLP, self).__init__()440 441 self.add_bias = config.add_bias_linear442 443 # Project to 4h. If using swiglu double the output width, see https://arxiv.org/pdf/2002.05202.pdf444 self.dense_h_to_4h = nn.Linear(445 config.hidden_size,446 config.ffn_hidden_size * 2,447 bias=self.add_bias,448 device=device,449 **_config_to_kwargs(config)450 )451 452 def swiglu(x):453 x = torch.chunk(x, 2, dim=-1)454 return F.silu(x[0]) * x[1]455 456 self.activation_func = swiglu457 458 # Project back to h.459 self.dense_4h_to_h = nn.Linear(460 config.ffn_hidden_size,461 config.hidden_size,462 bias=self.add_bias,463 device=device,464 **_config_to_kwargs(config)465 )466 467 def forward(self, hidden_states):468 # [s, b, 4hp]469 intermediate_parallel = self.dense_h_to_4h(hidden_states)470 intermediate_parallel = self.activation_func(intermediate_parallel)471 # [s, b, h]472 output = self.dense_4h_to_h(intermediate_parallel)473 return output474 475 476class GLMBlock(torch.nn.Module):477 """A single transformer layer.478 479 Transformer layer takes input with size [s, b, h] and returns an480 output of the same size.481 """482 483 def __init__(self, config: ChatGLMConfig, layer_number, device=None):484 super(GLMBlock, self).__init__()485 self.layer_number = layer_number486 487 self.apply_residual_connection_post_layernorm = config.apply_residual_connection_post_layernorm488 489 self.fp32_residual_connection = config.fp32_residual_connection490 491 LayerNormFunc = RMSNorm if config.rmsnorm else LayerNorm492 dtype = getattr(torch, config.torch_dtype) if isinstance(config.torch_dtype, str) else config.torch_dtype493 # Layernorm on the input data.494 self.input_layernorm = LayerNormFunc(config.hidden_size, eps=config.layernorm_epsilon, device=device,495 dtype=dtype)496 497 # Self attention.498 self.self_attention = SelfAttention(config, layer_number, device=device)499 self.hidden_dropout = config.hidden_dropout500 501 # Layernorm on the attention output502 self.post_attention_layernorm = LayerNormFunc(config.hidden_size, eps=config.layernorm_epsilon, device=device,503 dtype=dtype)504 505 # MLP506 self.mlp = MLP(config, device=device)507 508 def forward(509 self, hidden_states, attention_mask, rotary_pos_emb, kv_cache=None, use_cache=True,510 ):511 # hidden_states: [s, b, h]512 513 # Layer norm at the beginning of the transformer layer.514 layernorm_output = self.input_layernorm(hidden_states)515 # Self attention.516 attention_output, kv_cache = self.self_attention(517 layernorm_output,518 attention_mask,519 rotary_pos_emb,520 kv_cache=kv_cache,521 use_cache=use_cache522 )523 524 # Residual connection.525 if self.apply_residual_connection_post_layernorm:526 residual = layernorm_output527 else:528 residual = hidden_states529 530 layernorm_input = torch.nn.functional.dropout(attention_output, p=self.hidden_dropout, training=self.training)531 layernorm_input = residual + layernorm_input532 533 # Layer norm post the self attention.534 layernorm_output = self.post_attention_layernorm(layernorm_input)535 536 # MLP.537 mlp_output = self.mlp(layernorm_output)538 539 # Second residual connection.540 if self.apply_residual_connection_post_layernorm:541 residual = layernorm_output542 else:543 residual = layernorm_input544 545 output = torch.nn.functional.dropout(mlp_output, p=self.hidden_dropout, training=self.training)546 output = residual + output547 548 return output, kv_cache549 550 551class GLMTransformer(torch.nn.Module):552 """Transformer class."""553 554 def __init__(self, config: ChatGLMConfig, device=None):555 super(GLMTransformer, self).__init__()556 557 self.fp32_residual_connection = config.fp32_residual_connection558 self.post_layer_norm = config.post_layer_norm559 560 # Number of layers.561 self.num_layers = config.num_layers562 563 # Transformer layers.564 def build_layer(layer_number):565 return GLMBlock(config, layer_number, device=device)566 567 self.layers = torch.nn.ModuleList([build_layer(i + 1) for i in range(self.num_layers)])568 569 if self.post_layer_norm:570 LayerNormFunc = RMSNorm if config.rmsnorm else LayerNorm571 dtype = getattr(torch, config.torch_dtype) if isinstance(config.torch_dtype, str) else config.torch_dtype572 # Final layer norm before output.573 self.final_layernorm = LayerNormFunc(config.hidden_size, eps=config.layernorm_epsilon, device=device,574 dtype=dtype)575 576 self.gradient_checkpointing = False577 578 def _get_layer(self, layer_number):579 return self.layers[layer_number]580 581 def forward(582 self, hidden_states, attention_mask, rotary_pos_emb, kv_caches=None,583 use_cache: Optional[bool] = True,584 output_hidden_states: Optional[bool] = False,585 ):586 if not kv_caches:587 kv_caches = [None for _ in range(self.num_layers)]588 presents = () if use_cache else None589 if self.gradient_checkpointing and self.training:590 if use_cache:591 logger.warning_once(592 "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."593 )594 use_cache = False595 596 all_self_attentions = None597 all_hidden_states = () if output_hidden_states else None598 for index in range(self.num_layers):599 if output_hidden_states:600 all_hidden_states = all_hidden_states + (hidden_states,)601 602 layer = self._get_layer(index)603 if self.gradient_checkpointing and self.training:604 layer_ret = torch.utils.checkpoint.checkpoint(605 layer,606 hidden_states,607 attention_mask,608 rotary_pos_emb,609 kv_caches[index],610 use_cache,611 use_reentrant=False612 )613 else:614 layer_ret = layer(615 hidden_states,616 attention_mask,617 rotary_pos_emb,618 kv_cache=kv_caches[index],619 use_cache=use_cache620 )621 hidden_states, kv_cache = layer_ret622 if use_cache:623 # token by token decoding, use tuple format624 if kv_caches[0] is not None:625 presents = presents + (kv_cache,)626 # prefilling in decoding, use tensor format to save cuda memory627 else:628 if len(presents) == 0:629 presents = kv_cache630 else:631 presents = torch.cat((presents, kv_cache.to(presents.device)), dim=0)632 633 if output_hidden_states:634 all_hidden_states = all_hidden_states + (hidden_states,)635 636 # Final layer norm.637 if self.post_layer_norm:638 hidden_states = self.final_layernorm(hidden_states)639 640 return hidden_states, presents, all_hidden_states, all_self_attentions641 642 643class ChatGLMPreTrainedModel(PreTrainedModel):644 """645 An abstract class to handle weights initialization and646 a simple interface for downloading and loading pretrained models.647 """648 649 is_parallelizable = False650 supports_gradient_checkpointing = True651 config_class = ChatGLMConfig652 base_model_prefix = "transformer"653 _no_split_modules = ["GLMBlock"]654 655 def _init_weights(self, module: nn.Module):656 """Initialize the weights."""657 return658 659 def get_masks(self, input_ids, past_key_values, padding_mask=None):660 batch_size, seq_length = input_ids.shape661 full_attention_mask = torch.ones(batch_size, seq_length, seq_length, device=input_ids.device)662 full_attention_mask.tril_()663 past_length = 0664 if past_key_values:665 past_length = past_key_values[0][0].shape[2]666 if past_length:667 full_attention_mask = torch.cat((torch.ones(batch_size, seq_length, past_length,668 device=input_ids.device), full_attention_mask), dim=-1)669 if padding_mask is not None:670 full_attention_mask = full_attention_mask * padding_mask.unsqueeze(1)671 if not past_length and padding_mask is not None:672 full_attention_mask -= padding_mask.unsqueeze(-1) - 1673 full_attention_mask = (full_attention_mask < 0.5).bool()674 full_attention_mask.unsqueeze_(1)675 return full_attention_mask676 677 def get_position_ids(self, input_ids, device):678 batch_size, seq_length = input_ids.shape679 position_ids = torch.arange(seq_length, dtype=torch.long, device=device).unsqueeze(0).repeat(batch_size, 1)680 return position_ids681 682 def gradient_checkpointing_enable(self, gradient_checkpointing_kwargs=None):683 if not self.supports_gradient_checkpointing:684 raise ValueError(f"{self.__class__.__name__} does not support gradient checkpointing.")685 686 687class Embedding(torch.nn.Module):688 """Language model embeddings."""689 690 def __init__(self, config: ChatGLMConfig, device=None):691 super(Embedding, self).__init__()692 693 self.hidden_size = config.hidden_size694 # Word embeddings (parallel).695 dtype = getattr(torch, config.torch_dtype) if isinstance(config.torch_dtype, str) else config.torch_dtype696 self.word_embeddings = nn.Embedding(697 config.padded_vocab_size,698 self.hidden_size,699 dtype=dtype,700 device=device701 )702 self.fp32_residual_connection = config.fp32_residual_connection703 704 def forward(self, input_ids):705 # Embeddings.706 words_embeddings = self.word_embeddings(input_ids)707 embeddings = words_embeddings708 # If the input flag for fp32 residual connection is set, convert for float.709 if self.fp32_residual_connection:710 embeddings = embeddings.float()711 return embeddings712 713 714class ChatGLMModel(ChatGLMPreTrainedModel):715 def __init__(self, config: ChatGLMConfig, device=None, empty_init=False):716 super().__init__(config)717 if empty_init:718 init_method = skip_init719 else:720 init_method = default_init721 init_kwargs = {}722 if device is not None:723 init_kwargs["device"] = device if not isinstance(device, str) else torch.device(device)724 self.embedding = init_method(Embedding, config, **init_kwargs)725 self.num_layers = config.num_layers726 self.multi_query_group_num = config.multi_query_group_num727 self.kv_channels = config.kv_channels728 729 # Rotary positional embeddings730 self.seq_length = config.seq_length731 rotary_dim = (732 config.hidden_size // config.num_attention_heads if config.kv_channels is None else config.kv_channels733 )734 dtype = getattr(torch, config.torch_dtype) if isinstance(config.torch_dtype, str) else config.torch_dtype735 self.rotary_pos_emb = RotaryEmbedding(rotary_dim // 2, rope_ratio=config.rope_ratio, original_impl=config.original_rope, 736 device=device, dtype=dtype)737 self.encoder = init_method(GLMTransformer, config, **init_kwargs)738 self.output_layer = init_method(nn.Linear, config.hidden_size, config.padded_vocab_size, bias=False,739 dtype=dtype, **init_kwargs)740 741 def get_input_embeddings(self):742 return self.embedding.word_embeddings743 744 def set_input_embeddings(self, value):745 self.embedding.word_embeddings = value746 747 def forward(748 self,749 input_ids,750 position_ids: Optional[torch.Tensor] = None,751 attention_mask: Optional[torch.BoolTensor] = None,752 full_attention_mask: Optional[torch.BoolTensor] = None,753 past_key_values: Optional[Tuple[Tuple[torch.Tensor, torch.Tensor], ...]] = None,754 inputs_embeds: Optional[torch.Tensor] = None,755 use_cache: Optional[bool] = None,756 output_hidden_states: Optional[bool] = None,757 return_dict: Optional[bool] = None,758 ):759 output_hidden_states = (760 output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states761 )762 use_cache = use_cache if use_cache is not None else self.config.use_cache763 return_dict = return_dict if return_dict is not None else self.config.use_return_dict764 765 batch_size, seq_length = input_ids.shape766 767 if inputs_embeds is None:768 inputs_embeds = self.embedding(input_ids)769 770 if full_attention_mask is None:771 if (attention_mask is not None and not attention_mask.all()) or (past_key_values and seq_length != 1):772 full_attention_mask = self.get_masks(input_ids, past_key_values, padding_mask=attention_mask)773 774 # Rotary positional embeddings775 rotary_pos_emb = self.rotary_pos_emb(self.seq_length)776 if position_ids is not None:777 rotary_pos_emb = rotary_pos_emb[position_ids]778 else:779 rotary_pos_emb = rotary_pos_emb[None, :seq_length]780 781 # Run encoder.782 hidden_states, presents, all_hidden_states, all_self_attentions = self.encoder(783 inputs_embeds, full_attention_mask, rotary_pos_emb=rotary_pos_emb,784 kv_caches=past_key_values, use_cache=use_cache, output_hidden_states=output_hidden_states785 )786 if presents is not None and type(presents) is torch.Tensor:787 presents = presents.split(1, dim=0)788 presents = list(presents)789 presents = [list(x.squeeze(0).split(1, dim=0)) for x in presents]790 presents = [tuple([x.squeeze(0) for x in y]) for y in presents]791 presents = tuple(presents)792 793 if not return_dict:794 return tuple(v for v in [hidden_states, presents, all_hidden_states, all_self_attentions] if v is not None)795 796 return BaseModelOutputWithPast(797 last_hidden_state=hidden_states,798 past_key_values=presents,799 hidden_states=all_hidden_states,800 attentions=all_self_attentions,801 )802 803 804class ChatGLMForConditionalGeneration(ChatGLMPreTrainedModel):805 def __init__(self, config: ChatGLMConfig, empty_init=True, device=None):806 super().__init__(config)807 808 self.max_sequence_length = config.max_length809 self.transformer = ChatGLMModel(config, empty_init=empty_init, device=device)810 self.config = config811 812 def _update_model_kwargs_for_generation(813 self,814 outputs: ModelOutput,815 model_kwargs: Dict[str, Any],816 is_encoder_decoder: bool = False,817 standardize_cache_format: bool = False,818 ) -> Dict[str, Any]:819 820 if is_transformers_4_44_or_higher:821 model_kwargs["past_key_values"] = self._extract_past_from_model_output(822 outputs823 )[1]824 # update past_key_values825 elif is_transformers_4_42_or_higher:826 model_kwargs["past_key_values"] = self._extract_past_from_model_output(827 outputs, standardize_cache_format=standardize_cache_format828 )[1]829 else:830 model_kwargs["past_key_values"] = self._extract_past_from_model_output(831 outputs, standardize_cache_format=standardize_cache_format832 )833 834 # update attention mask835 if "attention_mask" in model_kwargs:836 attention_mask = model_kwargs["attention_mask"]837 model_kwargs["attention_mask"] = torch.cat(838 [attention_mask, attention_mask.new_ones((attention_mask.shape[0], 1))], dim=-1839 )840 841 # update position ids842 if "position_ids" in model_kwargs:843 position_ids = model_kwargs["position_ids"]844 new_position_id = position_ids[..., -1:].clone()845 new_position_id += 1846 model_kwargs["position_ids"] = torch.cat(847 [position_ids, new_position_id], dim=-1848 )849 850 model_kwargs["is_first_forward"] = False851 return model_kwargs852 853 def prepare_inputs_for_generation(854 self,855 input_ids: torch.LongTensor,856 past_key_values: Optional[torch.Tensor] = None,857 attention_mask: Optional[torch.Tensor] = None,858 position_ids: Optional[torch.Tensor] = None,859 use_cache: Optional[bool] = None,860 is_first_forward: bool = True,861 **kwargs862 ) -> dict:863 # only last token for input_ids if past is not None864 if position_ids is None:865 position_ids = self.get_position_ids(input_ids, device=input_ids.device)866 if not is_first_forward:867 if past_key_values is not None:868 position_ids = position_ids[..., -1:]869 input_ids = input_ids[:, -1:]870 return {871 "input_ids": input_ids,872 "past_key_values": past_key_values,873 "position_ids": position_ids,874 "attention_mask": attention_mask,875 "return_last_logit": True,876 "use_cache": use_cache877 }878 879 def forward(880 self,881 input_ids: Optional[torch.Tensor] = None,882 position_ids: Optional[torch.Tensor] = None,883 attention_mask: Optional[torch.Tensor] = None,884 past_key_values: Optional[Tuple[torch.FloatTensor]] = None,885 inputs_embeds: Optional[torch.Tensor] = None,886 labels: Optional[torch.Tensor] = None,887 use_cache: Optional[bool] = None,888 output_attentions: Optional[bool] = None,889 output_hidden_states: Optional[bool] = None,890 return_dict: Optional[bool] = None,891 return_last_logit: Optional[bool] = False,892 ):893 use_cache = use_cache if use_cache is not None else self.config.use_cache894 return_dict = return_dict if return_dict is not None else self.config.use_return_dict895 896 transformer_outputs = self.transformer(897 input_ids=input_ids,898 position_ids=position_ids,899 attention_mask=attention_mask,900 past_key_values=past_key_values,901 inputs_embeds=inputs_embeds,902 use_cache=use_cache,903 output_hidden_states=output_hidden_states,904 return_dict=return_dict,905 )906 907 hidden_states = transformer_outputs[0]908 if return_last_logit:909 hidden_states = hidden_states[:, -1:]910 lm_logits = self.transformer.output_layer(hidden_states)911 912 loss = None913 if labels is not None:914 lm_logits = lm_logits.to(torch.float32)915 916 # Shift so that tokens < n predict n917 shift_logits = lm_logits[..., :-1, :].contiguous()918 shift_labels = labels[..., 1:].contiguous()919 # Flatten the tokens920 loss_fct = CrossEntropyLoss(ignore_index=-100)921 loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))922 923 lm_logits = lm_logits.to(hidden_states.dtype)924 loss = loss.to(hidden_states.dtype)925 926 if not return_dict:927 output = (lm_logits,) + transformer_outputs[1:]928 return ((loss,) + output) if loss is not None else output929 930 return CausalLMOutputWithPast(931 loss=loss,932 logits=lm_logits,933 past_key_values=transformer_outputs.past_key_values,934 hidden_states=transformer_outputs.hidden_states,935 attentions=transformer_outputs.attentions,936 )937 938 @staticmethod939 def _reorder_cache(940 past: Tuple[Tuple[torch.Tensor, torch.Tensor], ...], beam_idx: torch.LongTensor941 ) -> Tuple[Tuple[torch.Tensor, torch.Tensor], ...]:942 """943 This function is used to re-order the `past_key_values` cache if [`~PreTrainedModel.beam_search`] or944 [`~PreTrainedModel.beam_sample`] is called. This is required to match `past_key_values` with the correct945 beam_idx at every generation step.946 947 Output shares the same memory storage as `past`.948 """949 return tuple(950 (951 layer_past[0].index_select(0, beam_idx.to(layer_past[0].device)),952 layer_past[1].index_select(0, beam_idx.to(layer_past[1].device)),953 )954 for layer_past in past955 )956 957 @staticmethod958 def _extract_past_from_model_output(outputs: ModelOutput, *args, **kwargs):959 past_key_values = None960 if "past_key_values" in outputs:961 past_key_values = outputs.past_key_values962 if is_transformers_4_42_or_higher:963 return None, past_key_values964 return past_key_values965 966 def process_response(self, output, history):967 content = ""968 history = deepcopy(history)969 for response in output.split("<|assistant|>"):970 if "\n" in response:971 metadata, content = response.split("\n", maxsplit=1)972 else:973 metadata, content = "", response974 if not metadata.strip():975 content = content.strip()976 history.append({"role": "assistant", "metadata": metadata, "content": content})977 content = content.replace("[[训练时间]]", "2023年")978 else:979 history.append({"role": "assistant", "metadata": metadata, "content": content})980 if history[0]["role"] == "system" and "tools" in history[0]:981 parameters = json.loads(content)982 content = {"name": metadata.strip(), "parameters": parameters}983 else:984 content = {"name": metadata.strip(), "content": content}985 return content, history986 987 @torch.inference_mode()988 def chat(self, tokenizer, query: str, history: List[Dict] = None, role: str = "user",989 max_length: int = 8192, num_beams=1, do_sample=True, top_p=0.8, temperature=0.8, logits_processor=None,990 **kwargs):991 if history is None:992 history = []993 if logits_processor is None:994 logits_processor = LogitsProcessorList()995 logits_processor.append(InvalidScoreLogitsProcessor())996 gen_kwargs = {"max_length": max_length, "num_beams": num_beams, "do_sample": do_sample, "top_p": top_p,997 "temperature": temperature, "logits_processor": logits_processor, **kwargs}998 history.append({"role": role, "content": query})999 inputs = tokenizer.apply_chat_template(history, add_generation_prompt=True, tokenize=True,1000 return_tensors="pt", return_dict=True)1001 inputs = inputs.to(self.device)1002 eos_token_id = [tokenizer.eos_token_id, tokenizer.convert_tokens_to_ids("<|user|>"),1003 tokenizer.convert_tokens_to_ids("<|observation|>")]1004 outputs = self.generate(**inputs, **gen_kwargs, eos_token_id=eos_token_id)1005 outputs = outputs.tolist()[0][len(inputs["input_ids"][0]):-1]1006 response = tokenizer.decode(outputs)1007 response, history = self.process_response(response, history)1008 return response, history1009 1010 @torch.inference_mode()1011 def stream_chat(self, tokenizer, query: str, history: List[Dict] = None, role: str = "user",1012 past_key_values=None, max_length: int = 8192, do_sample=True, top_p=0.8, temperature=0.8,1013 logits_processor=None, return_past_key_values=False, **kwargs):1014 if history is None:1015 history = []1016 if logits_processor is None:1017 logits_processor = LogitsProcessorList()1018 logits_processor.append(InvalidScoreLogitsProcessor())1019 eos_token_id = [tokenizer.eos_token_id, tokenizer.convert_tokens_to_ids("<|user|>"),1020 tokenizer.convert_tokens_to_ids("<|observation|>")]1021 gen_kwargs = {"max_length": max_length, "do_sample": do_sample, "top_p": top_p,1022 "temperature": temperature, "logits_processor": logits_processor, **kwargs}1023 if past_key_values is None:1024 inputs = tokenizer.apply_chat_template(history + [{"role": role, "content": query}],1025 add_generation_prompt=True, tokenize=True, return_tensors="pt",1026 return_dict=True)1027 else:1028 inputs = tokenizer.apply_chat_template([{"role": role, "content": query}], add_special_tokens=False,1029 add_generation_prompt=True, tokenize=True, return_tensors="pt",1030 return_dict=True)1031 inputs = inputs.to(self.device)1032 if past_key_values is not None:1033 past_length = past_key_values[0][0].shape[2]1034 inputs.position_ids += past_length1035 attention_mask = inputs.attention_mask1036 attention_mask = torch.cat((attention_mask.new_ones(1, past_length), attention_mask), dim=1)1037 inputs['attention_mask'] = attention_mask1038 history.append({"role": role, "content": query})1039 for outputs in self.stream_generate(**inputs, past_key_values=past_key_values,1040 eos_token_id=eos_token_id, return_past_key_values=return_past_key_values,1041 **gen_kwargs):1042 if return_past_key_values:1043 outputs, past_key_values = outputs1044 outputs = outputs.tolist()[0][len(inputs["input_ids"][0]):-1]1045 response = tokenizer.decode(outputs)1046 if response and response[-1] != "�":1047 response, new_history = self.process_response(response, history)1048 if return_past_key_values:1049 yield response, new_history, past_key_values1050 else:1051 yield response, new_history1052 1053 @torch.inference_mode()1054 def stream_generate(1055 self,1056 input_ids,1057 generation_config: Optional[GenerationConfig] = None,1058 logits_processor: Optional[LogitsProcessorList] = None,1059 stopping_criteria: Optional[StoppingCriteriaList] = None,1060 prefix_allowed_tokens_fn: Optional[Callable[[int, torch.Tensor], List[int]]] = None,1061 return_past_key_values=False,1062 **kwargs,1063 ):1064 batch_size, input_ids_seq_length = input_ids.shape[0], input_ids.shape[-1]1065 1066 if generation_config is None:1067 generation_config = self.generation_config1068 generation_config = copy.deepcopy(generation_config)1069 model_kwargs = generation_config.update(**kwargs)1070 model_kwargs["use_cache"] = generation_config.use_cache1071 bos_token_id, eos_token_id = generation_config.bos_token_id, generation_config.eos_token_id1072 1073 if isinstance(eos_token_id, int):1074 eos_token_id = [eos_token_id]1075 eos_token_id_tensor = torch.tensor(eos_token_id).to(input_ids.device) if eos_token_id is not None else None1076 1077 has_default_max_length = kwargs.get("max_length") is None and generation_config.max_length is not None1078 if has_default_max_length and generation_config.max_new_tokens is None:1079 warnings.warn(1080 f"Using `max_length`'s default ({generation_config.max_length}) to control the generation length. "1081 "This behaviour is deprecated and will be removed from the config in v5 of Transformers -- we"1082 " recommend using `max_new_tokens` to control the maximum length of the generation.",1083 UserWarning,1084 )1085 elif generation_config.max_new_tokens is not None:1086 generation_config.max_length = generation_config.max_new_tokens + input_ids_seq_length1087 if not has_default_max_length:1088 logger.warn(1089 f"Both `max_new_tokens` (={generation_config.max_new_tokens}) and `max_length`(="1090 f"{generation_config.max_length}) seem to have been set. `max_new_tokens` will take precedence. "1091 "Please refer to the documentation for more information. "1092 "(https://huggingface.co/docs/transformers/main/en/main_classes/text_generation)",1093 UserWarning,1094 )1095 1096 if input_ids_seq_length >= generation_config.max_length:1097 input_ids_string = "decoder_input_ids" if self.config.is_encoder_decoder else "input_ids"1098 logger.warning(1099 f"Input length of {input_ids_string} is {input_ids_seq_length}, but `max_length` is set to"1100 f" {generation_config.max_length}. This can lead to unexpected behavior. You should consider"1101 " increasing `max_new_tokens`."1102 )1103 1104 # 2. Set generation parameters if not already defined1105 logits_processor = logits_processor if logits_processor is not None else LogitsProcessorList()1106 stopping_criteria = stopping_criteria if stopping_criteria is not None else StoppingCriteriaList()1107 1108 logits_processor = self._get_logits_processor(1109 generation_config=generation_config,1110 input_ids_seq_length=input_ids_seq_length,1111 encoder_input_ids=input_ids,1112 prefix_allowed_tokens_fn=prefix_allowed_tokens_fn,1113 logits_processor=logits_processor,1114 )1115 1116 stopping_criteria = self._get_stopping_criteria(1117 generation_config=generation_config, stopping_criteria=stopping_criteria1118 )1119 logits_warper = self._get_logits_warper(generation_config)1120 1121 unfinished_sequences = input_ids.new(input_ids.shape[0]).fill_(1)1122 scores = None1123 while True:1124 model_inputs = self.prepare_inputs_for_generation(input_ids, **model_kwargs)1125 # forward pass to get next token1126 outputs = self(1127 **model_inputs,1128 return_dict=True,1129 output_attentions=False,1130 output_hidden_states=False,1131 )1132 1133 next_token_logits = outputs.logits[:, -1, :]1134 1135 # pre-process distribution1136 next_token_scores = logits_processor(input_ids, next_token_logits)1137 next_token_scores = logits_warper(input_ids, next_token_scores)1138 1139 # sample1140 probs = nn.functional.softmax(next_token_scores, dim=-1)1141 if generation_config.do_sample:1142 next_tokens = torch.multinomial(probs, num_samples=1).squeeze(1)1143 else:1144 next_tokens = torch.argmax(probs, dim=-1)1145 # update generated ids, model inputs, and length for next step1146 input_ids = torch.cat([input_ids, next_tokens[:, None]], dim=-1)1147 model_kwargs = self._update_model_kwargs_for_generation(1148 outputs, model_kwargs, is_encoder_decoder=self.config.is_encoder_decoder1149 )1150 unfinished_sequences = unfinished_sequences.mul(1151 next_tokens.tile(eos_token_id_tensor.shape[0], 1).ne(eos_token_id_tensor.unsqueeze(1)).prod(dim=0)1152 )1153 if return_past_key_values:1154 yield input_ids, outputs.past_key_values1155 else:1156 yield input_ids1157 # stop when each sentence is finished, or if we exceed the maximum length1158 if unfinished_sequences.max() == 0 or stopping_criteria(input_ids, scores):1159 break1160 1161 1162class ChatGLMForSequenceClassification(ChatGLMPreTrainedModel):1163 def __init__(self, config: ChatGLMConfig, empty_init=True, device=None):1164 super().__init__(config)1165 1166 self.num_labels = config.num_labels1167 self.transformer = ChatGLMModel(config, empty_init=empty_init, device=device)1168 dtype = getattr(torch, config.torch_dtype) if isinstance(config.torch_dtype, str) else config.torch_dtype1169 1170 self.classifier_head = nn.Linear(config.hidden_size, config.num_labels, bias=True, dtype=dtype)1171 if config.classifier_dropout is not None:1172 self.dropout = nn.Dropout(config.classifier_dropout)1173 else:1174 self.dropout = None1175 self.config = config1176 1177 def forward(1178 self,1179 input_ids: Optional[torch.LongTensor] = None,1180 position_ids: Optional[torch.LongTensor] = None,1181 attention_mask: Optional[torch.Tensor] = None,1182 full_attention_mask: Optional[torch.Tensor] = None,1183 past_key_values: Optional[Tuple[Tuple[torch.Tensor, torch.Tensor], ...]] = None,1184 inputs_embeds: Optional[torch.LongTensor] = None,1185 labels: Optional[torch.LongTensor] = None,1186 use_cache: Optional[bool] = None,1187 output_hidden_states: Optional[bool] = None,1188 return_dict: Optional[bool] = None,1189 ) -> Union[Tuple[torch.Tensor, ...], SequenceClassifierOutputWithPast]:1190 return_dict = return_dict if return_dict is not None else self.config.use_return_dict1191 1192 transformer_outputs = self.transformer(1193 input_ids=input_ids,1194 position_ids=position_ids,1195 attention_mask=attention_mask,1196 full_attention_mask=full_attention_mask,1197 past_key_values=past_key_values,1198 inputs_embeds=inputs_embeds,1199 use_cache=use_cache,1200 output_hidden_states=output_hidden_states,