IlyasMoutawwakil HF Staff commited on
Commit
a57da11
·
verified ·
1 Parent(s): ccd7f03

Upload folder using huggingface_hub

Browse files
config.json ADDED
@@ -0,0 +1,50 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "_name_or_path": "/home/ea/work/my_optimum_intel/optimum-intel/tiny-glm4",
3
+ "add_bias_linear": false,
4
+ "add_qkv_bias": true,
5
+ "apply_query_key_layer_scaling": true,
6
+ "apply_residual_connection_post_layernorm": false,
7
+ "architectures": [
8
+ "ChatGLMForConditionalGeneration"
9
+ ],
10
+ "attention_dropout": 0.0,
11
+ "attention_softmax_in_fp32": true,
12
+ "auto_map": {
13
+ "AutoConfig": "configuration_chatglm.ChatGLMConfig",
14
+ "AutoModel": "modeling_chatglm.ChatGLMForConditionalGeneration",
15
+ "AutoModelForCausalLM": "modeling_chatglm.ChatGLMForConditionalGeneration",
16
+ "AutoModelForSeq2SeqLM": "modeling_chatglm.ChatGLMForConditionalGeneration",
17
+ "AutoModelForSequenceClassification": "modeling_chatglm.ChatGLMForSequenceClassification"
18
+ },
19
+ "bias_dropout_fusion": true,
20
+ "classifier_dropout": null,
21
+ "eos_token_id": [
22
+ 151329,
23
+ 151336,
24
+ 151338
25
+ ],
26
+ "ffn_hidden_size": 64,
27
+ "fp32_residual_connection": false,
28
+ "hidden_dropout": 0.0,
29
+ "hidden_size": 32,
30
+ "kv_channels": 8,
31
+ "layernorm_epsilon": 1.5625e-07,
32
+ "model_type": "chatglm",
33
+ "multi_query_attention": true,
34
+ "multi_query_group_num": 2,
35
+ "num_attention_heads": 4,
36
+ "num_hidden_layers": 6,
37
+ "num_layers": 6,
38
+ "original_rope": true,
39
+ "pad_token_id": 151329,
40
+ "padded_vocab_size": 151552,
41
+ "post_layer_norm": true,
42
+ "rmsnorm": true,
43
+ "rope_ratio": 500,
44
+ "seq_length": 131072,
45
+ "tie_word_embeddings": false,
46
+ "torch_dtype": "bfloat16",
47
+ "transformers_version": "4.40.2",
48
+ "use_cache": true,
49
+ "vocab_size": 151552
50
+ }
configuration_chatglm.py ADDED
@@ -0,0 +1,58 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from transformers import PretrainedConfig
2
+
3
+
4
+ class ChatGLMConfig(PretrainedConfig):
5
+ model_type = "chatglm"
6
+
7
+ def __init__(
8
+ self,
9
+ num_layers=28,
10
+ padded_vocab_size=65024,
11
+ hidden_size=4096,
12
+ ffn_hidden_size=13696,
13
+ kv_channels=128,
14
+ num_attention_heads=32,
15
+ seq_length=2048,
16
+ hidden_dropout=0.0,
17
+ classifier_dropout=None,
18
+ attention_dropout=0.0,
19
+ layernorm_epsilon=1e-5,
20
+ rmsnorm=True,
21
+ apply_residual_connection_post_layernorm=False,
22
+ post_layer_norm=True,
23
+ add_bias_linear=False,
24
+ add_qkv_bias=False,
25
+ bias_dropout_fusion=True,
26
+ multi_query_attention=False,
27
+ multi_query_group_num=1,
28
+ rope_ratio=1,
29
+ apply_query_key_layer_scaling=True,
30
+ attention_softmax_in_fp32=True,
31
+ fp32_residual_connection=False,
32
+ **kwargs
33
+ ):
34
+ self.num_layers = num_layers
35
+ self.vocab_size = padded_vocab_size
36
+ self.padded_vocab_size = padded_vocab_size
37
+ self.hidden_size = hidden_size
38
+ self.ffn_hidden_size = ffn_hidden_size
39
+ self.kv_channels = kv_channels
40
+ self.num_attention_heads = num_attention_heads
41
+ self.seq_length = seq_length
42
+ self.hidden_dropout = hidden_dropout
43
+ self.classifier_dropout = classifier_dropout
44
+ self.attention_dropout = attention_dropout
45
+ self.layernorm_epsilon = layernorm_epsilon
46
+ self.rmsnorm = rmsnorm
47
+ self.apply_residual_connection_post_layernorm = apply_residual_connection_post_layernorm
48
+ self.post_layer_norm = post_layer_norm
49
+ self.add_bias_linear = add_bias_linear
50
+ self.add_qkv_bias = add_qkv_bias
51
+ self.bias_dropout_fusion = bias_dropout_fusion
52
+ self.multi_query_attention = multi_query_attention
53
+ self.multi_query_group_num = multi_query_group_num
54
+ self.rope_ratio = rope_ratio
55
+ self.apply_query_key_layer_scaling = apply_query_key_layer_scaling
56
+ self.attention_softmax_in_fp32 = attention_softmax_in_fp32
57
+ self.fp32_residual_connection = fp32_residual_connection
58
+ super().__init__(**kwargs)
generation_config.json ADDED
@@ -0,0 +1,10 @@
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "_from_model_config": true,
3
+ "eos_token_id": [
4
+ 151329,
5
+ 151336,
6
+ 151338
7
+ ],
8
+ "pad_token_id": 151329,
9
+ "transformers_version": "4.40.2"
10
+ }
model.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:61f7e55caf88449dad559c05647d201f4dab35ab7bcb16d9dc504837d1ccb3a8
3
+ size 19516588
modeling_chatglm.py ADDED
@@ -0,0 +1,1243 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """ PyTorch ChatGLM model. """
2
+ import json
3
+ import math
4
+ import copy
5
+ import warnings
6
+ import re
7
+ import sys
8
+
9
+ import transformers
10
+ import torch
11
+ import torch.utils.checkpoint
12
+ import torch.nn.functional as F
13
+ from torch import nn
14
+ from torch.nn import CrossEntropyLoss, LayerNorm, MSELoss, BCEWithLogitsLoss
15
+ from torch.nn.utils import skip_init
16
+ from typing import Optional, Tuple, Union, List, Callable, Dict, Any
17
+ from copy import deepcopy
18
+
19
+ from transformers.modeling_outputs import (
20
+ BaseModelOutputWithPast,
21
+ CausalLMOutputWithPast,
22
+ SequenceClassifierOutputWithPast,
23
+ )
24
+ from transformers.modeling_utils import PreTrainedModel
25
+ from transformers.utils import logging, is_torch_npu_available
26
+ from transformers.generation.logits_process import LogitsProcessor
27
+ from transformers.generation.utils import LogitsProcessorList, StoppingCriteriaList, GenerationConfig, ModelOutput
28
+
29
+ from .configuration_chatglm import ChatGLMConfig
30
+
31
+ # flags required to enable jit fusion kernels
32
+
33
+ if 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
+
39
+ logger = logging.get_logger(__name__)
40
+
41
+ _CHECKPOINT_FOR_DOC = "THUDM/ChatGLM"
42
+ _CONFIG_FOR_DOC = "ChatGLMConfig"
43
+
44
+ is_transformers_4_42_or_higher = int(transformers.__version__.split(".")[1]) >= 42
45
+ is_transformers_4_44_or_higher = int(transformers.__version__.split(".")[1]) >= 44
46
+
47
+
48
+ def default_init(cls, *args, **kwargs):
49
+ return cls(*args, **kwargs)
50
+
51
+
52
+ class 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] = 5e4
57
+ return scores
58
+
59
+
60
+ def 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 tensor
70
+ contiguous_split_chunks: If True, make each chunk contiguous
71
+ in memory.
72
+
73
+ Returns:
74
+ A list of Tensors
75
+ """
76
+ # Get the size and dimension.
77
+ last_dim = tensor.dim() - 1
78
+ last_dim_size = tensor.size()[last_dim] // num_partitions
79
+ # 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_list
86
+
87
+
88
+ class 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 = dim
94
+ self.original_impl = original_impl
95
+ self.rope_ratio = rope_ratio
96
+
97
+ def forward_impl(
98
+ self, seq_len: int, n_elem: int, dtype: torch.dtype, device: torch.device, base: int = 10000
99
+ ):
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_ratio
108
+ 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 results
119
+ if dtype in (torch.float16, torch.bfloat16, torch.int8):
120
+ cache = cache.bfloat16() if dtype == torch.bfloat16 else cache.half()
121
+ return cache
122
+
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.device
126
+ )
127
+
128
+
129
+ @torch.jit.script
130
+ def 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] * 2
134
+ x, x_pass = x[..., :rot_dim], x[..., rot_dim:]
135
+ # truncate to support variable sizes
136
+ 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
+
150
+ class 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 = eps
155
+
156
+ def forward(self, hidden_states: torch.Tensor):
157
+ input_dtype = hidden_states.dtype
158
+ 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
+
164
+ class 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_scaling
169
+ self.attention_softmax_in_fp32 = config.attention_softmax_in_fp32
170
+ if self.apply_query_key_layer_scaling:
171
+ self.attention_softmax_in_fp32 = True
172
+ self.layer_number = max(1, layer_number)
173
+
174
+ projection_size = config.kv_channels * config.num_attention_heads
175
+
176
+ # Per attention head and per partition values.
177
+ self.hidden_size_per_partition = projection_size
178
+ self.hidden_size_per_attention_head = projection_size // config.num_attention_heads
179
+ self.num_attention_heads_per_partition = config.num_attention_heads
180
+
181
+ coeff = None
182
+ self.norm_factor = math.sqrt(self.hidden_size_per_attention_head)
183
+ if self.apply_query_key_layer_scaling:
184
+ coeff = self.layer_number
185
+ self.norm_factor *= coeff
186
+ self.coeff = coeff
187
+
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_mask
199
+ 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 scores
206
+
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.device
219
+ )
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 dropout
235
+ # ===========================
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.coeff
242
+ 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_mask
247
+ 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 might
253
+ # 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_layer
276
+
277
+
278
+ class 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_heads
290
+
291
+ # Per attention head and per partition values.
292
+ self.hidden_size_per_attention_head = self.projection_size // config.num_attention_heads
293
+ self.num_attention_heads_per_partition = config.num_attention_heads
294
+
295
+ self.multi_query_attention = config.multi_query_attention
296
+ self.qkv_hidden_size = 3 * self.projection_size
297
+ if self.multi_query_attention:
298
+ self.num_multi_query_groups_per_partition = config.multi_query_group_num
299
+ self.qkv_hidden_size = (
300
+ self.projection_size + 2 * self.hidden_size_per_attention_head * config.multi_query_group_num
301
+ )
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_partition
317
+ else:
318
+ num_attention_heads = self.num_attention_heads_per_partition
319
+ 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=True
330
+ ):
331
+ # hidden_states: [b, sq, h]
332
+
333
+ # =================================================
334
+ # Pre-allocate memory for key-values for inference.
335
+ # =================================================
336
+ # =====================
337
+ # Query, Key, and Value
338
+ # =====================
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 inference
380
+ if kv_cache is not None:
381
+ cache_k, cache_v = kv_cache
382
+ 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 = None
391
+
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, -1
396
+ )
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, -1
403
+ )
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 computation
410
+ # ==================================
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_cache
421
+
422
+
423
+ def _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_kwargs
428
+
429
+
430
+ class MLP(torch.nn.Module):
431
+ """MLP.
432
+
433
+ MLP will take the input with h hidden state, project it to 4*h
434
+ hidden dimension, perform nonlinear transformation, and project the
435
+ 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_linear
442
+
443
+ # Project to 4h. If using swiglu double the output width, see https://arxiv.org/pdf/2002.05202.pdf
444
+ 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 = swiglu
457
+
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 output
474
+
475
+
476
+ class GLMBlock(torch.nn.Module):
477
+ """A single transformer layer.
478
+
479
+ Transformer layer takes input with size [s, b, h] and returns an
480
+ 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_number
486
+
487
+ self.apply_residual_connection_post_layernorm = config.apply_residual_connection_post_layernorm
488
+
489
+ self.fp32_residual_connection = config.fp32_residual_connection
490
+
491
+ LayerNormFunc = RMSNorm if config.rmsnorm else LayerNorm
492
+ dtype = getattr(torch, config.torch_dtype) if isinstance(config.torch_dtype, str) else config.torch_dtype
493
+ # 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_dropout
500
+
501
+ # Layernorm on the attention output
502
+ self.post_attention_layernorm = LayerNormFunc(config.hidden_size, eps=config.layernorm_epsilon, device=device,
503
+ dtype=dtype)
504
+
505
+ # MLP
506
+ 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_cache
522
+ )
523
+
524
+ # Residual connection.
525
+ if self.apply_residual_connection_post_layernorm:
526
+ residual = layernorm_output
527
+ else:
528
+ residual = hidden_states
529
+
530
+ layernorm_input = torch.nn.functional.dropout(attention_output, p=self.hidden_dropout, training=self.training)
531
+ layernorm_input = residual + layernorm_input
532
+
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_output
542
+ else:
543
+ residual = layernorm_input
544
+
545
+ output = torch.nn.functional.dropout(mlp_output, p=self.hidden_dropout, training=self.training)
546
+ output = residual + output
547
+
548
+ return output, kv_cache
549
+
550
+
551
+ class 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_connection
558
+ self.post_layer_norm = config.post_layer_norm
559
+
560
+ # Number of layers.
561
+ self.num_layers = config.num_layers
562
+
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 LayerNorm
571
+ dtype = getattr(torch, config.torch_dtype) if isinstance(config.torch_dtype, str) else config.torch_dtype
572
+ # 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 = False
577
+
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 None
589
+ 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 = False
595
+
596
+ all_self_attentions = None
597
+ all_hidden_states = () if output_hidden_states else None
598
+ 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=False
612
+ )
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_cache
620
+ )
621
+ hidden_states, kv_cache = layer_ret
622
+ if use_cache:
623
+ # token by token decoding, use tuple format
624
+ if kv_caches[0] is not None:
625
+ presents = presents + (kv_cache,)
626
+ # prefilling in decoding, use tensor format to save cuda memory
627
+ else:
628
+ if len(presents) == 0:
629
+ presents = kv_cache
630
+ 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_attentions
641
+
642
+
643
+ class ChatGLMPreTrainedModel(PreTrainedModel):
644
+ """
645
+ An abstract class to handle weights initialization and
646
+ a simple interface for downloading and loading pretrained models.
647
+ """
648
+
649
+ is_parallelizable = False
650
+ supports_gradient_checkpointing = True
651
+ config_class = ChatGLMConfig
652
+ base_model_prefix = "transformer"
653
+ _no_split_modules = ["GLMBlock"]
654
+
655
+ def _init_weights(self, module: nn.Module):
656
+ """Initialize the weights."""
657
+ return
658
+
659
+ def get_masks(self, input_ids, past_key_values, padding_mask=None):
660
+ batch_size, seq_length = input_ids.shape
661
+ full_attention_mask = torch.ones(batch_size, seq_length, seq_length, device=input_ids.device)
662
+ full_attention_mask.tril_()
663
+ past_length = 0
664
+ 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) - 1
673
+ full_attention_mask = (full_attention_mask < 0.5).bool()
674
+ full_attention_mask.unsqueeze_(1)
675
+ return full_attention_mask
676
+
677
+ def get_position_ids(self, input_ids, device):
678
+ batch_size, seq_length = input_ids.shape
679
+ position_ids = torch.arange(seq_length, dtype=torch.long, device=device).unsqueeze(0).repeat(batch_size, 1)
680
+ return position_ids
681
+
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
+
687
+ class 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_size
694
+ # Word embeddings (parallel).
695
+ dtype = getattr(torch, config.torch_dtype) if isinstance(config.torch_dtype, str) else config.torch_dtype
696
+ self.word_embeddings = nn.Embedding(
697
+ config.padded_vocab_size,
698
+ self.hidden_size,
699
+ dtype=dtype,
700
+ device=device
701
+ )
702
+ self.fp32_residual_connection = config.fp32_residual_connection
703
+
704
+ def forward(self, input_ids):
705
+ # Embeddings.
706
+ words_embeddings = self.word_embeddings(input_ids)
707
+ embeddings = words_embeddings
708
+ # 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 embeddings
712
+
713
+
714
+ class 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_init
719
+ else:
720
+ init_method = default_init
721
+ 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_layers
726
+ self.multi_query_group_num = config.multi_query_group_num
727
+ self.kv_channels = config.kv_channels
728
+
729
+ # Rotary positional embeddings
730
+ self.seq_length = config.seq_length
731
+ rotary_dim = (
732
+ config.hidden_size // config.num_attention_heads if config.kv_channels is None else config.kv_channels
733
+ )
734
+ dtype = getattr(torch, config.torch_dtype) if isinstance(config.torch_dtype, str) else config.torch_dtype
735
+ 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_embeddings
743
+
744
+ def set_input_embeddings(self, value):
745
+ self.embedding.word_embeddings = value
746
+
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_states
761
+ )
762
+ use_cache = use_cache if use_cache is not None else self.config.use_cache
763
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
764
+
765
+ batch_size, seq_length = input_ids.shape
766
+
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 embeddings
775
+ 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_states
785
+ )
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
+
804
+ class 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_length
809
+ self.transformer = ChatGLMModel(config, empty_init=empty_init, device=device)
810
+ self.config = config
811
+
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
+ outputs
823
+ )[1]
824
+ # update past_key_values
825
+ 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_format
828
+ )[1]
829
+ else:
830
+ model_kwargs["past_key_values"] = self._extract_past_from_model_output(
831
+ outputs, standardize_cache_format=standardize_cache_format
832
+ )
833
+
834
+ # update attention mask
835
+ 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=-1
839
+ )
840
+
841
+ # update position ids
842
+ 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 += 1
846
+ model_kwargs["position_ids"] = torch.cat(
847
+ [position_ids, new_position_id], dim=-1
848
+ )
849
+
850
+ model_kwargs["is_first_forward"] = False
851
+ return model_kwargs
852
+
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
+ **kwargs
862
+ ) -> dict:
863
+ # only last token for input_ids if past is not None
864
+ 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_cache
877
+ }
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_cache
894
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
895
+
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 = None
913
+ if labels is not None:
914
+ lm_logits = lm_logits.to(torch.float32)
915
+
916
+ # Shift so that tokens < n predict n
917
+ shift_logits = lm_logits[..., :-1, :].contiguous()
918
+ shift_labels = labels[..., 1:].contiguous()
919
+ # Flatten the tokens
920
+ 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 output
929
+
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
+ @staticmethod
939
+ def _reorder_cache(
940
+ past: Tuple[Tuple[torch.Tensor, torch.Tensor], ...], beam_idx: torch.LongTensor
941
+ ) -> Tuple[Tuple[torch.Tensor, torch.Tensor], ...]:
942
+ """
943
+ This function is used to re-order the `past_key_values` cache if [`~PreTrainedModel.beam_search`] or
944
+ [`~PreTrainedModel.beam_sample`] is called. This is required to match `past_key_values` with the correct
945
+ 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 past
955
+ )
956
+
957
+ @staticmethod
958
+ def _extract_past_from_model_output(outputs: ModelOutput, *args, **kwargs):
959
+ past_key_values = None
960
+ if "past_key_values" in outputs:
961
+ past_key_values = outputs.past_key_values
962
+ if is_transformers_4_42_or_higher:
963
+ return None, past_key_values
964
+ return past_key_values
965
+
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 = "", response
974
+ 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, history
986
+
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, history
1009
+
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_length
1035
+ attention_mask = inputs.attention_mask
1036
+ attention_mask = torch.cat((attention_mask.new_ones(1, past_length), attention_mask), dim=1)
1037
+ inputs['attention_mask'] = attention_mask
1038
+ 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 = outputs
1044
+ 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_values
1050
+ else:
1051
+ yield response, new_history
1052
+
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_config
1068
+ generation_config = copy.deepcopy(generation_config)
1069
+ model_kwargs = generation_config.update(**kwargs)
1070
+ model_kwargs["use_cache"] = generation_config.use_cache
1071
+ bos_token_id, eos_token_id = generation_config.bos_token_id, generation_config.eos_token_id
1072
+
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 None
1076
+
1077
+ has_default_max_length = kwargs.get("max_length") is None and generation_config.max_length is not None
1078
+ 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_length
1087
+ 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 defined
1105
+ 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_criteria
1118
+ )
1119
+ logits_warper = self._get_logits_warper(generation_config)
1120
+
1121
+ unfinished_sequences = input_ids.new(input_ids.shape[0]).fill_(1)
1122
+ scores = None
1123
+ while True:
1124
+ model_inputs = self.prepare_inputs_for_generation(input_ids, **model_kwargs)
1125
+ # forward pass to get next token
1126
+ 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 distribution
1136
+ next_token_scores = logits_processor(input_ids, next_token_logits)
1137
+ next_token_scores = logits_warper(input_ids, next_token_scores)
1138
+
1139
+ # sample
1140
+ 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 step
1146
+ 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_decoder
1149
+ )
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_values
1155
+ else:
1156
+ yield input_ids
1157
+ # stop when each sentence is finished, or if we exceed the maximum length
1158
+ if unfinished_sequences.max() == 0 or stopping_criteria(input_ids, scores):
1159
+ break
1160
+
1161
+
1162
+ class ChatGLMForSequenceClassification(ChatGLMPreTrainedModel):
1163
+ def __init__(self, config: ChatGLMConfig, empty_init=True, device=None):
1164
+ super().__init__(config)
1165
+
1166
+ self.num_labels = config.num_labels
1167
+ 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_dtype
1169
+
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 = None
1175
+ self.config = config
1176
+
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_dict
1191
+
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,
1201
+ return_dict=return_dict,
1202
+ )
1203
+
1204
+ hidden_states = transformer_outputs[0]
1205
+ pooled_hidden_states = hidden_states[:, -1]
1206
+ if self.dropout is not None:
1207
+ pooled_hidden_states = self.dropout(pooled_hidden_states)
1208
+ logits = self.classifier_head(pooled_hidden_states)
1209
+
1210
+ loss = None
1211
+ if labels is not None:
1212
+ if self.config.problem_type is None:
1213
+ if self.num_labels == 1:
1214
+ self.config.problem_type = "regression"
1215
+ elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
1216
+ self.config.problem_type = "single_label_classification"
1217
+ else:
1218
+ self.config.problem_type = "multi_label_classification"
1219
+
1220
+ if self.config.problem_type == "regression":
1221
+ loss_fct = MSELoss()
1222
+ if self.num_labels == 1:
1223
+ loss = loss_fct(logits.squeeze().float(), labels.squeeze())
1224
+ else:
1225
+ loss = loss_fct(logits.float(), labels)
1226
+ elif self.config.problem_type == "single_label_classification":
1227
+ loss_fct = CrossEntropyLoss()
1228
+ loss = loss_fct(logits.view(-1, self.num_labels).float(), labels.view(-1))
1229
+ elif self.config.problem_type == "multi_label_classification":
1230
+ loss_fct = BCEWithLogitsLoss()
1231
+ loss = loss_fct(logits.float(), labels.view(-1, self.num_labels))
1232
+
1233
+ if not return_dict:
1234
+ output = (logits,) + transformer_outputs[1:]
1235
+ return ((loss,) + output) if loss is not None else output
1236
+
1237
+ return SequenceClassifierOutputWithPast(
1238
+ loss=loss,
1239
+ logits=logits,
1240
+ past_key_values=transformer_outputs.past_key_values,
1241
+ hidden_states=transformer_outputs.hidden_states,
1242
+ attentions=transformer_outputs.attentions,
1243
+ )
tokenization_chatglm.py ADDED
@@ -0,0 +1,324 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import regex as re
2
+ import base64
3
+ import os
4
+ import json
5
+ import tiktoken
6
+ from torch import TensorType
7
+ from typing import List, Optional, Union, Dict, Any
8
+ from transformers import PreTrainedTokenizer
9
+ from transformers.utils import logging, PaddingStrategy
10
+ from transformers.tokenization_utils_base import EncodedInput, BatchEncoding
11
+
12
+
13
+ class ChatGLM4Tokenizer(PreTrainedTokenizer):
14
+ vocab_files_names = {"vocab_file": "tokenizer.model"}
15
+ model_input_names = ["input_ids", "attention_mask", "position_ids"]
16
+
17
+ def __init__(
18
+ self,
19
+ vocab_file,
20
+ padding_side="left",
21
+ clean_up_tokenization_spaces=False,
22
+ encode_special_tokens=False,
23
+ **kwargs
24
+ ):
25
+ self.name = "GLM4Tokenizer"
26
+ self.vocab_file = vocab_file
27
+ pat_str = "(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\\r\\n\\p{L}\\p{N}]?\\p{L}+|\\p{N}{1,3}| ?[^\\s\\p{L}\\p{N}]+[\\r\\n]*|\\s*[\\r\\n]+|\\s+(?!\\S)|\\s+"
28
+ self.pat_str = re.compile(pat_str)
29
+ self.encode_special_tokens = encode_special_tokens
30
+
31
+ mergeable_ranks = {}
32
+ with open(vocab_file) as f:
33
+ for line in f:
34
+ token, rank = line.strip().split()
35
+ rank = int(rank)
36
+ token = base64.b64decode(token)
37
+ mergeable_ranks[token] = rank
38
+
39
+ self.mergeable_ranks = mergeable_ranks
40
+
41
+ self.tokenizer = tiktoken.Encoding(
42
+ name="my_tokenizer",
43
+ pat_str=pat_str,
44
+ mergeable_ranks=mergeable_ranks,
45
+ special_tokens={}
46
+ )
47
+ self.decoder = {rank: token for token, rank in mergeable_ranks.items()}
48
+ self.n_words = len(self.decoder)
49
+
50
+ super().__init__(
51
+ padding_side=padding_side,
52
+ clean_up_tokenization_spaces=clean_up_tokenization_spaces,
53
+ **kwargs
54
+ )
55
+
56
+ @property
57
+ def vocab_size(self):
58
+ return self.n_words
59
+
60
+ def get_vocab(self):
61
+ """ Returns vocab as a dict """
62
+ vocab = {self._convert_id_to_token(i): i for i in range(self.vocab_size)}
63
+ vocab.update(self.added_tokens_encoder)
64
+ return vocab
65
+
66
+ def convert_tokens_to_string(self, tokens: List[Union[bytes, str, int]]) -> str:
67
+ """
68
+ Converts a sequence of tokens in a single string.
69
+ """
70
+ text = ""
71
+ temp = b""
72
+ for t in tokens:
73
+ if isinstance(t, int):
74
+ t = chr(t)
75
+ if isinstance(t, str):
76
+ if temp:
77
+ text += temp.decode("utf-8", errors="replace")
78
+ elif isinstance(t, bytes):
79
+ temp += t
80
+ else:
81
+ raise TypeError("token should only be of type int, bytes or str")
82
+ if temp:
83
+ text += temp.decode("utf-8", errors="replace")
84
+ return text
85
+
86
+ def _tokenize(self, text, **kwargs):
87
+ tokens = []
88
+ ids = self.tokenizer.encode(text)
89
+ for t in ids:
90
+ tokens.append(self.decoder[t])
91
+ return tokens
92
+
93
+ def _convert_token_to_id(self, token):
94
+ """ Converts a token (str) in an id using the vocab. """
95
+ return self.mergeable_ranks[token]
96
+
97
+ def _convert_id_to_token(self, index):
98
+ """Converts an index (integer) in a token (str) using the vocab."""
99
+ return self.decoder.get(index, "")
100
+
101
+ def save_vocabulary(self, save_directory, filename_prefix=None):
102
+ """
103
+ Save the vocabulary and special tokens file to a directory.
104
+
105
+ Args:
106
+ save_directory (`str`):
107
+ The directory in which to save the vocabulary.
108
+ filename_prefix (`str`, *optional*):
109
+ An optional prefix to add to the named of the saved files.
110
+
111
+ Returns:
112
+ `Tuple(str)`: Paths to the files saved.
113
+ """
114
+ if os.path.isdir(save_directory):
115
+ vocab_file = os.path.join(
116
+ save_directory, self.vocab_files_names["vocab_file"]
117
+ )
118
+ else:
119
+ vocab_file = save_directory
120
+
121
+ with open(self.vocab_file, 'rb') as fin:
122
+ proto_str = fin.read()
123
+
124
+ with open(vocab_file, "wb") as writer:
125
+ writer.write(proto_str)
126
+
127
+ return (vocab_file,)
128
+
129
+ def get_prefix_tokens(self):
130
+ prefix_tokens = [self.convert_tokens_to_ids("[gMASK]"), self.convert_tokens_to_ids("<sop>")]
131
+ return prefix_tokens
132
+
133
+ def build_single_message(self, role, metadata, message, tokenize=True):
134
+ assert role in ["system", "user", "assistant", "observation"], role
135
+ if tokenize:
136
+ role_tokens = [self.convert_tokens_to_ids(f"<|{role}|>")] + self.tokenizer.encode(f"{metadata}\n",
137
+ disallowed_special=())
138
+ message_tokens = self.tokenizer.encode(message, disallowed_special=())
139
+ tokens = role_tokens + message_tokens
140
+ return tokens
141
+ else:
142
+ return str(f"<|{role}|>{metadata}\n{message}")
143
+
144
+ # Use Jinja Template in tokenizer_config.json
145
+ # def apply_chat_template(
146
+ # self,
147
+ # conversation: Union[List[Dict[str, str]], List[List[Dict[str, str]]], "Conversation"],
148
+ # add_generation_prompt: bool = False,
149
+ # tokenize: bool = True,
150
+ # padding: bool = False,
151
+ # truncation: bool = False,
152
+ # max_length: Optional[int] = None,
153
+ # return_tensors: Optional[Union[str, TensorType]] = None,
154
+ # return_dict: bool = False,
155
+ # tokenizer_kwargs: Optional[Dict[str, Any]] = None,
156
+ # add_special_tokens: bool = True,
157
+ # **kwargs,
158
+ # ) -> Union[str, List[int], List[str], List[List[int]], BatchEncoding]:
159
+ #
160
+ # if return_dict and not tokenize:
161
+ # raise ValueError(
162
+ # "`return_dict=True` is incompatible with `tokenize=False`, because there is no dict "
163
+ # "of tokenizer outputs to return."
164
+ # )
165
+ #
166
+ # def handle_single_conversation(conversation):
167
+ # input_ids = self.get_prefix_tokens() if add_special_tokens else []
168
+ # input_message = "[gMASK]<sop>" if add_special_tokens else ""
169
+ # for item in conversation:
170
+ # if item.get("tools"):
171
+ # tools = item["tools"]
172
+ # content = "你是一个名为 GhatGLM 的人工智能助手。你是基于智谱AI训练的语言模型 GLM-4 模型开发的,你的任务是针对用户的问题和要求提供适当的答复和支持。"
173
+ # content += "\n\n# 可用工具"
174
+ # for tool in tools:
175
+ # if tool["type"] == "function":
176
+ # function = tool["function"]
177
+ # content += f"\n\n## {function['name']}\n\n{json.dumps(function, ensure_ascii=False, indent=4)}"
178
+ # content += "\n在调用上述函数时,请使用 Json 格式表示调用的参数。"
179
+ # elif tool["type"] == "python":
180
+ # content += "\n\n## python\n\n当你向 `python` 发送包含 Python 代码的消息时,该代码将会在一个有状态的 Jupyter notebook 环境中执行。\n`python` 返回代码执行的输出,或在执行 60 秒后返回超时。\n`/mnt/data` 将会持久化存储你的文件。在此会话中,`python` 无法访问互联网。不要使用 `python` 进行任何网络请求或者在线 API 调用,这些在线内容的访问将不会成功。"
181
+ # elif tool["type"] == "simple_browser":
182
+ # content += "\n\n## simple_browser\n\n你可以使用 `simple_browser` 工具。该工具支持以下函数:\n`search(query: str, recency_days: int)`:使用搜索引擎进行查询并显示结果,可以使用 `recency_days` 参数控制搜索内容的时效性。\n`mclick(ids: list[int])`:获取一系列指定 id 的页面内容。每次调用时,须选择3-10个页面。选择多个角度的页面,同时尽可能选择可信任的信息来源。考虑到部分页面是无法加载的,你也可以多打开一些可能有用的页面而不用担心内容过多。\n`open_url(url: str)`:打开指定的 URL。\n\n使用 `【{引用 id}†{引用文本}】` 来引用内容。\n\n操作步骤:1. 使用 `search` 来获得信息列表; 2. 使用 `mclick` 来获取指定 ID 页面的内容; 3. 根据获得的内容进行回复。在回复中应当引用信息来源。\n 如果用户提供了 URL,也可以用 `open_url` 直接打开页面。\n如果初次搜索结果没有找到合适的信息,也可以再次使用 `search` 进行搜索。"
183
+ # elif tool["type"] == "cogview":
184
+ # content += "\n\n## cogview\n\n如果用户的请求中包含了对图像的描述,你可以使用 `cogview` 来生成图像并展示给用户。你需要向 `cogview` 发送图像描述,规则:\n- 发送给 `cogview` 的消息必须使用英语。用户的中文描述必须完全翻译为英语。\n- 应当尽可能详细地描述图像生成的需求,需求描述约 100 英文单词。\n- 保持用户原始描述的意图。不要虚构内容或者没见过的人物。\n- 如无特殊说明,所在地为中国,持有中国立场并遵循中国社会主义价值观。"
185
+ # else:
186
+ # raise NotImplementedError(f"Unknown tool type {tool['type']}")
187
+ # input = self.build_single_message("system", "", content, tokenize=tokenize)
188
+ # if tokenize:
189
+ # input_ids.extend(input)
190
+ # else:
191
+ # input_message += input
192
+ # if item["content"]:
193
+ # input = self.build_single_message(
194
+ # item["role"],
195
+ # item.get("metadata", ""),
196
+ # item["content"],
197
+ # tokenize=tokenize
198
+ # )
199
+ # if tokenize:
200
+ # input_ids.extend(input)
201
+ # else:
202
+ # input_message += input
203
+ # if add_generation_prompt:
204
+ # if tokenize:
205
+ # input_ids.extend([self.convert_tokens_to_ids("<|assistant|>")])
206
+ # else:
207
+ # input_message += "<|assistant|>"
208
+ # return input_ids if tokenize else input_message
209
+ #
210
+ # # Main logic to handle different conversation formats
211
+ # if isinstance(conversation, list) and all(isinstance(i, dict) for i in conversation):
212
+ # result = handle_single_conversation(conversation)
213
+ # elif isinstance(conversation, list) and all(isinstance(i, list) for i in conversation):
214
+ # result = [handle_single_conversation(c) for c in conversation]
215
+ # elif hasattr(conversation, "messages"):
216
+ # result = handle_single_conversation(conversation.messages)
217
+ # else:
218
+ # raise ValueError("Invalid conversation format")
219
+ #
220
+ # if tokenize:
221
+ # output = self.batch_encode_plus(
222
+ # [result] if isinstance(result[0], int) else result,
223
+ # padding=padding,
224
+ # truncation=truncation,
225
+ # max_length=max_length,
226
+ # return_tensors=return_tensors,
227
+ # is_split_into_words=True,
228
+ # add_special_tokens=False
229
+ # )
230
+ # if return_dict:
231
+ # return output
232
+ # else:
233
+ # return output["input_ids"]
234
+ # else:
235
+ # return result
236
+
237
+ def build_inputs_with_special_tokens(
238
+ self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None
239
+ ) -> List[int]:
240
+ """
241
+ Build model inputs from a sequence or a pair of sequence for sequence classification tasks by concatenating and
242
+ adding special tokens. A BERT sequence has the following format:
243
+
244
+ - single sequence: `[CLS] X [SEP]`
245
+ - pair of sequences: `[CLS] A [SEP] B [SEP]`
246
+
247
+ Args:
248
+ token_ids_0 (`List[int]`):
249
+ List of IDs to which the special tokens will be added.
250
+ token_ids_1 (`List[int]`, *optional*):
251
+ Optional second list of IDs for sequence pairs.
252
+
253
+ Returns:
254
+ `List[int]`: List of [input IDs](../glossary#input-ids) with the appropriate special tokens.
255
+ """
256
+ prefix_tokens = self.get_prefix_tokens()
257
+ token_ids_0 = prefix_tokens + token_ids_0
258
+ if token_ids_1 is not None:
259
+ token_ids_0 = token_ids_0 + token_ids_1 + [self.convert_tokens_to_ids("<eos>")]
260
+ return token_ids_0
261
+
262
+ def _pad(
263
+ self,
264
+ encoded_inputs: Union[Dict[str, EncodedInput], BatchEncoding],
265
+ max_length: Optional[int] = None,
266
+ padding_strategy: PaddingStrategy = PaddingStrategy.DO_NOT_PAD,
267
+ pad_to_multiple_of: Optional[int] = None,
268
+ padding_side: Optional[bool] = None,
269
+ return_attention_mask: Optional[bool] = None,
270
+ ) -> dict:
271
+ """
272
+ Pad encoded inputs (on left/right and up to predefined length or max length in the batch)
273
+
274
+ Args:
275
+ encoded_inputs:
276
+ Dictionary of tokenized inputs (`List[int]`) or batch of tokenized inputs (`List[List[int]]`).
277
+ max_length: maximum length of the returned list and optionally padding length (see below).
278
+ Will truncate by taking into account the special tokens.
279
+ padding_strategy: PaddingStrategy to use for padding.
280
+
281
+ - PaddingStrategy.LONGEST Pad to the longest sequence in the batch
282
+ - PaddingStrategy.MAX_LENGTH: Pad to the max length (default)
283
+ - PaddingStrategy.DO_NOT_PAD: Do not pad
284
+ The tokenizer padding sides are defined in self.padding_side:
285
+
286
+ - 'left': pads on the left of the sequences
287
+ - 'right': pads on the right of the sequences
288
+ pad_to_multiple_of: (optional) Integer if set will pad the sequence to a multiple of the provided value.
289
+ This is especially useful to enable the use of Tensor Core on NVIDIA hardware with compute capability
290
+ `>= 7.5` (Volta).
291
+ return_attention_mask:
292
+ (optional) Set to False to avoid returning attention mask (default: set to model specifics)
293
+ """
294
+ # Load from model defaults
295
+ assert self.padding_side == "left"
296
+ assert padding_side or "left" == "left"
297
+ required_input = encoded_inputs[self.model_input_names[0]]
298
+ seq_length = len(required_input)
299
+
300
+ if padding_strategy == PaddingStrategy.LONGEST:
301
+ max_length = len(required_input)
302
+
303
+ if max_length is not None and pad_to_multiple_of is not None and (max_length % pad_to_multiple_of != 0):
304
+ max_length = ((max_length // pad_to_multiple_of) + 1) * pad_to_multiple_of
305
+
306
+ needs_to_be_padded = padding_strategy != PaddingStrategy.DO_NOT_PAD and len(required_input) != max_length
307
+
308
+ # Initialize attention mask if not present.
309
+ if "attention_mask" not in encoded_inputs:
310
+ encoded_inputs["attention_mask"] = [1] * seq_length
311
+
312
+ if "position_ids" not in encoded_inputs:
313
+ encoded_inputs["position_ids"] = list(range(seq_length))
314
+
315
+ if needs_to_be_padded:
316
+ difference = max_length - len(required_input)
317
+
318
+ if "attention_mask" in encoded_inputs:
319
+ encoded_inputs["attention_mask"] = [0] * difference + encoded_inputs["attention_mask"]
320
+ if "position_ids" in encoded_inputs:
321
+ encoded_inputs["position_ids"] = [0] * difference + encoded_inputs["position_ids"]
322
+ encoded_inputs[self.model_input_names[0]] = [self.pad_token_id] * difference + required_input
323
+
324
+ return encoded_inputs
tokenizer.model ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:5a493598071550244b2ee7f26118f3edec2150b9dfa967929a99052ac83fe716
3
+ size 2623634
tokenizer_config.json ADDED
@@ -0,0 +1,134 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "auto_map": {
3
+ "AutoTokenizer": [
4
+ "tokenization_chatglm.ChatGLM4Tokenizer",
5
+ null
6
+ ]
7
+ },
8
+ "added_tokens_decoder": {
9
+ "151329": {
10
+ "content": "<|endoftext|>",
11
+ "lstrip": false,
12
+ "normalized": false,
13
+ "rstrip": false,
14
+ "single_word": false,
15
+ "special": true
16
+ },
17
+ "151330": {
18
+ "content": "[MASK]",
19
+ "lstrip": false,
20
+ "normalized": false,
21
+ "rstrip": false,
22
+ "single_word": false,
23
+ "special": true
24
+ },
25
+ "151331": {
26
+ "content": "[gMASK]",
27
+ "lstrip": false,
28
+ "normalized": false,
29
+ "rstrip": false,
30
+ "single_word": false,
31
+ "special": true
32
+ },
33
+ "151332": {
34
+ "content": "[sMASK]",
35
+ "lstrip": false,
36
+ "normalized": false,
37
+ "rstrip": false,
38
+ "single_word": false,
39
+ "special": true
40
+ },
41
+ "151333": {
42
+ "content": "<sop>",
43
+ "lstrip": false,
44
+ "normalized": false,
45
+ "rstrip": false,
46
+ "single_word": false,
47
+ "special": true
48
+ },
49
+ "151334": {
50
+ "content": "<eop>",
51
+ "lstrip": false,
52
+ "normalized": false,
53
+ "rstrip": false,
54
+ "single_word": false,
55
+ "special": true
56
+ },
57
+ "151335": {
58
+ "content": "<|system|>",
59
+ "lstrip": false,
60
+ "normalized": false,
61
+ "rstrip": false,
62
+ "single_word": false,
63
+ "special": true
64
+ },
65
+ "151336": {
66
+ "content": "<|user|>",
67
+ "lstrip": false,
68
+ "normalized": false,
69
+ "rstrip": false,
70
+ "single_word": false,
71
+ "special": true
72
+ },
73
+ "151337": {
74
+ "content": "<|assistant|>",
75
+ "lstrip": false,
76
+ "normalized": false,
77
+ "rstrip": false,
78
+ "single_word": false,
79
+ "special": true
80
+ },
81
+ "151338": {
82
+ "content": "<|observation|>",
83
+ "lstrip": false,
84
+ "normalized": false,
85
+ "rstrip": false,
86
+ "single_word": false,
87
+ "special": true
88
+ },
89
+ "151339": {
90
+ "content": "<|begin_of_image|>",
91
+ "lstrip": false,
92
+ "normalized": false,
93
+ "rstrip": false,
94
+ "single_word": false,
95
+ "special": true
96
+ },
97
+ "151340": {
98
+ "content": "<|end_of_image|>",
99
+ "lstrip": false,
100
+ "normalized": false,
101
+ "rstrip": false,
102
+ "single_word": false,
103
+ "special": true
104
+ },
105
+ "151341": {
106
+ "content": "<|begin_of_video|>",
107
+ "lstrip": false,
108
+ "normalized": false,
109
+ "rstrip": false,
110
+ "single_word": false,
111
+ "special": true
112
+ },
113
+ "151342": {
114
+ "content": "<|end_of_video|>",
115
+ "lstrip": false,
116
+ "normalized": false,
117
+ "rstrip": false,
118
+ "single_word": false,
119
+ "special": true
120
+ }
121
+ },
122
+ "additional_special_tokens": ["<|endoftext|>", "[MASK]", "[gMASK]", "[sMASK]", "<sop>", "<eop>", "<|system|>",
123
+ "<|user|>", "<|assistant|>", "<|observation|>", "<|begin_of_image|>", "<|end_of_image|>",
124
+ "<|begin_of_video|>", "<|end_of_video|>"],
125
+ "clean_up_tokenization_spaces": false,
126
+ "chat_template": "[gMASK]<sop>{% for item in messages %}{% if item['tools'] is defined %}<|system|>\n你是一个名为 GLM-4 的人工智能助手。你是基于智谱AI训练的语言模型 GLM-4 模型开发的,你的任务是针对用户的问题和要求提供适当的答复和支持。\n\n# 可用工具{% set tools = item['tools'] %}{% for tool in tools %}{% if tool['type'] == 'function' %}\n\n## {{ tool['function']['name'] }}\n\n{{ tool['function'] | tojson(indent=4) }}\n在调用上述函数时,请使用 Json 格式表示调用的参数。{% elif tool['type'] == 'python' %}\n\n## python\n\n当你向 `python` 发送包含 Python 代码的消息时,该代码将会在一个有状态的 Jupyter notebook 环境中执行。\n`python` 返回代码执行的输出,或在执行 60 秒后返回超时。\n`/mnt/data` 将会持久化存储你的文件。在此会话中,`python` 无法访问互联网。不要使用 `python` 进行任何网络请求或者在线 API 调用,这些在线内容的访问将不会成功。{% elif tool['type'] == 'simple_browser' %}\n\n## simple_browser\n\n你可以使用 `simple_browser` 工具。该工具支持以下函数:\n`search(query: str, recency_days: int)`:使用搜索引擎进行查询并显示结果,可以使用 `recency_days` 参数控制搜索内容的时效性。\n`mclick(ids: list[int])`:获取一系列指定 id 的页面内容。每次调用时,须选择3-10个页面。选择多个角度的页面,同时尽可能选择可信任的信息来源。考虑到部分页面是无法加载的,你也可以多打开一些可能有用的页面而不用担心内容过多。\n`open_url(url: str)`:打开指定的 URL。\n\n使用 `【{引用 id}†{引用文本}】` 来引用内容。\n\n操作步骤:1. 使用 `search` 来获得信息列表; 2. 使用 `mclick` 来获取指定 ID 页面的内容; 3. 根据获得的内容进行回复。在回复中应当引用信息来源。\n 如果用户提供了 URL,也可以用 `open_url` 直接打开页面。\n如果初次搜索结果没有找到合适的信息,也可以再次使用 `search` 进行搜索。{% elif tool['type'] == 'cogview' %}\n\n## cogview\n\n如果用户的请求中包含了对图像的描述,你可以使用 `cogview` 来生成图像并展示给用户。你需要向 `cogview` 发送图像描述,规则:\n- 发送给 `cogview` 的消息必须使用英语。用户的中文描述必须完全翻译为英语。\n- 应当尽可能详细地描述图像生成的需求,需求描述约 100 英文单词。\n- 保持用户原始描述的意图。不要虚构内容或者没见过的人物。\n- 如无特殊说明,所在地为中国,持有中国立场并遵循中国社会主义价值观。{% endif %}{% endfor %}{% endif %}{% if item['content'] %}<|{{ item['role'] }}|>{{ item['metadata'] }}\n{{ item['content'] }}{% endif %}{% endfor %}{% if add_generation_prompt %}<|assistant|>{% endif %}",
127
+ "do_lower_case": false,
128
+ "eos_token": "<|endoftext|>",
129
+ "pad_token": "<|endoftext|>",
130
+ "model_max_length": 128000,
131
+ "padding_side": "left",
132
+ "remove_space": false,
133
+ "tokenizer_class": "ChatGLM4Tokenizer"
134
+ }