dummy-foo/ChatGLM3-Japanese-Zero
039
1""" PyTorch ChatGLM model. """
2
3import math
4import copy
5import warnings
6import re
7import sys
8
9import torch
10import torch.utils.checkpoint
11import torch.nn.functional as F
12from torch import nn
13from torch.nn import CrossEntropyLoss, LayerNorm, MSELoss, BCEWithLogitsLoss
14from torch.nn.utils import skip_init
15from typing import Optional, Tuple, Union, List, Callable, Dict, Any
16from copy import deepcopy
17
18from transformers.modeling_outputs import (
19 BaseModelOutputWithPast,
20 CausalLMOutputWithPast,
21 SequenceClassifierOutputWithPast,
22)
23from transformers.modeling_utils import PreTrainedModel
24from transformers.utils import logging
25from transformers.generation.logits_process import LogitsProcessor
26from transformers.generation.utils import LogitsProcessorList, StoppingCriteriaList, GenerationConfig, ModelOutput
27
28from .configuration_chatglm import ChatGLMConfig
29
30# flags required to enable jit fusion kernels
31
32if sys.platform != 'darwin':
33 torch._C._jit_set_profiling_mode(False)
34 torch._C._jit_set_profiling_executor(False)
35 torch._C._jit_override_can_fuse_on_cpu(True)
36 torch._C._jit_override_can_fuse_on_gpu(True)
37
38logger = logging.get_logger(__name__)
39
40_CHECKPOINT_FOR_DOC = "THUDM/ChatGLM"
41_CONFIG_FOR_DOC = "ChatGLMConfig"
42
43CHATGLM_6B_PRETRAINED_MODEL_ARCHIVE_LIST = [
44 "THUDM/chatglm3-6b",
45 # See all ChatGLM models at https://huggingface.co/models?filter=chatglm
46]
47
48
49def default_init(cls, *args, **kwargs):
50 return cls(*args, **kwargs)
51
52
53class InvalidScoreLogitsProcessor(LogitsProcessor):
54 def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:
55 if torch.isnan(scores).any() or torch.isinf(scores).any():
56 scores.zero_()
57 scores[..., 5] = 5e4
58 return scores
59
60
61class PrefixEncoder(torch.nn.Module):
62 """
63 The torch.nn model to encode the prefix
64 Input shape: (batch-size, prefix-length)
65 Output shape: (batch-size, prefix-length, 2*layers*hidden)
66 """
67
68 def __init__(self, config: ChatGLMConfig):
69 super().__init__()
70 self.prefix_projection = config.prefix_projection
71 if self.prefix_projection:
72 # Use a two-layer MLP to encode the prefix
73 kv_size = config.num_layers * config.kv_channels * config.multi_query_group_num * 2
74 self.embedding = torch.nn.Embedding(config.pre_seq_len, kv_size)
75 self.trans = torch.nn.Sequential(
76 torch.nn.Linear(kv_size, config.hidden_size),
77 torch.nn.Tanh(),
78 torch.nn.Linear(config.hidden_size, kv_size)
79 )
80 else:
81 self.embedding = torch.nn.Embedding(config.pre_seq_len,
82 config.num_layers * config.kv_channels * config.multi_query_group_num * 2)
83
84 def forward(self, prefix: torch.Tensor):
85 if self.prefix_projection:
86 prefix_tokens = self.embedding(prefix)
87 past_key_values = self.trans(prefix_tokens)
88 else:
89 past_key_values = self.embedding(prefix)
90 return past_key_values
91
92
93def split_tensor_along_last_dim(
94 tensor: torch.Tensor,
95 num_partitions: int,
96 contiguous_split_chunks: bool = False,
97) -> List[torch.Tensor]:
98 """Split a tensor along its last dimension.
99
100 Arguments:
101 tensor: input tensor.
102 num_partitions: number of partitions to split the tensor
103 contiguous_split_chunks: If True, make each chunk contiguous
104 in memory.
105
106 Returns:
107 A list of Tensors
108 """
109 # Get the size and dimension.
110 last_dim = tensor.dim() - 1
111 last_dim_size = tensor.size()[last_dim] // num_partitions
112 # Split.
113 tensor_list = torch.split(tensor, last_dim_size, dim=last_dim)
114 # Note: torch.split does not create contiguous tensors by default.
115 if contiguous_split_chunks:
116 return tuple(chunk.contiguous() for chunk in tensor_list)
117
118 return tensor_list
119
120
121class RotaryEmbedding(nn.Module):
122 def __init__(self, dim, original_impl=False, device=None, dtype=None):
123 super().__init__()
124 inv_freq = 1.0 / (10000 ** (torch.arange(0, dim, 2, device=device).to(dtype=dtype) / dim))
125 self.register_buffer("inv_freq", inv_freq)
126 self.dim = dim
127 self.original_impl = original_impl
128
129 def forward_impl(
130 self, seq_len: int, n_elem: int, dtype: torch.dtype, device: torch.device, base: int = 10000
131 ):
132 """Enhanced Transformer with Rotary Position Embedding.
133
134 Derived from: https://github.com/labmlai/annotated_deep_learning_paper_implementations/blob/master/labml_nn/
135 transformers/rope/__init__.py. MIT License:
136 https://github.com/labmlai/annotated_deep_learning_paper_implementations/blob/master/license.
137 """
138 # $\Theta = {\theta_i = 10000^{\frac{2(i-1)}{d}}, i \in [1, 2, ..., \frac{d}{2}]}$
139 theta = 1.0 / (base ** (torch.arange(0, n_elem, 2, dtype=torch.float, device=device) / n_elem))
140
141 # Create position indexes `[0, 1, ..., seq_len - 1]`
142 seq_idx = torch.arange(seq_len, dtype=torch.float, device=device)
143
144 # Calculate the product of position index and $\theta_i$
145 idx_theta = torch.outer(seq_idx, theta).float()
146
147 cache = torch.stack([torch.cos(idx_theta), torch.sin(idx_theta)], dim=-1)
148
149 # this is to mimic the behaviour of complex32, else we will get different results
150 if dtype in (torch.float16, torch.bfloat16, torch.int8):
151 cache = cache.bfloat16() if dtype == torch.bfloat16 else cache.half()
152 return cache
153
154 def forward(self, max_seq_len, offset=0):
155 return self.forward_impl(
156 max_seq_len, self.dim, dtype=self.inv_freq.dtype, device=self.inv_freq.device
157 )
158
159
160@torch.jit.script
161def apply_rotary_pos_emb(x: torch.Tensor, rope_cache: torch.Tensor) -> torch.Tensor:
162 # x: [sq, b, np, hn]
163 sq, b, np, hn = x.size(0), x.size(1), x.size(2), x.size(3)
164 rot_dim = rope_cache.shape[-2] * 2
165 x, x_pass = x[..., :rot_dim], x[..., rot_dim:]
166 # truncate to support variable sizes
167 rope_cache = rope_cache[:sq]
168 xshaped = x.reshape(sq, -1, np, rot_dim // 2, 2)
169 rope_cache = rope_cache.view(sq, -1, 1, xshaped.size(3), 2)
170 x_out2 = torch.stack(
171 [
172 xshaped[..., 0] * rope_cache[..., 0] - xshaped[..., 1] * rope_cache[..., 1],
173 xshaped[..., 1] * rope_cache[..., 0] + xshaped[..., 0] * rope_cache[..., 1],
174 ],
175 -1,
176 )
177 x_out2 = x_out2.flatten(3)
178 return torch.cat((x_out2, x_pass), dim=-1)
179
180
181class RMSNorm(torch.nn.Module):
182 def __init__(self, normalized_shape, eps=1e-5, device=None, dtype=None, **kwargs):
183 super().__init__()
184 self.weight = torch.nn.Parameter(torch.empty(normalized_shape, device=device, dtype=dtype))
185 self.eps = eps
186
187 def forward(self, hidden_states: torch.Tensor):
188 input_dtype = hidden_states.dtype
189 variance = hidden_states.to(torch.float32).pow(2).mean(-1, keepdim=True)
190 hidden_states = hidden_states * torch.rsqrt(variance + self.eps)
191
192 return (self.weight * hidden_states).to(input_dtype)
193
194
195class CoreAttention(torch.nn.Module):
196 def __init__(self, config: ChatGLMConfig, layer_number):
197 super(CoreAttention, self).__init__()
198
199 self.apply_query_key_layer_scaling = config.apply_query_key_layer_scaling
200 self.attention_softmax_in_fp32 = config.attention_softmax_in_fp32
201 if self.apply_query_key_layer_scaling:
202 self.attention_softmax_in_fp32 = True
203 self.layer_number = max(1, layer_number)
204
205 projection_size = config.kv_channels * config.num_attention_heads
206
207 # Per attention head and per partition values.
208 self.hidden_size_per_partition = projection_size
209 self.hidden_size_per_attention_head = projection_size // config.num_attention_heads
210 self.num_attention_heads_per_partition = config.num_attention_heads
211
212 coeff = None
213 self.norm_factor = math.sqrt(self.hidden_size_per_attention_head)
214 if self.apply_query_key_layer_scaling:
215 coeff = self.layer_number
216 self.norm_factor *= coeff
217 self.coeff = coeff
218
219 self.attention_dropout = torch.nn.Dropout(config.attention_dropout)
220
221 def forward(self, query_layer, key_layer, value_layer, attention_mask):
222 pytorch_major_version = int(torch.__version__.split('.')[0])
223 if pytorch_major_version >= 2:
224 query_layer, key_layer, value_layer = [k.permute(1, 2, 0, 3) for k in [query_layer, key_layer, value_layer]]
225 if attention_mask is None and query_layer.shape[2] == key_layer.shape[2]:
226 context_layer = torch.nn.functional.scaled_dot_product_attention(query_layer, key_layer, value_layer,
227 is_causal=True)
228 else:
229 if attention_mask is not None:
230 attention_mask = ~attention_mask
231 context_layer = torch.nn.functional.scaled_dot_product_attention(query_layer, key_layer, value_layer,
232 attention_mask)
233 context_layer = context_layer.permute(2, 0, 1, 3)
234 new_context_layer_shape = context_layer.size()[:-2] + (self.hidden_size_per_partition,)
235 context_layer = context_layer.reshape(*new_context_layer_shape)
236 else:
237 # Raw attention scores
238
239 # [b, np, sq, sk]
240 output_size = (query_layer.size(1), query_layer.size(2), query_layer.size(0), key_layer.size(0))
241
242 # [sq, b, np, hn] -> [sq, b * np, hn]
243 query_layer = query_layer.view(output_size[2], output_size[0] * output_size[1], -1)
244 # [sk, b, np, hn] -> [sk, b * np, hn]
245 key_layer = key_layer.view(output_size[3], output_size[0] * output_size[1], -1)
246
247 # preallocting input tensor: [b * np, sq, sk]
248 matmul_input_buffer = torch.empty(
249 output_size[0] * output_size[1], output_size[2], output_size[3], dtype=query_layer.dtype,
250 device=query_layer.device
251 )
252
253 # Raw attention scores. [b * np, sq, sk]
254 matmul_result = torch.baddbmm(
255 matmul_input_buffer,
256 query_layer.transpose(0, 1), # [b * np, sq, hn]
257 key_layer.transpose(0, 1).transpose(1, 2), # [b * np, hn, sk]
258 beta=0.0,
259 alpha=(1.0 / self.norm_factor),
260 )
261
262 # change view to [b, np, sq, sk]
263 attention_scores = matmul_result.view(*output_size)
264
265 # ===========================
266 # Attention probs and dropout
267 # ===========================
268
269 # attention scores and attention mask [b, np, sq, sk]
270 if self.attention_softmax_in_fp32:
271 attention_scores = attention_scores.float()
272 if self.coeff is not None:
273 attention_scores = attention_scores * self.coeff
274 if attention_mask is None and attention_scores.shape[2] == attention_scores.shape[3]:
275 attention_mask = torch.ones(output_size[0], 1, output_size[2], output_size[3],
276 device=attention_scores.device, dtype=torch.bool)
277 attention_mask.tril_()
278 attention_mask = ~attention_mask
279 if attention_mask is not None:
280 attention_scores = attention_scores.masked_fill(attention_mask, float("-inf"))
281 attention_probs = F.softmax(attention_scores, dim=-1)
282 attention_probs = attention_probs.type_as(value_layer)
283
284 # This is actually dropping out entire tokens to attend to, which might
285 # seem a bit unusual, but is taken from the original Transformer paper.
286 attention_probs = self.attention_dropout(attention_probs)
287 # =========================
288 # Context layer. [sq, b, hp]
289 # =========================
290
291 # value_layer -> context layer.
292 # [sk, b, np, hn] --> [b, np, sq, hn]
293
294 # context layer shape: [b, np, sq, hn]
295 output_size = (value_layer.size(1), value_layer.size(2), query_layer.size(0), value_layer.size(3))
296 # change view [sk, b * np, hn]
297 value_layer = value_layer.view(value_layer.size(0), output_size[0] * output_size[1], -1)
298 # change view [b * np, sq, sk]
299 attention_probs = attention_probs.view(output_size[0] * output_size[1], output_size[2], -1)
300 # matmul: [b * np, sq, hn]
301 context_layer = torch.bmm(attention_probs, value_layer.transpose(0, 1))
302 # change view [b, np, sq, hn]
303 context_layer = context_layer.view(*output_size)
304 # [b, np, sq, hn] --> [sq, b, np, hn]
305 context_layer = context_layer.permute(2, 0, 1, 3).contiguous()
306 # [sq, b, np, hn] --> [sq, b, hp]
307 new_context_layer_shape = context_layer.size()[:-2] + (self.hidden_size_per_partition,)
308 context_layer = context_layer.view(*new_context_layer_shape)
309
310 return context_layer
311
312
313class SelfAttention(torch.nn.Module):
314 """Parallel self-attention layer abstract class.
315
316 Self-attention layer takes input with size [s, b, h]
317 and returns output of the same size.
318 """
319
320 def __init__(self, config: ChatGLMConfig, layer_number, device=None):
321 super(SelfAttention, self).__init__()
322 self.layer_number = max(1, layer_number)
323
324 self.projection_size = config.kv_channels * config.num_attention_heads
325
326 # Per attention head and per partition values.
327 self.hidden_size_per_attention_head = self.projection_size // config.num_attention_heads
328 self.num_attention_heads_per_partition = config.num_attention_heads
329
330 self.multi_query_attention = config.multi_query_attention
331 self.qkv_hidden_size = 3 * self.projection_size
332 if self.multi_query_attention:
333 self.num_multi_query_groups_per_partition = config.multi_query_group_num
334 self.qkv_hidden_size = (
335 self.projection_size + 2 * self.hidden_size_per_attention_head * config.multi_query_group_num
336 )
337 self.query_key_value = nn.Linear(config.hidden_size, self.qkv_hidden_size,
338 bias=config.add_bias_linear or config.add_qkv_bias,
339 device=device, **_config_to_kwargs(config)
340 )
341
342 self.core_attention = CoreAttention(config, self.layer_number)
343
344 # Output.
345 self.dense = nn.Linear(self.projection_size, config.hidden_size, bias=config.add_bias_linear,
346 device=device, **_config_to_kwargs(config)
347 )
348
349 def _allocate_memory(self, inference_max_sequence_len, batch_size, device=None, dtype=None):
350 if self.multi_query_attention:
351 num_attention_heads = self.num_multi_query_groups_per_partition
352 else:
353 num_attention_heads = self.num_attention_heads_per_partition
354 return torch.empty(
355 inference_max_sequence_len,
356 batch_size,
357 num_attention_heads,
358 self.hidden_size_per_attention_head,
359 dtype=dtype,
360 device=device,
361 )
362
363 def forward(
364 self, hidden_states, attention_mask, rotary_pos_emb, kv_cache=None, use_cache=True
365 ):
366 # hidden_states: [sq, b, h]
367
368 # =================================================
369 # Pre-allocate memory for key-values for inference.
370 # =================================================
371 # =====================
372 # Query, Key, and Value
373 # =====================
374
375 # Attention heads [sq, b, h] --> [sq, b, (np * 3 * hn)]
376 mixed_x_layer = self.query_key_value(hidden_states)
377
378 if self.multi_query_attention:
379 (query_layer, key_layer, value_layer) = mixed_x_layer.split(
380 [
381 self.num_attention_heads_per_partition * self.hidden_size_per_attention_head,
382 self.num_multi_query_groups_per_partition * self.hidden_size_per_attention_head,
383 self.num_multi_query_groups_per_partition * self.hidden_size_per_attention_head,
384 ],
385 dim=-1,
386 )
387 query_layer = query_layer.view(
388 query_layer.size()[:-1] + (self.num_attention_heads_per_partition, self.hidden_size_per_attention_head)
389 )
390 key_layer = key_layer.view(
391 key_layer.size()[:-1] + (self.num_multi_query_groups_per_partition, self.hidden_size_per_attention_head)
392 )
393 value_layer = value_layer.view(
394 value_layer.size()[:-1]
395 + (self.num_multi_query_groups_per_partition, self.hidden_size_per_attention_head)
396 )
397 else:
398 new_tensor_shape = mixed_x_layer.size()[:-1] + \
399 (self.num_attention_heads_per_partition,
400 3 * self.hidden_size_per_attention_head)
401 mixed_x_layer = mixed_x_layer.view(*new_tensor_shape)
402
403 # [sq, b, np, 3 * hn] --> 3 [sq, b, np, hn]
404 (query_layer, key_layer, value_layer) = split_tensor_along_last_dim(mixed_x_layer, 3)
405
406 # apply relative positional encoding (rotary embedding)
407 if rotary_pos_emb is not None:
408 query_layer = apply_rotary_pos_emb(query_layer, rotary_pos_emb)
409 key_layer = apply_rotary_pos_emb(key_layer, rotary_pos_emb)
410
411 # adjust key and value for inference
412 if kv_cache is not None:
413 cache_k, cache_v = kv_cache
414 key_layer = torch.cat((cache_k, key_layer), dim=0)
415 value_layer = torch.cat((cache_v, value_layer), dim=0)
416 if use_cache:
417 kv_cache = (key_layer, value_layer)
418 else:
419 kv_cache = None
420
421 if self.multi_query_attention:
422 key_layer = key_layer.unsqueeze(-2)
423 key_layer = key_layer.expand(
424 -1, -1, -1, self.num_attention_heads_per_partition // self.num_multi_query_groups_per_partition, -1
425 )
426 key_layer = key_layer.contiguous().view(
427 key_layer.size()[:2] + (self.num_attention_heads_per_partition, self.hidden_size_per_attention_head)
428 )
429 value_layer = value_layer.unsqueeze(-2)
430 value_layer = value_layer.expand(
431 -1, -1, -1, self.num_attention_heads_per_partition // self.num_multi_query_groups_per_partition, -1
432 )
433 value_layer = value_layer.contiguous().view(
434 value_layer.size()[:2] + (self.num_attention_heads_per_partition, self.hidden_size_per_attention_head)
435 )
436
437 # ==================================
438 # core attention computation
439 # ==================================
440
441 context_layer = self.core_attention(query_layer, key_layer, value_layer, attention_mask)
442
443 # =================
444 # Output. [sq, b, h]
445 # =================
446
447 output = self.dense(context_layer)
448
449 return output, kv_cache
450
451
452def _config_to_kwargs(args):
453 common_kwargs = {
454 "dtype": args.torch_dtype,
455 }
456 return common_kwargs
457
458
459class MLP(torch.nn.Module):
460 """MLP.
461
462 MLP will take the input with h hidden state, project it to 4*h
463 hidden dimension, perform nonlinear transformation, and project the
464 state back into h hidden dimension.
465 """
466
467 def __init__(self, config: ChatGLMConfig, device=None):
468 super(MLP, self).__init__()
469
470 self.add_bias = config.add_bias_linear
471
472 # Project to 4h. If using swiglu double the output width, see https://arxiv.org/pdf/2002.05202.pdf
473 self.dense_h_to_4h = nn.Linear(
474 config.hidden_size,
475 config.ffn_hidden_size * 2,
476 bias=self.add_bias,
477 device=device,
478 **_config_to_kwargs(config)
479 )
480
481 def swiglu(x):
482 x = torch.chunk(x, 2, dim=-1)
483 return F.silu(x[0]) * x[1]
484
485 self.activation_func = swiglu
486
487 # Project back to h.
488 self.dense_4h_to_h = nn.Linear(
489 config.ffn_hidden_size,
490 config.hidden_size,
491 bias=self.add_bias,
492 device=device,
493 **_config_to_kwargs(config)
494 )
495
496 def forward(self, hidden_states):
497 # [s, b, 4hp]
498 intermediate_parallel = self.dense_h_to_4h(hidden_states)
499 intermediate_parallel = self.activation_func(intermediate_parallel)
500 # [s, b, h]
501 output = self.dense_4h_to_h(intermediate_parallel)
502 return output
503
504
505class GLMBlock(torch.nn.Module):
506 """A single transformer layer.
507
508 Transformer layer takes input with size [s, b, h] and returns an
509 output of the same size.
510 """
511
512 def __init__(self, config: ChatGLMConfig, layer_number, device=None):
513 super(GLMBlock, self).__init__()
514 self.layer_number = layer_number
515
516 self.apply_residual_connection_post_layernorm = config.apply_residual_connection_post_layernorm
517
518 self.fp32_residual_connection = config.fp32_residual_connection
519
520 LayerNormFunc = RMSNorm if config.rmsnorm else LayerNorm
521 # Layernorm on the input data.
522 self.input_layernorm = LayerNormFunc(config.hidden_size, eps=config.layernorm_epsilon, device=device,
523 dtype=config.torch_dtype)
524
525 # Self attention.
526 self.self_attention = SelfAttention(config, layer_number, device=device)
527 self.hidden_dropout = config.hidden_dropout
528
529 # Layernorm on the attention output
530 self.post_attention_layernorm = LayerNormFunc(config.hidden_size, eps=config.layernorm_epsilon, device=device,
531 dtype=config.torch_dtype)
532
533 # MLP
534 self.mlp = MLP(config, device=device)
535
536 def forward(
537 self, hidden_states, attention_mask, rotary_pos_emb, kv_cache=None, use_cache=True,
538 ):
539 # hidden_states: [s, b, h]
540
541 # Layer norm at the beginning of the transformer layer.
542 layernorm_output = self.input_layernorm(hidden_states)
543 # Self attention.
544 attention_output, kv_cache = self.self_attention(
545 layernorm_output,
546 attention_mask,
547 rotary_pos_emb,
548 kv_cache=kv_cache,
549 use_cache=use_cache
550 )
551
552 # Residual connection.
553 if self.apply_residual_connection_post_layernorm:
554 residual = layernorm_output
555 else:
556 residual = hidden_states
557
558 layernorm_input = torch.nn.functional.dropout(attention_output, p=self.hidden_dropout, training=self.training)
559 layernorm_input = residual + layernorm_input
560
561 # Layer norm post the self attention.
562 layernorm_output = self.post_attention_layernorm(layernorm_input)
563
564 # MLP.
565 mlp_output = self.mlp(layernorm_output)
566
567 # Second residual connection.
568 if self.apply_residual_connection_post_layernorm:
569 residual = layernorm_output
570 else:
571 residual = layernorm_input
572
573 output = torch.nn.functional.dropout(mlp_output, p=self.hidden_dropout, training=self.training)
574 output = residual + output
575
576 return output, kv_cache
577
578
579class GLMTransformer(torch.nn.Module):
580 """Transformer class."""
581
582 def __init__(self, config: ChatGLMConfig, device=None):
583 super(GLMTransformer, self).__init__()
584
585 self.fp32_residual_connection = config.fp32_residual_connection
586 self.post_layer_norm = config.post_layer_norm
587
588 # Number of layers.
589 self.num_layers = config.num_layers
590
591 # Transformer layers.
592 def build_layer(layer_number):
593 return GLMBlock(config, layer_number, device=device)
594
595 self.layers = torch.nn.ModuleList([build_layer(i + 1) for i in range(self.num_layers)])
596
597 if self.post_layer_norm:
598 LayerNormFunc = RMSNorm if config.rmsnorm else LayerNorm
599 # Final layer norm before output.
600 self.final_layernorm = LayerNormFunc(config.hidden_size, eps=config.layernorm_epsilon, device=device,
601 dtype=config.torch_dtype)
602
603 self.gradient_checkpointing = False
604
605 def _get_layer(self, layer_number):
606 return self.layers[layer_number]
607
608 def forward(
609 self, hidden_states, attention_mask, rotary_pos_emb, kv_caches=None,
610 use_cache: Optional[bool] = True,
611 output_hidden_states: Optional[bool] = False,
612 ):
613 if not kv_caches:
614 kv_caches = [None for _ in range(self.num_layers)]
615 presents = () if use_cache else None
616 if self.gradient_checkpointing and self.training:
617 if use_cache:
618 logger.warning_once(
619 "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
620 )
621 use_cache = False
622
623 all_self_attentions = None
624 all_hidden_states = () if output_hidden_states else None
625 for index in range(self.num_layers):
626 if output_hidden_states:
627 all_hidden_states = all_hidden_states + (hidden_states,)
628
629 layer = self._get_layer(index)
630 if self.gradient_checkpointing and self.training:
631 layer_ret = torch.utils.checkpoint.checkpoint(
632 layer,
633 hidden_states,
634 attention_mask,
635 rotary_pos_emb,
636 kv_caches[index],
637 use_cache,
638 use_reentrant=False
639 )
640 else:
641 layer_ret = layer(
642 hidden_states,
643 attention_mask,
644 rotary_pos_emb,
645 kv_cache=kv_caches[index],
646 use_cache=use_cache
647 )
648 hidden_states, kv_cache = layer_ret
649 if use_cache:
650 presents = presents + (kv_cache,)
651
652 if output_hidden_states:
653 all_hidden_states = all_hidden_states + (hidden_states,)
654
655 # Final layer norm.
656 if self.post_layer_norm:
657 hidden_states = self.final_layernorm(hidden_states)
658
659 return hidden_states, presents, all_hidden_states, all_self_attentions
660
661
662class ChatGLMPreTrainedModel(PreTrainedModel):
663 """
664 An abstract class to handle weights initialization and
665 a simple interface for downloading and loading pretrained models.
666 """
667
668 is_parallelizable = False
669 supports_gradient_checkpointing = True
670 config_class = ChatGLMConfig
671 base_model_prefix = "transformer"
672 _no_split_modules = ["GLMBlock"]
673
674 def _init_weights(self, module: nn.Module):
675 """Initialize the weights."""
676 return
677
678 def get_masks(self, input_ids, past_key_values, padding_mask=None):
679 batch_size, seq_length = input_ids.shape
680 full_attention_mask = torch.ones(batch_size, seq_length, seq_length, device=input_ids.device)
681 full_attention_mask.tril_()
682 past_length = 0
683 if past_key_values:
684 past_length = past_key_values[0][0].shape[0]
685 if past_length:
686 full_attention_mask = torch.cat((torch.ones(batch_size, seq_length, past_length,
687 device=input_ids.device), full_attention_mask), dim=-1)
688 if padding_mask is not None:
689 full_attention_mask = full_attention_mask * padding_mask.unsqueeze(1)
690 if not past_length and padding_mask is not None:
691 full_attention_mask -= padding_mask.unsqueeze(-1) - 1
692 full_attention_mask = (full_attention_mask < 0.5).bool()
693 full_attention_mask.unsqueeze_(1)
694 return full_attention_mask
695
696 def get_position_ids(self, input_ids, device):
697 batch_size, seq_length = input_ids.shape
698 position_ids = torch.arange(seq_length, dtype=torch.long, device=device).unsqueeze(0).repeat(batch_size, 1)
699 return position_ids
700
701 def gradient_checkpointing_enable(self, gradient_checkpointing_kwargs=None):
702 if not self.supports_gradient_checkpointing:
703 raise ValueError(f"{self.__class__.__name__} does not support gradient checkpointing.")
704
705
706class Embedding(torch.nn.Module):
707 """Language model embeddings."""
708
709 def __init__(self, config: ChatGLMConfig, device=None):
710 super(Embedding, self).__init__()
711
712 self.hidden_size = config.hidden_size
713 # Word embeddings (parallel).
714 self.word_embeddings = nn.Embedding(
715 config.padded_vocab_size,
716 self.hidden_size,
717 dtype=config.torch_dtype,
718 device=device
719 )
720 self.fp32_residual_connection = config.fp32_residual_connection
721
722 def forward(self, input_ids):
723 # Embeddings.
724 words_embeddings = self.word_embeddings(input_ids)
725 embeddings = words_embeddings
726 # Data format change to avoid explicit tranposes : [b s h] --> [s b h].
727 embeddings = embeddings.transpose(0, 1).contiguous()
728 # If the input flag for fp32 residual connection is set, convert for float.
729 if self.fp32_residual_connection:
730 embeddings = embeddings.float()
731 return embeddings
732
733
734class ChatGLMModel(ChatGLMPreTrainedModel):
735 def __init__(self, config: ChatGLMConfig, device=None, empty_init=True):
736 super().__init__(config)
737 if empty_init:
738 init_method = skip_init
739 else:
740 init_method = default_init
741 init_kwargs = {}
742 if device is not None:
743 init_kwargs["device"] = device
744 self.embedding = init_method(Embedding, config, **init_kwargs)
745 self.num_layers = config.num_layers
746 self.multi_query_group_num = config.multi_query_group_num
747 self.kv_channels = config.kv_channels
748
749 # Rotary positional embeddings
750 self.seq_length = config.seq_length
751 rotary_dim = (
752 config.hidden_size // config.num_attention_heads if config.kv_channels is None else config.kv_channels
753 )
754
755 self.rotary_pos_emb = RotaryEmbedding(rotary_dim // 2, original_impl=config.original_rope, device=device,
756 dtype=config.torch_dtype)
757 self.encoder = init_method(GLMTransformer, config, **init_kwargs)
758 self.output_layer = init_method(nn.Linear, config.hidden_size, config.padded_vocab_size, bias=False,
759 dtype=config.torch_dtype, **init_kwargs)
760 self.pre_seq_len = config.pre_seq_len
761 self.prefix_projection = config.prefix_projection
762 if self.pre_seq_len is not None:
763 for param in self.parameters():
764 param.requires_grad = False
765 self.prefix_tokens = torch.arange(self.pre_seq_len).long()
766 self.prefix_encoder = PrefixEncoder(config)
767 self.dropout = torch.nn.Dropout(0.1)
768
769 def get_input_embeddings(self):
770 return self.embedding.word_embeddings
771
772 def set_input_embeddings(self, value):
773 self.embedding.word_embeddings = value
774
775 def get_prompt(self, batch_size, device, dtype=torch.half):
776 prefix_tokens = self.prefix_tokens.unsqueeze(0).expand(batch_size, -1).to(device)
777 past_key_values = self.prefix_encoder(prefix_tokens).type(dtype)
778 past_key_values = past_key_values.view(
779 batch_size,
780 self.pre_seq_len,
781 self.num_layers * 2,
782 self.multi_query_group_num,
783 self.kv_channels
784 )
785 # seq_len, b, nh, hidden_size
786 past_key_values = self.dropout(past_key_values)
787 past_key_values = past_key_values.permute([2, 1, 0, 3, 4]).split(2)
788 return past_key_values
789
790 def forward(
791 self,
792 input_ids,
793 position_ids: Optional[torch.Tensor] = None,
794 attention_mask: Optional[torch.BoolTensor] = None,
795 full_attention_mask: Optional[torch.BoolTensor] = None,
796 past_key_values: Optional[Tuple[Tuple[torch.Tensor, torch.Tensor], ...]] = None,
797 inputs_embeds: Optional[torch.Tensor] = None,
798 use_cache: Optional[bool] = None,
799 output_hidden_states: Optional[bool] = None,
800 return_dict: Optional[bool] = None,
801 ):
802 output_hidden_states = (
803 output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
804 )
805 use_cache = use_cache if use_cache is not None else self.config.use_cache
806 return_dict = return_dict if return_dict is not None else self.config.use_return_dict
807
808 batch_size, seq_length = input_ids.shape
809
810 if inputs_embeds is None:
811 inputs_embeds = self.embedding(input_ids)
812
813 if self.pre_seq_len is not None:
814 if past_key_values is None:
815 past_key_values = self.get_prompt(batch_size=batch_size, device=input_ids.device,
816 dtype=inputs_embeds.dtype)
817 if attention_mask is not None:
818 attention_mask = torch.cat([attention_mask.new_ones((batch_size, self.pre_seq_len)),
819 attention_mask], dim=-1)
820
821 if full_attention_mask is None:
822 if (attention_mask is not None and not attention_mask.all()) or (past_key_values and seq_length != 1):
823 full_attention_mask = self.get_masks(input_ids, past_key_values, padding_mask=attention_mask)
824
825 # Rotary positional embeddings
826 rotary_pos_emb = self.rotary_pos_emb(self.seq_length)
827 if position_ids is not None:
828 rotary_pos_emb = rotary_pos_emb[position_ids]
829 else:
830 rotary_pos_emb = rotary_pos_emb[None, :seq_length]
831 rotary_pos_emb = rotary_pos_emb.transpose(0, 1).contiguous()
832
833 # Run encoder.
834 hidden_states, presents, all_hidden_states, all_self_attentions = self.encoder(
835 inputs_embeds, full_attention_mask, rotary_pos_emb=rotary_pos_emb,
836 kv_caches=past_key_values, use_cache=use_cache, output_hidden_states=output_hidden_states
837 )
838
839 if not return_dict:
840 return tuple(v for v in [hidden_states, presents, all_hidden_states, all_self_attentions] if v is not None)
841
842 return BaseModelOutputWithPast(
843 last_hidden_state=hidden_states,
844 past_key_values=presents,
845 hidden_states=all_hidden_states,
846 attentions=all_self_attentions,
847 )
848
849 def quantize(self, weight_bit_width: int):
850 from .quantization import quantize
851 quantize(self.encoder, weight_bit_width)
852 return self
853
854
855class ChatGLMForConditionalGeneration(ChatGLMPreTrainedModel):
856 def __init__(self, config: ChatGLMConfig, empty_init=True, device=None):
857 super().__init__(config)
858
859 self.max_sequence_length = config.max_length
860 self.transformer = ChatGLMModel(config, empty_init=empty_init, device=device)
861 self.config = config
862 self.quantized = False
863
864 if self.config.quantization_bit:
865 self.quantize(self.config.quantization_bit, empty_init=True)
866
867 def get_input_embeddings(self): # add for resize embedding
868 return self.transformer.embedding.word_embeddings
869
870 def set_input_embeddings(self, value): # add for resize embedding
871 self.transformer.embedding.word_embeddings = value
872
873 def get_output_embeddings(self): # add for resize embedding
874 return self.transformer.output_layer
875
876 def set_output_embeddings(self, value): # add for resize embedding
877 self.transformer.output_layer = value
878
879 def _update_model_kwargs_for_generation(
880 self,
881 outputs: ModelOutput,
882 model_kwargs: Dict[str, Any],
883 is_encoder_decoder: bool = False,
884 standardize_cache_format: bool = False,
885 ) -> Dict[str, Any]:
886 # update past_key_values
887 model_kwargs["past_key_values"] = self._extract_past_from_model_output(
888 outputs, standardize_cache_format=standardize_cache_format
889 )
890
891 # update attention mask
892 if "attention_mask" in model_kwargs:
893 attention_mask = model_kwargs["attention_mask"]
894 model_kwargs["attention_mask"] = torch.cat(
895 [attention_mask, attention_mask.new_ones((attention_mask.shape[0], 1))], dim=-1
896 )
897
898 # update position ids
899 if "position_ids" in model_kwargs:
900 position_ids = model_kwargs["position_ids"]
901 new_position_id = position_ids[..., -1:].clone()
902 new_position_id += 1
903 model_kwargs["position_ids"] = torch.cat(
904 [position_ids, new_position_id], dim=-1
905 )
906
907 model_kwargs["is_first_forward"] = False
908 return model_kwargs
909
910 def prepare_inputs_for_generation(
911 self,
912 input_ids: torch.LongTensor,
913 past_key_values: Optional[torch.Tensor] = None,
914 attention_mask: Optional[torch.Tensor] = None,
915 position_ids: Optional[torch.Tensor] = None,
916 use_cache: Optional[bool] = None,
917 is_first_forward: bool = True,
918 **kwargs
919 ) -> dict:
920 # only last token for input_ids if past is not None
921 if position_ids is None:
922 position_ids = self.get_position_ids(input_ids, device=input_ids.device)
923 if not is_first_forward:
924 if past_key_values is not None:
925 position_ids = position_ids[..., -1:]
926 input_ids = input_ids[:, -1:]
927 return {
928 "input_ids": input_ids,
929 "past_key_values": past_key_values,
930 "position_ids": position_ids,
931 "attention_mask": attention_mask,
932 "return_last_logit": True,
933 "use_cache": use_cache
934 }
935
936 def forward(
937 self,
938 input_ids: Optional[torch.Tensor] = None,
939 position_ids: Optional[torch.Tensor] = None,
940 attention_mask: Optional[torch.Tensor] = None,
941 past_key_values: Optional[Tuple[torch.FloatTensor]] = None,
942 inputs_embeds: Optional[torch.Tensor] = None,
943 labels: Optional[torch.Tensor] = None,
944 use_cache: Optional[bool] = None,
945 output_attentions: Optional[bool] = None,
946 output_hidden_states: Optional[bool] = None,
947 return_dict: Optional[bool] = None,
948 return_last_logit: Optional[bool] = False,
949 ):
950 use_cache = use_cache if use_cache is not None else self.config.use_cache
951 return_dict = return_dict if return_dict is not None else self.config.use_return_dict
952
953 transformer_outputs = self.transformer(
954 input_ids=input_ids,
955 position_ids=position_ids,
956 attention_mask=attention_mask,
957 past_key_values=past_key_values,
958 inputs_embeds=inputs_embeds,
959 use_cache=use_cache,
960 output_hidden_states=output_hidden_states,
961 return_dict=return_dict,
962 )
963
964 hidden_states = transformer_outputs[0]
965 if return_last_logit:
966 hidden_states = hidden_states[-1:]
967 lm_logits = self.transformer.output_layer(hidden_states)
968 lm_logits = lm_logits.transpose(0, 1).contiguous()
969
970 loss = None
971 if labels is not None:
972 lm_logits = lm_logits.to(torch.float32)
973
974 # Shift so that tokens < n predict n
975 shift_logits = lm_logits[..., :-1, :].contiguous()
976 shift_labels = labels[..., 1:].contiguous()
977 # Flatten the tokens
978 loss_fct = CrossEntropyLoss(ignore_index=-100)
979 loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))
980
981 lm_logits = lm_logits.to(hidden_states.dtype)
982 loss = loss.to(hidden_states.dtype)
983
984 if not return_dict:
985 output = (lm_logits,) + transformer_outputs[1:]
986 return ((loss,) + output) if loss is not None else output
987
988 return CausalLMOutputWithPast(
989 loss=loss,
990 logits=lm_logits,
991 past_key_values=transformer_outputs.past_key_values,
992 hidden_states=transformer_outputs.hidden_states,
993 attentions=transformer_outputs.attentions,
994 )
995
996 @staticmethod
997 def _reorder_cache(
998 past: Tuple[Tuple[torch.Tensor, torch.Tensor], ...], beam_idx: torch.LongTensor
999 ) -> Tuple[Tuple[torch.Tensor, torch.Tensor], ...]:
1000 """
1001 This function is used to re-order the `past_key_values` cache if [`~PreTrainedModel.beam_search`] or
1002 [`~PreTrainedModel.beam_sample`] is called. This is required to match `past_key_values` with the correct
1003 beam_idx at every generation step.
1004
1005 Output shares the same memory storage as `past`.
1006 """
1007 return tuple(
1008 (
1009 layer_past[0].index_select(1, beam_idx.to(layer_past[0].device)),
1010 layer_past[1].index_select(1, beam_idx.to(layer_past[1].device)),
1011 )
1012 for layer_past in past
1013 )
1014
1015 def process_response(self, output, history):
1016 content = ""
1017 history = deepcopy(history)
1018 for response in output.split("<|assistant|>"):
1019 if "\n" in response:
1020 metadata, content = response.split("\n", maxsplit=1)
1021 else:
1022 metadata, content = "", response
1023 if not metadata.strip():
1024 content = content.strip()
1025 history.append({"role": "assistant", "metadata": metadata, "content": content})
1026 content = content.replace("[[训练时间]]", "2023年")
1027 else:
1028 history.append({"role": "assistant", "metadata": metadata, "content": content})
1029 if history[0]["role"] == "system" and "tools" in history[0]:
1030 content = "\n".join(content.split("\n")[1:-1])
1031 def tool_call(**kwargs):
1032 return kwargs
1033 parameters = eval(content)
1034 content = {"name": metadata.strip(), "parameters": parameters}
1035 else:
1036 content = {"name": metadata.strip(), "content": content}
1037 return content, history
1038
1039 @torch.inference_mode()
1040 def chat(self, tokenizer, query: str, history: List[Dict] = None, role: str = "user",
1041 max_length: int = 8192, num_beams=1, do_sample=True, top_p=0.8, temperature=0.8, logits_processor=None,
1042 **kwargs):
1043 if history is None:
1044 history = []
1045 if logits_processor is None:
1046 logits_processor = LogitsProcessorList()
1047 logits_processor.append(InvalidScoreLogitsProcessor())
1048 gen_kwargs = {"max_length": max_length, "num_beams": num_beams, "do_sample": do_sample, "top_p": top_p,
1049 "temperature": temperature, "logits_processor": logits_processor, **kwargs}
1050 inputs = tokenizer.build_chat_input(query, history=history, role=role)
1051 inputs = inputs.to(self.device)
1052 eos_token_id = [tokenizer.eos_token_id, tokenizer.get_command("<|user|>"),
1053 tokenizer.get_command("<|observation|>")]
1054 outputs = self.generate(**inputs, **gen_kwargs, eos_token_id=eos_token_id)
1055 outputs = outputs.tolist()[0][len(inputs["input_ids"][0]):-1]
1056 response = tokenizer.decode(outputs)
1057 history.append({"role": role, "content": query})
1058 response, history = self.process_response(response, history)
1059 return response, history
1060
1061 @torch.inference_mode()
1062 def stream_chat(self, tokenizer, query: str, history: List[Dict] = None, role: str = "user",
1063 past_key_values=None,max_length: int = 8192, do_sample=True, top_p=0.8, temperature=0.8,
1064 logits_processor=None, return_past_key_values=False, **kwargs):
1065 if history is None:
1066 history = []
1067 if logits_processor is None:
1068 logits_processor = LogitsProcessorList()
1069 logits_processor.append(InvalidScoreLogitsProcessor())
1070 eos_token_id = [tokenizer.eos_token_id, tokenizer.get_command("<|user|>"),
1071 tokenizer.get_command("<|observation|>")]
1072 gen_kwargs = {"max_length": max_length, "do_sample": do_sample, "top_p": top_p,
1073 "temperature": temperature, "logits_processor": logits_processor, **kwargs}
1074 if past_key_values is None:
1075 inputs = tokenizer.build_chat_input(query, history=history, role=role)
1076 else:
1077 inputs = tokenizer.build_chat_input(query, role=role)
1078 inputs = inputs.to(self.device)
1079 if past_key_values is not None:
1080 past_length = past_key_values[0][0].shape[0]
1081 if self.transformer.pre_seq_len is not None:
1082 past_length -= self.transformer.pre_seq_len
1083 inputs.position_ids += past_length
1084 attention_mask = inputs.attention_mask
1085 attention_mask = torch.cat((attention_mask.new_ones(1, past_length), attention_mask), dim=1)
1086 inputs['attention_mask'] = attention_mask
1087 history.append({"role": role, "content": query})
1088 for outputs in self.stream_generate(**inputs, past_key_values=past_key_values,
1089 eos_token_id=eos_token_id, return_past_key_values=return_past_key_values,
1090 **gen_kwargs):
1091 if return_past_key_values:
1092 outputs, past_key_values = outputs
1093 outputs = outputs.tolist()[0][len(inputs["input_ids"][0]):-1]
1094 response = tokenizer.decode(outputs)
1095 if response and response[-1] != "�":
1096 response, new_history = self.process_response(response, history)
1097 if return_past_key_values:
1098 yield response, new_history, past_key_values
1099 else:
1100 yield response, new_history
1101
1102 @torch.inference_mode()
1103 def stream_generate(
1104 self,
1105 input_ids,
1106 generation_config: Optional[GenerationConfig] = None,
1107 logits_processor: Optional[LogitsProcessorList] = None,
1108 stopping_criteria: Optional[StoppingCriteriaList] = None,
1109 prefix_allowed_tokens_fn: Optional[Callable[[int, torch.Tensor], List[int]]] = None,
1110 return_past_key_values=False,
1111 **kwargs,
1112 ):
1113 batch_size, input_ids_seq_length = input_ids.shape[0], input_ids.shape[-1]
1114
1115 if generation_config is None:
1116 generation_config = self.generation_config
1117 generation_config = copy.deepcopy(generation_config)
1118 model_kwargs = generation_config.update(**kwargs)
1119 model_kwargs["use_cache"] = generation_config.use_cache
1120 bos_token_id, eos_token_id = generation_config.bos_token_id, generation_config.eos_token_id
1121
1122 if isinstance(eos_token_id, int):
1123 eos_token_id = [eos_token_id]
1124 eos_token_id_tensor = torch.tensor(eos_token_id).to(input_ids.device) if eos_token_id is not None else None
1125
1126 has_default_max_length = kwargs.get("max_length") is None and generation_config.max_length is not None
1127 if has_default_max_length and generation_config.max_new_tokens is None:
1128 warnings.warn(
1129 f"Using `max_length`'s default ({generation_config.max_length}) to control the generation length. "
1130 "This behaviour is deprecated and will be removed from the config in v5 of Transformers -- we"
1131 " recommend using `max_new_tokens` to control the maximum length of the generation.",
1132 UserWarning,
1133 )
1134 elif generation_config.max_new_tokens is not None:
1135 generation_config.max_length = generation_config.max_new_tokens + input_ids_seq_length
1136 if not has_default_max_length:
1137 logger.warn(
1138 f"Both `max_new_tokens` (={generation_config.max_new_tokens}) and `max_length`(="
1139 f"{generation_config.max_length}) seem to have been set. `max_new_tokens` will take precedence. "
1140 "Please refer to the documentation for more information. "
1141 "(https://huggingface.co/docs/transformers/main/en/main_classes/text_generation)",
1142 UserWarning,
1143 )
1144
1145 if input_ids_seq_length >= generation_config.max_length:
1146 input_ids_string = "decoder_input_ids" if self.config.is_encoder_decoder else "input_ids"
1147 logger.warning(
1148 f"Input length of {input_ids_string} is {input_ids_seq_length}, but `max_length` is set to"
1149 f" {generation_config.max_length}. This can lead to unexpected behavior. You should consider"
1150 " increasing `max_new_tokens`."
1151 )
1152
1153 # 2. Set generation parameters if not already defined
1154 logits_processor = logits_processor if logits_processor is not None else LogitsProcessorList()
1155 stopping_criteria = stopping_criteria if stopping_criteria is not None else StoppingCriteriaList()
1156
1157 logits_processor = self._get_logits_processor(
1158 generation_config=generation_config,
1159 input_ids_seq_length=input_ids_seq_length,
1160 encoder_input_ids=input_ids,
1161 prefix_allowed_tokens_fn=prefix_allowed_tokens_fn,
1162 logits_processor=logits_processor,
1163 )
1164
1165 stopping_criteria = self._get_stopping_criteria(
1166 generation_config=generation_config, stopping_criteria=stopping_criteria
1167 )
1168 logits_warper = self._get_logits_warper(generation_config)
1169
1170 unfinished_sequences = input_ids.new(input_ids.shape[0]).fill_(1)
1171 scores = None
1172 while True:
1173 model_inputs = self.prepare_inputs_for_generation(input_ids, **model_kwargs)
1174 # forward pass to get next token
1175 outputs = self(
1176 **model_inputs,
1177 return_dict=True,
1178 output_attentions=False,
1179 output_hidden_states=False,
1180 )
1181
1182 next_token_logits = outputs.logits[:, -1, :]
1183
1184 # pre-process distribution
1185 next_token_scores = logits_processor(input_ids, next_token_logits)
1186 next_token_scores = logits_warper(input_ids, next_token_scores)
1187
1188 # sample
1189 probs = nn.functional.softmax(next_token_scores, dim=-1)
1190 if generation_config.do_sample:
1191 next_tokens = torch.multinomial(probs, num_samples=1).squeeze(1)
1192 else:
1193 next_tokens = torch.argmax(probs, dim=-1)
1194 # update generated ids, model inputs, and length for next step
1195 input_ids = torch.cat([input_ids, next_tokens[:, None]], dim=-1)
1196 model_kwargs = self._update_model_kwargs_for_generation(
1197 outputs, model_kwargs, is_encoder_decoder=self.config.is_encoder_decoder
1198 )
1199 unfinished_sequences = unfinished_sequences.mul(
1200 next_tokens.tile(eos_token_id_tensor.shape[0], 1).ne(eos_token_id_tensor.unsqueeze(1)).prod(dim=0)