chujiezheng wandz commited on
Commit
da34456
1 Parent(s): 366733e

add model (#1)

Browse files

- add model (9ccb83034fdcdedde374be8482fbf0af8ed0fb1d)


Co-authored-by: one dozon <[email protected]>

.gitattributes CHANGED
@@ -33,3 +33,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
33
  *.zip filter=lfs diff=lfs merge=lfs -text
34
  *.zst filter=lfs diff=lfs merge=lfs -text
35
  *tfevents* filter=lfs diff=lfs merge=lfs -text
 
 
 
33
  *.zip filter=lfs diff=lfs merge=lfs -text
34
  *.zst filter=lfs diff=lfs merge=lfs -text
35
  *tfevents* filter=lfs diff=lfs merge=lfs -text
36
+ pytorch_model-00001-of-00002.bin filter=lfs diff=lfs merge=lfs -text
37
+ pytorch_model-00002-of-00002.bin filter=lfs diff=lfs merge=lfs -text
README.md ADDED
@@ -0,0 +1,11 @@
 
 
 
 
 
 
 
 
 
 
 
 
1
+
2
+ 代码示例
3
+
4
+ ```python
5
+ >>> from transformers import AutoTokenizer, AutoModel
6
+ >>> tokenizer = AutoTokenizer.from_pretrained("LingxinAI/CharacterGLM-6b", trust_remote_code=True)
7
+ >>> model = AutoModel.from_pretrained("LingxinAI/CharacterGLM-6b", trust_remote_code=True, device='cuda')
8
+ >>> model = model.eval()
9
+ >>> session_meta = {'user_info': '我是陆星辰,是一个男性,是一位知名导演,也是苏梦远的合作导演。我擅长拍摄音乐题材的电影。苏梦远对我的态度是尊敬的,并视我为良师益友。', 'bot_info': '苏梦远,本名苏远心,是一位当红的国内女歌手及演员。在参加选秀节目后,凭借独特的嗓音及出众的舞台魅力迅速成名,进入娱乐圈。她外表美丽动人,但真正的魅力在于她的才华和勤奋。苏梦远是音乐学院毕业的优秀生,善于创作,拥有多首热门原创歌曲。除了音乐方面的成就,她还热衷于慈善事业,积极参加公益活动,用实际行动传递正能量。在工作中,她对待工作非常敬业,拍戏时总是全身心投入角色,赢得了业内人士的赞誉和粉丝的喜爱。虽然在娱乐圈,但她始终保持低调、谦逊的态度,深得同行尊重。在表达时,苏梦远喜欢使用“我们”和“一起”,强调团队精神。', 'bot_name': '苏梦远', 'user_name': '陆星辰'}
10
+ >>> response, history = model.chat(tokenizer, session_meta, "你好", history=[])
11
+ ```
characterglm_generation_utils.py ADDED
@@ -0,0 +1,121 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ from typing import TypedDict, Literal, List, Optional, Tuple, Iterator
3
+
4
+
5
+ #### data types #########
6
+ # 下面的数据类型定义与CharacterGLM API一致,但与modeling_chatglm.py的chat方法不一致
7
+ # 参考 https://open.bigmodel.cn/dev/api#characterglm
8
+ RoleType = Literal["user", "assistant"]
9
+
10
+ class Msg(TypedDict):
11
+ role: RoleType
12
+ content: str
13
+
14
+
15
+ class SessionMeta(TypedDict):
16
+ user_name: str
17
+ bot_name: str
18
+ bot_info: str
19
+ user_info: Optional[str]
20
+
21
+
22
+ HistoryType = List[Msg]
23
+
24
+
25
+ class CharacterGLMGenerationUtils:
26
+ @staticmethod
27
+ def convert_chatglm_history_to_characterglm_history(user_query: str, history: List[Tuple[str, str]]) -> HistoryType:
28
+ characterglm_history: HistoryType = []
29
+ for i, (query, response) in enumerate(history):
30
+ if i == 0 and query == '':
31
+ # first empty query is an placeholder
32
+ pass
33
+ else:
34
+ characterglm_history.append({
35
+ "role": "user",
36
+ "content": query
37
+ })
38
+ characterglm_history.append({
39
+ "role": "assistant",
40
+ "content": response
41
+ })
42
+
43
+ characterglm_history.append({
44
+ "role": "user",
45
+ "content": user_query
46
+ })
47
+ return characterglm_history
48
+
49
+ @staticmethod
50
+ def build_inputs(session_meta: SessionMeta, history: HistoryType) -> str:
51
+ """
52
+ 注意:这里假设history最后一条消息是用户query
53
+ """
54
+ texts = []
55
+ texts.append(
56
+ f"以下是一段{session_meta['bot_name']}和{session_meta['user_name']}之间的对话。")
57
+ if session_meta.get("bot_info"):
58
+ texts.append(f"关于{session_meta['bot_name']}的信息:{session_meta['bot_info']}")
59
+ if session_meta.get("user_info"):
60
+ texts.append(
61
+ f"关于{session_meta['user_name']}的信息:{session_meta['user_info']}")
62
+
63
+ assert history and history[-1]['role'] == 'user'
64
+ for msg in history:
65
+ name = session_meta['user_name'] if msg['role'] == 'user' else session_meta['bot_name']
66
+ texts.append(f"[{name}]" + msg['content'].strip())
67
+
68
+ texts = [text.replace('\n', ' ') for text in texts]
69
+ texts.append(f"[{session_meta['bot_name']}]")
70
+ return '\n'.join(texts)
71
+
72
+
73
+ class CharacterGLMAPI:
74
+ @staticmethod
75
+ def build_api_arguments(session_meta: SessionMeta, history: HistoryType) -> dict:
76
+ return {
77
+ "model": "characterglm",
78
+ "meta": session_meta,
79
+ "prompt": history
80
+ }
81
+
82
+ @classmethod
83
+ def async_invoke(cls, session_meta: SessionMeta, history: HistoryType):
84
+ """
85
+ 注意:
86
+ 1. 先设置zhipuai.api_key
87
+ 2. 建议传入`return_type='text'`,否则返回结果是json字符串
88
+
89
+ 参考:
90
+ https://open.bigmodel.cn/dev/api#characterglm
91
+ """
92
+ import zhipuai
93
+ kwargs = cls.build_api_arguments(session_meta, history)
94
+ return zhipuai.model_api.async_invoke(**kwargs, return_type='text')
95
+
96
+ @classmethod
97
+ def invoke(cls, session_meta: SessionMeta, history: HistoryType):
98
+ """
99
+ 注意:
100
+ 1. 先设置zhipuai.api_key
101
+ 2. 建议传入`return_type='text'`,否则返回结果是json字符串
102
+ 3. 需要再次调用`zhipuai.model_api.query_async_invoke_result`才能获取生成结果
103
+
104
+ 参考:
105
+ https://open.bigmodel.cn/dev/api#characterglm
106
+ """
107
+ import zhipuai
108
+ kwargs = cls.build_api_arguments(session_meta, history)
109
+ return zhipuai.model_api.invoke(**kwargs, return_type='text')
110
+
111
+ @classmethod
112
+ def generate(cls, session_meta: SessionMeta, history: HistoryType) -> str:
113
+ result = cls.invoke(session_meta, history)
114
+ if not result['success']:
115
+ raise RuntimeError(result)
116
+ return result['data']['choices'][0]['content']
117
+
118
+ @classmethod
119
+ def stream_generate(cls, session_meta: SessionMeta, history: HistoryType) -> Iterator[str]:
120
+ # 伪流式生成
121
+ return iter(cls.generate(session_meta, history))
config.json ADDED
@@ -0,0 +1,42 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "_name_or_path": "thu-coai/CharacterGLM-6B",
3
+ "model_type": "chatglm",
4
+ "architectures": [
5
+ "ChatGLMModel"
6
+ ],
7
+ "auto_map": {
8
+ "AutoConfig": "configuration_chatglm.ChatGLMConfig",
9
+ "AutoModel": "modeling_characterglm.CharacterGLMForConditionalGeneration",
10
+ "AutoModelForCausalLM": "modeling_characterglm.CharacterGLMForConditionalGeneration",
11
+ "AutoModelForSeq2SeqLM": "modeling_characterglm.CharacterGLMForConditionalGeneration",
12
+ "AutoModelForSequenceClassification": "modeling_chatglm.ChatGLMForSequenceClassification"
13
+ },
14
+ "add_bias_linear": false,
15
+ "add_qkv_bias": true,
16
+ "apply_query_key_layer_scaling": true,
17
+ "apply_residual_connection_post_layernorm": false,
18
+ "attention_dropout": 0.0,
19
+ "attention_softmax_in_fp32": true,
20
+ "bias_dropout_fusion": true,
21
+ "ffn_hidden_size": 13696,
22
+ "fp32_residual_connection": false,
23
+ "hidden_dropout": 0.0,
24
+ "hidden_size": 4096,
25
+ "kv_channels": 128,
26
+ "layernorm_epsilon": 1e-05,
27
+ "multi_query_attention": true,
28
+ "multi_query_group_num": 2,
29
+ "num_attention_heads": 32,
30
+ "num_layers": 28,
31
+ "original_rope": true,
32
+ "padded_vocab_size": 65024,
33
+ "post_layer_norm": true,
34
+ "rmsnorm": true,
35
+ "seq_length": 32768,
36
+ "use_cache": true,
37
+ "torch_dtype": "float16",
38
+ "transformers_version": "4.27.1",
39
+ "tie_word_embeddings": false,
40
+ "eos_token_id": 2,
41
+ "pad_token_id": 0
42
+ }
configuration_chatglm.py ADDED
@@ -0,0 +1,61 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from transformers import PretrainedConfig
2
+
3
+
4
+ class ChatGLMConfig(PretrainedConfig):
5
+ model_type = "chatglm"
6
+ def __init__(
7
+ self,
8
+ num_layers=28,
9
+ padded_vocab_size=65024,
10
+ hidden_size=4096,
11
+ ffn_hidden_size=13696,
12
+ kv_channels=128,
13
+ num_attention_heads=32,
14
+ seq_length=2048,
15
+ hidden_dropout=0.0,
16
+ classifier_dropout=None,
17
+ attention_dropout=0.0,
18
+ layernorm_epsilon=1e-5,
19
+ rmsnorm=True,
20
+ apply_residual_connection_post_layernorm=False,
21
+ post_layer_norm=True,
22
+ add_bias_linear=False,
23
+ add_qkv_bias=False,
24
+ bias_dropout_fusion=True,
25
+ multi_query_attention=False,
26
+ multi_query_group_num=1,
27
+ apply_query_key_layer_scaling=True,
28
+ attention_softmax_in_fp32=True,
29
+ fp32_residual_connection=False,
30
+ quantization_bit=0,
31
+ pre_seq_len=None,
32
+ prefix_projection=False,
33
+ **kwargs
34
+ ):
35
+ self.num_layers = num_layers
36
+ self.vocab_size = padded_vocab_size
37
+ self.padded_vocab_size = padded_vocab_size
38
+ self.hidden_size = hidden_size
39
+ self.ffn_hidden_size = ffn_hidden_size
40
+ self.kv_channels = kv_channels
41
+ self.num_attention_heads = num_attention_heads
42
+ self.seq_length = seq_length
43
+ self.hidden_dropout = hidden_dropout
44
+ self.classifier_dropout = classifier_dropout
45
+ self.attention_dropout = attention_dropout
46
+ self.layernorm_epsilon = layernorm_epsilon
47
+ self.rmsnorm = rmsnorm
48
+ self.apply_residual_connection_post_layernorm = apply_residual_connection_post_layernorm
49
+ self.post_layer_norm = post_layer_norm
50
+ self.add_bias_linear = add_bias_linear
51
+ self.add_qkv_bias = add_qkv_bias
52
+ self.bias_dropout_fusion = bias_dropout_fusion
53
+ self.multi_query_attention = multi_query_attention
54
+ self.multi_query_group_num = multi_query_group_num
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
+ self.quantization_bit = quantization_bit
59
+ self.pre_seq_len = pre_seq_len
60
+ self.prefix_projection = prefix_projection
61
+ super().__init__(**kwargs)
generation_config.json ADDED
@@ -0,0 +1,6 @@
 
 
 
 
 
 
 
1
+ {
2
+ "_from_model_config": true,
3
+ "eos_token_id": 2,
4
+ "pad_token_id": 0,
5
+ "transformers_version": "4.31.0"
6
+ }
modeling_characterglm.py ADDED
@@ -0,0 +1,218 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import copy
2
+ import warnings
3
+ import logging
4
+ from typing import List, Tuple, Optional, Callable
5
+
6
+ import torch
7
+ from torch import nn
8
+ from transformers.utils import logging
9
+ from transformers.generation.utils import LogitsProcessorList, StoppingCriteriaList, GenerationConfig
10
+
11
+ from .modeling_chatglm import ChatGLMForConditionalGeneration, InvalidScoreLogitsProcessor
12
+ from .characterglm_generation_utils import CharacterGLMGenerationUtils, SessionMeta
13
+
14
+
15
+ logger = logging.get_logger(__name__)
16
+ default_generation_config = {
17
+ "do_sample": True,
18
+ "top_k": 100,
19
+ "top_p": 0.9,
20
+ "no_repeat_ngram_size": 0,
21
+ "temperature": 0.9,
22
+ "num_beams": 1,
23
+ "length_penalty": 1.6,
24
+ "repetition_penalty": 1.3,
25
+ "eos_token_id": 13
26
+ }
27
+
28
+
29
+ class CharacterGLMForConditionalGeneration(ChatGLMForConditionalGeneration):
30
+ """
31
+ CharacterGLM的prompt格式与chatglm有差异。
32
+ CharacterGLMForConditionalGeneration复用了ChatGLMForConditionalGeneration的forward方法,
33
+ 重新实现了`build_inputs`和`build_stream_inputs`,
34
+ 调整了`chat`和`stream_chat`方法的函数签名,增加session_meta参数,并修改解码参数的默认值。
35
+ """
36
+
37
+ def build_inputs(self, tokenizer, session_meta: SessionMeta, query: str, history: Optional[List[Tuple[str, str]]] = None):
38
+ character_glm_history = CharacterGLMGenerationUtils.convert_chatglm_history_to_characterglm_history(query, history or [])
39
+ prompt = CharacterGLMGenerationUtils.build_inputs(session_meta, character_glm_history)
40
+ inputs = tokenizer([prompt], return_tensors="pt")
41
+ inputs = inputs.to(self.device)
42
+ return inputs
43
+
44
+ def build_stream_inputs(self, tokenizer, session_meta: SessionMeta, query: str, history: Optional[List[Tuple[str, str]]] = None):
45
+ prompt = "\n[{}]{}\n[{}]".format(
46
+ session_meta['user_name'],
47
+ query.replace('\n', ' '),
48
+ session_meta['bot_name']
49
+ )
50
+ input_ids = tokenizer.encode(prompt, add_special_tokens=False)
51
+ input_ids = input_ids[1:]
52
+ inputs = tokenizer.batch_encode_plus([(input_ids, None)], return_tensors="pt", add_special_tokens=False)
53
+ inputs = inputs.to(self.device)
54
+ return inputs
55
+
56
+ @torch.inference_mode()
57
+ def chat(self, tokenizer, session_meta: SessionMeta, query: str, history: List[Tuple[str, str]] = None, max_length: int = 8192, num_beams=1,
58
+ do_sample=True, top_p=0.9, temperature=0.9, repetition_penalty=1.6, logits_processor=None, **kwargs):
59
+ if history is None:
60
+ history = []
61
+ if logits_processor is None:
62
+ logits_processor = LogitsProcessorList()
63
+ logits_processor.append(InvalidScoreLogitsProcessor())
64
+ gen_kwargs = {"max_length": max_length, "num_beams": num_beams, "do_sample": do_sample, "top_p": top_p,
65
+ "temperature": temperature, "logits_processor": logits_processor, "repetition_penalty": repetition_penalty, **kwargs}
66
+ gen_kwargs.update({k: v for k, v in default_generation_config.items() if k not in gen_kwargs})
67
+ inputs = self.build_inputs(tokenizer, session_meta, query, history=history)
68
+ outputs = self.generate(**inputs, **gen_kwargs)
69
+ outputs = outputs.tolist()[0][len(inputs["input_ids"][0]):]
70
+ response = tokenizer.decode(outputs)
71
+ response = self.process_response(response)
72
+ history = history + [(query, response)]
73
+ return response, history
74
+
75
+ @torch.inference_mode()
76
+ def stream_chat(self, tokenizer, session_meta: SessionMeta, query: str, history: List[Tuple[str, str]] = None, past_key_values=None,
77
+ max_length: int = 8192, do_sample=True, top_p=0.9, temperature=0.9, repetition_penalty=1.0, logits_processor=None,
78
+ return_past_key_values=False, **kwargs):
79
+ if history is None:
80
+ history = []
81
+ if logits_processor is None:
82
+ logits_processor = LogitsProcessorList()
83
+ logits_processor.append(InvalidScoreLogitsProcessor())
84
+ gen_kwargs = {"max_length": max_length, "do_sample": do_sample, "top_p": top_p,
85
+ "temperature": temperature, "logits_processor": logits_processor, "repetition_penalty": repetition_penalty, **kwargs}
86
+ gen_kwargs.update({k: v for k, v in default_generation_config.items() if k not in gen_kwargs})
87
+ gen_kwargs.pop('repetition_penalty', None)
88
+ if past_key_values is None:
89
+ inputs = self.build_inputs(tokenizer, session_meta, query, history=history)
90
+ else:
91
+ inputs = self.build_stream_inputs(tokenizer, session_meta, query, history=history)
92
+ if past_key_values is not None:
93
+ past_length = past_key_values[0][0].shape[0]
94
+ if self.transformer.pre_seq_len is not None:
95
+ past_length -= self.transformer.pre_seq_len
96
+ inputs.position_ids += past_length
97
+ attention_mask = inputs.attention_mask
98
+ attention_mask = torch.cat((attention_mask.new_ones(1, past_length), attention_mask), dim=1)
99
+ inputs['attention_mask'] = attention_mask
100
+ for outputs in self.stream_generate(**inputs, past_key_values=past_key_values,
101
+ return_past_key_values=return_past_key_values, **gen_kwargs):
102
+ if return_past_key_values:
103
+ outputs, past_key_values = outputs
104
+ outputs = outputs.tolist()[0][len(inputs["input_ids"][0]):]
105
+ response = tokenizer.decode(outputs)
106
+ if response and response[-1] != "�":
107
+ response = self.process_response(response)
108
+ new_history = history + [(query, response)]
109
+ if return_past_key_values:
110
+ yield response, new_history, past_key_values
111
+ else:
112
+ yield response, new_history
113
+
114
+ @torch.inference_mode()
115
+ def stream_generate(
116
+ self,
117
+ input_ids,
118
+ generation_config: Optional[GenerationConfig] = None,
119
+ logits_processor: Optional[LogitsProcessorList] = None,
120
+ stopping_criteria: Optional[StoppingCriteriaList] = None,
121
+ prefix_allowed_tokens_fn: Optional[Callable[[int, torch.Tensor], List[int]]] = None,
122
+ return_past_key_values=False,
123
+ **kwargs,
124
+ ):
125
+ batch_size, input_ids_seq_length = input_ids.shape[0], input_ids.shape[-1]
126
+
127
+ if generation_config is None:
128
+ generation_config = self.generation_config
129
+ generation_config = copy.deepcopy(generation_config)
130
+ model_kwargs = generation_config.update(**kwargs)
131
+ model_kwargs["use_cache"] = generation_config.use_cache
132
+ bos_token_id, eos_token_id = generation_config.bos_token_id, generation_config.eos_token_id
133
+
134
+ if isinstance(eos_token_id, int):
135
+ eos_token_id = [eos_token_id]
136
+
137
+ has_default_max_length = kwargs.get("max_length") is None and generation_config.max_length is not None
138
+ if has_default_max_length and generation_config.max_new_tokens is None:
139
+ warnings.warn(
140
+ f"Using `max_length`'s default ({generation_config.max_length}) to control the generation length. "
141
+ "This behaviour is deprecated and will be removed from the config in v5 of Transformers -- we"
142
+ " recommend using `max_new_tokens` to control the maximum length of the generation.",
143
+ UserWarning,
144
+ )
145
+ elif generation_config.max_new_tokens is not None:
146
+ generation_config.max_length = generation_config.max_new_tokens + input_ids_seq_length
147
+ if not has_default_max_length:
148
+ logger.warn(
149
+ f"Both `max_new_tokens` (={generation_config.max_new_tokens}) and `max_length`(="
150
+ f"{generation_config.max_length}) seem to have been set. `max_new_tokens` will take precedence. "
151
+ "Please refer to the documentation for more information. "
152
+ "(https://huggingface.co/docs/transformers/main/en/main_classes/text_generation)",
153
+ UserWarning,
154
+ )
155
+
156
+ if input_ids_seq_length >= generation_config.max_length:
157
+ input_ids_string = "decoder_input_ids" if self.config.is_encoder_decoder else "input_ids"
158
+ logger.warning(
159
+ f"Input length of {input_ids_string} is {input_ids_seq_length}, but `max_length` is set to"
160
+ f" {generation_config.max_length}. This can lead to unexpected behavior. You should consider"
161
+ " increasing `max_new_tokens`."
162
+ )
163
+
164
+ # 2. Set generation parameters if not already defined
165
+ logits_processor = logits_processor if logits_processor is not None else LogitsProcessorList()
166
+ stopping_criteria = stopping_criteria if stopping_criteria is not None else StoppingCriteriaList()
167
+
168
+ logits_processor = self._get_logits_processor(
169
+ generation_config=generation_config,
170
+ input_ids_seq_length=input_ids_seq_length,
171
+ encoder_input_ids=input_ids,
172
+ prefix_allowed_tokens_fn=prefix_allowed_tokens_fn,
173
+ logits_processor=logits_processor,
174
+ )
175
+
176
+ stopping_criteria = self._get_stopping_criteria(
177
+ generation_config=generation_config, stopping_criteria=stopping_criteria
178
+ )
179
+ logits_warper = self._get_logits_warper(generation_config)
180
+
181
+ unfinished_sequences = input_ids.new(input_ids.shape[0]).fill_(1)
182
+ scores = None
183
+ while True:
184
+ model_inputs = self.prepare_inputs_for_generation(input_ids, **model_kwargs)
185
+ # forward pass to get next token
186
+ outputs = self(
187
+ **model_inputs,
188
+ return_dict=True,
189
+ output_attentions=False,
190
+ output_hidden_states=False,
191
+ )
192
+
193
+ next_token_logits = outputs.logits[:, -1, :]
194
+
195
+ # pre-process distribution
196
+ next_token_scores = logits_processor(input_ids, next_token_logits)
197
+ next_token_scores = logits_warper(input_ids, next_token_scores)
198
+
199
+ # sample
200
+ probs = nn.functional.softmax(next_token_scores, dim=-1)
201
+ if generation_config.do_sample:
202
+ next_tokens = torch.multinomial(probs, num_samples=1).squeeze(1)
203
+ else:
204
+ next_tokens = torch.argmax(probs, dim=-1)
205
+
206
+ # update generated ids, model inputs, and length for next step
207
+ input_ids = torch.cat([input_ids, next_tokens[:, None]], dim=-1)
208
+ model_kwargs = self._update_model_kwargs_for_generation(
209
+ outputs, model_kwargs, is_encoder_decoder=self.config.is_encoder_decoder
210
+ )
211
+ unfinished_sequences = unfinished_sequences.mul((sum(next_tokens != i for i in eos_token_id)).long())
212
+ if return_past_key_values:
213
+ yield input_ids, outputs.past_key_values
214
+ else:
215
+ yield input_ids
216
+ # stop when each sentence is finished, or if we exceed the maximum length
217
+ if unfinished_sequences.max() == 0 or stopping_criteria(input_ids, scores):
218
+ break
modeling_chatglm.py ADDED
@@ -0,0 +1,1285 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """ PyTorch ChatGLM model. """
2
+
3
+ import math
4
+ import copy
5
+ import warnings
6
+ import re
7
+ import sys
8
+
9
+ import torch
10
+ import torch.utils.checkpoint
11
+ import torch.nn.functional as F
12
+ from torch import nn
13
+ from torch.nn import CrossEntropyLoss, LayerNorm
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
+
18
+ from transformers.modeling_outputs import (
19
+ BaseModelOutputWithPast,
20
+ CausalLMOutputWithPast,
21
+ SequenceClassifierOutputWithPast,
22
+ )
23
+ from transformers.modeling_utils import PreTrainedModel
24
+ from transformers.utils import logging
25
+ from transformers.generation.logits_process import LogitsProcessor
26
+ from transformers.generation.utils import LogitsProcessorList, StoppingCriteriaList, GenerationConfig, ModelOutput
27
+
28
+ from .configuration_chatglm import ChatGLMConfig
29
+
30
+ # flags required to enable jit fusion kernels
31
+
32
+ if 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
+
38
+ logger = logging.get_logger(__name__)
39
+
40
+ _CHECKPOINT_FOR_DOC = "THUDM/ChatGLM2-6B"
41
+ _CONFIG_FOR_DOC = "ChatGLM6BConfig"
42
+
43
+ CHATGLM_6B_PRETRAINED_MODEL_ARCHIVE_LIST = [
44
+ "THUDM/chatglm2-6b",
45
+ # See all ChatGLM models at https://huggingface.co/models?filter=chatglm
46
+ ]
47
+
48
+
49
+ def default_init(cls, *args, **kwargs):
50
+ return cls(*args, **kwargs)
51
+
52
+
53
+ class 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
+
61
+ class 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
+
93
+ def 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
+
121
+ class 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=dtype, device=device) / n_elem))
140
+
141
+ # Create position indexes `[0, 1, ..., seq_len - 1]`
142
+ seq_idx = torch.arange(seq_len, dtype=dtype, 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
161
+ def 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
+
181
+ class 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
+
195
+ class 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
+
313
+ class 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
+
452
+ def _config_to_kwargs(args):
453
+ common_kwargs = {
454
+ "dtype": args.torch_dtype,
455
+ }
456
+ return common_kwargs
457
+
458
+
459
+ class 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
+
505
+ class 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
+
579
+ class 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
+ )
639
+ else:
640
+ layer_ret = layer(
641
+ hidden_states,
642
+ attention_mask,
643
+ rotary_pos_emb,
644
+ kv_cache=kv_caches[index],
645
+ use_cache=use_cache
646
+ )
647
+ hidden_states, kv_cache = layer_ret
648
+ if use_cache:
649
+ presents = presents + (kv_cache,)
650
+
651
+ if output_hidden_states:
652
+ all_hidden_states = all_hidden_states + (hidden_states,)
653
+
654
+ # Final layer norm.
655
+ if self.post_layer_norm:
656
+ hidden_states = self.final_layernorm(hidden_states)
657
+
658
+ return hidden_states, presents, all_hidden_states, all_self_attentions
659
+
660
+
661
+ class ChatGLMPreTrainedModel(PreTrainedModel):
662
+ """
663
+ An abstract class to handle weights initialization and
664
+ a simple interface for downloading and loading pretrained models.
665
+ """
666
+
667
+ is_parallelizable = False
668
+ supports_gradient_checkpointing = True
669
+ config_class = ChatGLMConfig
670
+ base_model_prefix = "transformer"
671
+ _no_split_modules = ["GLMBlock"]
672
+
673
+ def _init_weights(self, module: nn.Module):
674
+ """Initialize the weights."""
675
+ return
676
+
677
+ def get_masks(self, input_ids, past_key_values, padding_mask=None):
678
+ batch_size, seq_length = input_ids.shape
679
+ full_attention_mask = torch.ones(batch_size, seq_length, seq_length, device=input_ids.device)
680
+ full_attention_mask.tril_()
681
+ past_length = 0
682
+ if past_key_values:
683
+ past_length = past_key_values[0][0].shape[0]
684
+ if past_length:
685
+ full_attention_mask = torch.cat((torch.ones(batch_size, seq_length, past_length,
686
+ device=input_ids.device), full_attention_mask), dim=-1)
687
+ if padding_mask is not None:
688
+ full_attention_mask = full_attention_mask * padding_mask.unsqueeze(1)
689
+ if not past_length and padding_mask is not None:
690
+ full_attention_mask -= padding_mask.unsqueeze(-1) - 1
691
+ full_attention_mask = (full_attention_mask < 0.5).bool()
692
+ full_attention_mask.unsqueeze_(1)
693
+ return full_attention_mask
694
+
695
+ def get_position_ids(self, input_ids, device):
696
+ batch_size, seq_length = input_ids.shape
697
+ position_ids = torch.arange(seq_length, dtype=torch.long, device=device).unsqueeze(0).repeat(batch_size, 1)
698
+ return position_ids
699
+
700
+ def _set_gradient_checkpointing(self, module, value=False):
701
+ if isinstance(module, GLMTransformer):
702
+ module.gradient_checkpointing = value
703
+
704
+
705
+ class Embedding(torch.nn.Module):
706
+ """Language model embeddings."""
707
+
708
+ def __init__(self, config: ChatGLMConfig, device=None):
709
+ super(Embedding, self).__init__()
710
+
711
+ self.hidden_size = config.hidden_size
712
+ # Word embeddings (parallel).
713
+ self.word_embeddings = nn.Embedding(
714
+ config.padded_vocab_size,
715
+ self.hidden_size,
716
+ dtype=config.torch_dtype,
717
+ device=device
718
+ )
719
+ self.fp32_residual_connection = config.fp32_residual_connection
720
+
721
+ def forward(self, input_ids):
722
+ # Embeddings.
723
+ words_embeddings = self.word_embeddings(input_ids)
724
+ embeddings = words_embeddings
725
+ # Data format change to avoid explicit tranposes : [b s h] --> [s b h].
726
+ embeddings = embeddings.transpose(0, 1).contiguous()
727
+ # If the input flag for fp32 residual connection is set, convert for float.
728
+ if self.fp32_residual_connection:
729
+ embeddings = embeddings.float()
730
+ return embeddings
731
+
732
+
733
+ class ChatGLMModel(ChatGLMPreTrainedModel):
734
+ def __init__(self, config: ChatGLMConfig, device=None, empty_init=True):
735
+ super().__init__(config)
736
+ if empty_init:
737
+ init_method = skip_init
738
+ else:
739
+ init_method = default_init
740
+ init_kwargs = {}
741
+ if device is not None:
742
+ init_kwargs["device"] = device
743
+ self.embedding = init_method(Embedding, config, **init_kwargs)
744
+ self.num_layers = config.num_layers
745
+ self.multi_query_group_num = config.multi_query_group_num
746
+ self.kv_channels = config.kv_channels
747
+
748
+ # Rotary positional embeddings
749
+ self.seq_length = config.seq_length
750
+ rotary_dim = (
751
+ config.hidden_size // config.num_attention_heads if config.kv_channels is None else config.kv_channels
752
+ )
753
+
754
+ self.rotary_pos_emb = RotaryEmbedding(rotary_dim // 2, original_impl=config.original_rope, device=device,
755
+ dtype=config.torch_dtype)
756
+ self.encoder = init_method(GLMTransformer, config, **init_kwargs)
757
+ self.output_layer = init_method(nn.Linear, config.hidden_size, config.padded_vocab_size, bias=False,
758
+ dtype=config.torch_dtype, **init_kwargs)
759
+ self.pre_seq_len = config.pre_seq_len
760
+ self.prefix_projection = config.prefix_projection
761
+ if self.pre_seq_len is not None:
762
+ for param in self.parameters():
763
+ param.requires_grad = False
764
+ self.prefix_tokens = torch.arange(self.pre_seq_len).long()
765
+ self.prefix_encoder = PrefixEncoder(config)
766
+ self.dropout = torch.nn.Dropout(0.1)
767
+
768
+ def get_input_embeddings(self):
769
+ return self.embedding.word_embeddings
770
+
771
+ def get_prompt(self, batch_size, device, dtype=torch.half):
772
+ prefix_tokens = self.prefix_tokens.unsqueeze(0).expand(batch_size, -1).to(device)
773
+ past_key_values = self.prefix_encoder(prefix_tokens).type(dtype)
774
+ past_key_values = past_key_values.view(
775
+ batch_size,
776
+ self.pre_seq_len,
777
+ self.num_layers * 2,
778
+ self.multi_query_group_num,
779
+ self.kv_channels
780
+ )
781
+ # seq_len, b, nh, hidden_size
782
+ past_key_values = self.dropout(past_key_values)
783
+ past_key_values = past_key_values.permute([2, 1, 0, 3, 4]).split(2)
784
+ return past_key_values
785
+
786
+ def forward(
787
+ self,
788
+ input_ids,
789
+ position_ids: Optional[torch.Tensor] = None,
790
+ attention_mask: Optional[torch.BoolTensor] = None,
791
+ full_attention_mask: Optional[torch.BoolTensor] = None,
792
+ past_key_values: Optional[Tuple[Tuple[torch.Tensor, torch.Tensor], ...]] = None,
793
+ inputs_embeds: Optional[torch.Tensor] = None,
794
+ use_cache: Optional[bool] = None,
795
+ output_hidden_states: Optional[bool] = None,
796
+ return_dict: Optional[bool] = None,
797
+ ):
798
+ output_hidden_states = (
799
+ output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
800
+ )
801
+ use_cache = use_cache if use_cache is not None else self.config.use_cache
802
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
803
+
804
+ batch_size, seq_length = input_ids.shape
805
+
806
+ if inputs_embeds is None:
807
+ inputs_embeds = self.embedding(input_ids)
808
+
809
+ if self.pre_seq_len is not None:
810
+ if past_key_values is None:
811
+ past_key_values = self.get_prompt(batch_size=batch_size, device=input_ids.device,
812
+ dtype=inputs_embeds.dtype)
813
+ if attention_mask is not None:
814
+ attention_mask = torch.cat([attention_mask.new_ones((batch_size, self.pre_seq_len)),
815
+ attention_mask], dim=-1)
816
+
817
+ if full_attention_mask is None:
818
+ if (attention_mask is not None and not attention_mask.all()) or (past_key_values and seq_length != 1):
819
+ full_attention_mask = self.get_masks(input_ids, past_key_values, padding_mask=attention_mask)
820
+
821
+ # Rotary positional embeddings
822
+ rotary_pos_emb = self.rotary_pos_emb(self.seq_length)
823
+ if position_ids is not None:
824
+ rotary_pos_emb = rotary_pos_emb[position_ids]
825
+ else:
826
+ rotary_pos_emb = rotary_pos_emb[None, :seq_length]
827
+ rotary_pos_emb = rotary_pos_emb.transpose(0, 1).contiguous()
828
+
829
+ # Run encoder.
830
+ hidden_states, presents, all_hidden_states, all_self_attentions = self.encoder(
831
+ inputs_embeds, full_attention_mask, rotary_pos_emb=rotary_pos_emb,
832
+ kv_caches=past_key_values, use_cache=use_cache, output_hidden_states=output_hidden_states
833
+ )
834
+
835
+ if not return_dict:
836
+ return tuple(v for v in [hidden_states, presents, all_hidden_states, all_self_attentions] if v is not None)
837
+
838
+ return BaseModelOutputWithPast(
839
+ last_hidden_state=hidden_states,
840
+ past_key_values=presents,
841
+ hidden_states=all_hidden_states,
842
+ attentions=all_self_attentions,
843
+ )
844
+
845
+ def quantize(self, weight_bit_width: int):
846
+ from .quantization import quantize
847
+ quantize(self.encoder, weight_bit_width)
848
+ return self
849
+
850
+
851
+ class ChatGLMForConditionalGeneration(ChatGLMPreTrainedModel):
852
+ def __init__(self, config: ChatGLMConfig, empty_init=True, device=None):
853
+ super().__init__(config)
854
+
855
+ self.max_sequence_length = config.max_length
856
+ self.transformer = ChatGLMModel(config, empty_init=empty_init, device=device)
857
+ self.config = config
858
+ self.quantized = False
859
+
860
+ if self.config.quantization_bit:
861
+ self.quantize(self.config.quantization_bit, empty_init=True)
862
+
863
+ def _update_model_kwargs_for_generation(
864
+ self,
865
+ outputs: ModelOutput,
866
+ model_kwargs: Dict[str, Any],
867
+ is_encoder_decoder: bool = False,
868
+ standardize_cache_format: bool = False,
869
+ ) -> Dict[str, Any]:
870
+ # update past_key_values
871
+ model_kwargs["past_key_values"] = self._extract_past_from_model_output(
872
+ outputs, standardize_cache_format=standardize_cache_format
873
+ )
874
+
875
+ # update attention mask
876
+ if "attention_mask" in model_kwargs:
877
+ attention_mask = model_kwargs["attention_mask"]
878
+ model_kwargs["attention_mask"] = torch.cat(
879
+ [attention_mask, attention_mask.new_ones((attention_mask.shape[0], 1))], dim=-1
880
+ )
881
+
882
+ # update position ids
883
+ if "position_ids" in model_kwargs:
884
+ position_ids = model_kwargs["position_ids"]
885
+ new_position_id = position_ids[..., -1:].clone()
886
+ new_position_id += 1
887
+ model_kwargs["position_ids"] = torch.cat(
888
+ [position_ids, new_position_id], dim=-1
889
+ )
890
+
891
+ model_kwargs["is_first_forward"] = False
892
+ return model_kwargs
893
+
894
+ def prepare_inputs_for_generation(
895
+ self,
896
+ input_ids: torch.LongTensor,
897
+ past_key_values: Optional[torch.Tensor] = None,
898
+ attention_mask: Optional[torch.Tensor] = None,
899
+ position_ids: Optional[torch.Tensor] = None,
900
+ use_cache: Optional[bool] = None,
901
+ is_first_forward: bool = True,
902
+ **kwargs
903
+ ) -> dict:
904
+ # only last token for input_ids if past is not None
905
+ if position_ids is None:
906
+ position_ids = self.get_position_ids(input_ids, device=input_ids.device)
907
+ if not is_first_forward:
908
+ if past_key_values is not None:
909
+ position_ids = position_ids[..., -1:]
910
+ input_ids = input_ids[:, -1:]
911
+ return {
912
+ "input_ids": input_ids,
913
+ "past_key_values": past_key_values,
914
+ "position_ids": position_ids,
915
+ "attention_mask": attention_mask,
916
+ "return_last_logit": True,
917
+ "use_cache": use_cache
918
+ }
919
+
920
+ def forward(
921
+ self,
922
+ input_ids: Optional[torch.Tensor] = None,
923
+ position_ids: Optional[torch.Tensor] = None,
924
+ attention_mask: Optional[torch.Tensor] = None,
925
+ past_key_values: Optional[Tuple[torch.FloatTensor]] = None,
926
+ inputs_embeds: Optional[torch.Tensor] = None,
927
+ labels: Optional[torch.Tensor] = None,
928
+ use_cache: Optional[bool] = None,
929
+ output_attentions: Optional[bool] = None,
930
+ output_hidden_states: Optional[bool] = None,
931
+ return_dict: Optional[bool] = None,
932
+ return_last_logit: Optional[bool] = False,
933
+ ):
934
+ use_cache = use_cache if use_cache is not None else self.config.use_cache
935
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
936
+
937
+ transformer_outputs = self.transformer(
938
+ input_ids=input_ids,
939
+ position_ids=position_ids,
940
+ attention_mask=attention_mask,
941
+ past_key_values=past_key_values,
942
+ inputs_embeds=inputs_embeds,
943
+ use_cache=use_cache,
944
+ output_hidden_states=output_hidden_states,
945
+ return_dict=return_dict,
946
+ )
947
+
948
+ hidden_states = transformer_outputs[0]
949
+ if return_last_logit:
950
+ hidden_states = hidden_states[-1:]
951
+ lm_logits = self.transformer.output_layer(hidden_states)
952
+ lm_logits = lm_logits.transpose(0, 1).contiguous()
953
+
954
+ loss = None
955
+ if labels is not None:
956
+ lm_logits = lm_logits.to(torch.float32)
957
+
958
+ # Shift so that tokens < n predict n
959
+ shift_logits = lm_logits[..., :-1, :].contiguous()
960
+ shift_labels = labels[..., 1:].contiguous()
961
+ # Flatten the tokens
962
+ loss_fct = CrossEntropyLoss(ignore_index=-100)
963
+ loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))
964
+
965
+ lm_logits = lm_logits.to(hidden_states.dtype)
966
+ loss = loss.to(hidden_states.dtype)
967
+
968
+ if not return_dict:
969
+ output = (lm_logits,) + transformer_outputs[1:]
970
+ return ((loss,) + output) if loss is not None else output
971
+
972
+ return CausalLMOutputWithPast(
973
+ loss=loss,
974
+ logits=lm_logits,
975
+ past_key_values=transformer_outputs.past_key_values,
976
+ hidden_states=transformer_outputs.hidden_states,
977
+ attentions=transformer_outputs.attentions,
978
+ )
979
+
980
+ @staticmethod
981
+ def _reorder_cache(
982
+ past: Tuple[Tuple[torch.Tensor, torch.Tensor], ...], beam_idx: torch.LongTensor
983
+ ) -> Tuple[Tuple[torch.Tensor, torch.Tensor], ...]:
984
+ """
985
+ This function is used to re-order the `past_key_values` cache if [`~PreTrainedModel.beam_search`] or
986
+ [`~PreTrainedModel.beam_sample`] is called. This is required to match `past_key_values` with the correct
987
+ beam_idx at every generation step.
988
+
989
+ Output shares the same memory storage as `past`.
990
+ """
991
+ return tuple(
992
+ (
993
+ layer_past[0].index_select(1, beam_idx.to(layer_past[0].device)),
994
+ layer_past[1].index_select(1, beam_idx.to(layer_past[1].device)),
995
+ )
996
+ for layer_past in past
997
+ )
998
+
999
+ def process_response(self, response):
1000
+ response = response.strip()
1001
+ response = response.replace("[[训练时间]]", "2023年")
1002
+ return response
1003
+
1004
+ def build_inputs(self, tokenizer, query: str, history: List[Tuple[str, str]] = None):
1005
+ prompt = tokenizer.build_prompt(query, history=history)
1006
+ inputs = tokenizer([prompt], return_tensors="pt")
1007
+ inputs = inputs.to(self.device)
1008
+ return inputs
1009
+
1010
+ def build_stream_inputs(self, tokenizer, query: str, history: List[Tuple[str, str]] = None):
1011
+ if history:
1012
+ prompt = "\n\n[Round {}]\n\n问:{}\n\n答:".format(len(history) + 1, query)
1013
+ input_ids = tokenizer.encode(prompt, add_special_tokens=False)
1014
+ input_ids = input_ids[1:]
1015
+ inputs = tokenizer.batch_encode_plus([(input_ids, None)], return_tensors="pt", add_special_tokens=False)
1016
+ else:
1017
+ prompt = "[Round {}]\n\n问:{}\n\n答:".format(len(history) + 1, query)
1018
+ inputs = tokenizer([prompt], return_tensors="pt")
1019
+ inputs = inputs.to(self.device)
1020
+ return inputs
1021
+
1022
+ @torch.inference_mode()
1023
+ def chat(self, tokenizer, query: str, history: List[Tuple[str, str]] = None, max_length: int = 8192, num_beams=1,
1024
+ do_sample=True, top_p=0.8, temperature=0.8, logits_processor=None, **kwargs):
1025
+ if history is None:
1026
+ history = []
1027
+ if logits_processor is None:
1028
+ logits_processor = LogitsProcessorList()
1029
+ logits_processor.append(InvalidScoreLogitsProcessor())
1030
+ gen_kwargs = {"max_length": max_length, "num_beams": num_beams, "do_sample": do_sample, "top_p": top_p,
1031
+ "temperature": temperature, "logits_processor": logits_processor, **kwargs}
1032
+ inputs = self.build_inputs(tokenizer, query, history=history)
1033
+ outputs = self.generate(**inputs, **gen_kwargs)
1034
+ outputs = outputs.tolist()[0][len(inputs["input_ids"][0]):]
1035
+ response = tokenizer.decode(outputs)
1036
+ response = self.process_response(response)
1037
+ history = history + [(query, response)]
1038
+ return response, history
1039
+
1040
+ @torch.inference_mode()
1041
+ def stream_chat(self, tokenizer, query: str, history: List[Tuple[str, str]] = None, past_key_values=None,
1042
+ max_length: int = 8192, do_sample=True, top_p=0.8, temperature=0.8, logits_processor=None,
1043
+ return_past_key_values=False, **kwargs):
1044
+ if history is None:
1045
+ history = []
1046
+ if logits_processor is None:
1047
+ logits_processor = LogitsProcessorList()
1048
+ logits_processor.append(InvalidScoreLogitsProcessor())
1049
+ gen_kwargs = {"max_length": max_length, "do_sample": do_sample, "top_p": top_p,
1050
+ "temperature": temperature, "logits_processor": logits_processor, **kwargs}
1051
+ if past_key_values is None and not return_past_key_values:
1052
+ inputs = self.build_inputs(tokenizer, query, history=history)
1053
+ else:
1054
+ inputs = self.build_stream_inputs(tokenizer, query, history=history)
1055
+ if past_key_values is not None:
1056
+ past_length = past_key_values[0][0].shape[0]
1057
+ if self.transformer.pre_seq_len is not None:
1058
+ past_length -= self.transformer.pre_seq_len
1059
+ inputs.position_ids += past_length
1060
+ attention_mask = inputs.attention_mask
1061
+ attention_mask = torch.cat((attention_mask.new_ones(1, past_length), attention_mask), dim=1)
1062
+ inputs['attention_mask'] = attention_mask
1063
+ for outputs in self.stream_generate(**inputs, past_key_values=past_key_values,
1064
+ return_past_key_values=return_past_key_values, **gen_kwargs):
1065
+ if return_past_key_values:
1066
+ outputs, past_key_values = outputs
1067
+ outputs = outputs.tolist()[0][len(inputs["input_ids"][0]):]
1068
+ response = tokenizer.decode(outputs)
1069
+ if response and response[-1] != "�":
1070
+ response = self.process_response(response)
1071
+ new_history = history + [(query, response)]
1072
+ if return_past_key_values:
1073
+ yield response, new_history, past_key_values
1074
+ else:
1075
+ yield response, new_history
1076
+
1077
+ @torch.inference_mode()
1078
+ def stream_generate(
1079
+ self,
1080
+ input_ids,
1081
+ generation_config: Optional[GenerationConfig] = None,
1082
+ logits_processor: Optional[LogitsProcessorList] = None,
1083
+ stopping_criteria: Optional[StoppingCriteriaList] = None,
1084
+ prefix_allowed_tokens_fn: Optional[Callable[[int, torch.Tensor], List[int]]] = None,
1085
+ return_past_key_values=False,
1086
+ **kwargs,
1087
+ ):
1088
+ batch_size, input_ids_seq_length = input_ids.shape[0], input_ids.shape[-1]
1089
+
1090
+ if generation_config is None:
1091
+ generation_config = self.generation_config
1092
+ generation_config = copy.deepcopy(generation_config)
1093
+ model_kwargs = generation_config.update(**kwargs)
1094
+ model_kwargs["use_cache"] = generation_config.use_cache
1095
+ bos_token_id, eos_token_id = generation_config.bos_token_id, generation_config.eos_token_id
1096
+
1097
+ if isinstance(eos_token_id, int):
1098
+ eos_token_id = [eos_token_id]
1099
+
1100
+ has_default_max_length = kwargs.get("max_length") is None and generation_config.max_length is not None
1101
+ if has_default_max_length and generation_config.max_new_tokens is None:
1102
+ warnings.warn(
1103
+ f"Using `max_length`'s default ({generation_config.max_length}) to control the generation length. "
1104
+ "This behaviour is deprecated and will be removed from the config in v5 of Transformers -- we"
1105
+ " recommend using `max_new_tokens` to control the maximum length of the generation.",
1106
+ UserWarning,
1107
+ )
1108
+ elif generation_config.max_new_tokens is not None:
1109
+ generation_config.max_length = generation_config.max_new_tokens + input_ids_seq_length
1110
+ if not has_default_max_length:
1111
+ logger.warn(
1112
+ f"Both `max_new_tokens` (={generation_config.max_new_tokens}) and `max_length`(="
1113
+ f"{generation_config.max_length}) seem to have been set. `max_new_tokens` will take precedence. "
1114
+ "Please refer to the documentation for more information. "
1115
+ "(https://huggingface.co/docs/transformers/main/en/main_classes/text_generation)",
1116
+ UserWarning,
1117
+ )
1118
+
1119
+ if input_ids_seq_length >= generation_config.max_length:
1120
+ input_ids_string = "decoder_input_ids" if self.config.is_encoder_decoder else "input_ids"
1121
+ logger.warning(
1122
+ f"Input length of {input_ids_string} is {input_ids_seq_length}, but `max_length` is set to"
1123
+ f" {generation_config.max_length}. This can lead to unexpected behavior. You should consider"
1124
+ " increasing `max_new_tokens`."
1125
+ )
1126
+
1127
+ # 2. Set generation parameters if not already defined
1128
+ logits_processor = logits_processor if logits_processor is not None else LogitsProcessorList()
1129
+ stopping_criteria = stopping_criteria if stopping_criteria is not None else StoppingCriteriaList()
1130
+
1131
+ logits_processor = self._get_logits_processor(
1132
+ generation_config=generation_config,
1133
+ input_ids_seq_length=input_ids_seq_length,
1134
+ encoder_input_ids=input_ids,
1135
+ prefix_allowed_tokens_fn=prefix_allowed_tokens_fn,
1136
+ logits_processor=logits_processor,
1137
+ )
1138
+
1139
+ stopping_criteria = self._get_stopping_criteria(
1140
+ generation_config=generation_config, stopping_criteria=stopping_criteria
1141
+ )
1142
+ logits_warper = self._get_logits_warper(generation_config)
1143
+
1144
+ unfinished_sequences = input_ids.new(input_ids.shape[0]).fill_(1)
1145
+ scores = None
1146
+ while True:
1147
+ model_inputs = self.prepare_inputs_for_generation(input_ids, **model_kwargs)
1148
+ # forward pass to get next token
1149
+ outputs = self(
1150
+ **model_inputs,
1151
+ return_dict=True,
1152
+ output_attentions=False,
1153
+ output_hidden_states=False,
1154
+ )
1155
+
1156
+ next_token_logits = outputs.logits[:, -1, :]
1157
+
1158
+ # pre-process distribution
1159
+ next_token_scores = logits_processor(input_ids, next_token_logits)
1160
+ next_token_scores = logits_warper(input_ids, next_token_scores)
1161
+
1162
+ # sample
1163
+ probs = nn.functional.softmax(next_token_scores, dim=-1)
1164
+ if generation_config.do_sample:
1165
+ next_tokens = torch.multinomial(probs, num_samples=1).squeeze(1)
1166
+ else:
1167
+ next_tokens = torch.argmax(probs, dim=-1)
1168
+
1169
+ # update generated ids, model inputs, and length for next step
1170
+ input_ids = torch.cat([input_ids, next_tokens[:, None]], dim=-1)
1171
+ model_kwargs = self._update_model_kwargs_for_generation(
1172
+ outputs, model_kwargs, is_encoder_decoder=self.config.is_encoder_decoder
1173
+ )
1174
+ unfinished_sequences = unfinished_sequences.mul((sum(next_tokens != i for i in eos_token_id)).long())
1175
+ if return_past_key_values:
1176
+ yield input_ids, outputs.past_key_values
1177
+ else:
1178
+ yield input_ids
1179
+ # stop when each sentence is finished, or if we exceed the maximum length
1180
+ if unfinished_sequences.max() == 0 or stopping_criteria(input_ids, scores):
1181
+ break
1182
+
1183
+ def quantize(self, bits: int, empty_init=False, device=None, **kwargs):
1184
+ if bits == 0:
1185
+ return
1186
+
1187
+ from .quantization import quantize
1188
+
1189
+ if self.quantized:
1190
+ logger.info("Already quantized.")
1191
+ return self
1192
+
1193
+ self.quantized = True
1194
+
1195
+ self.config.quantization_bit = bits
1196
+
1197
+ self.transformer.encoder = quantize(self.transformer.encoder, bits, empty_init=empty_init, device=device,
1198
+ **kwargs)
1199
+ return self
1200
+
1201
+
1202
+ class ChatGLMForSequenceClassification(ChatGLMPreTrainedModel):
1203
+ def __init__(self, config: ChatGLMConfig, empty_init=True, device=None):
1204
+ super().__init__(config)
1205
+
1206
+ self.num_labels = config.num_labels
1207
+ self.transformer = ChatGLMModel(config, empty_init=empty_init, device=device)
1208
+
1209
+ self.classifier_head = nn.Linear(config.hidden_size, config.num_labels, bias=True, dtype=torch.half)
1210
+ if config.classifier_dropout is not None:
1211
+ self.dropout = nn.Dropout(config.classifier_dropout)
1212
+ else:
1213
+ self.dropout = None
1214
+ self.config = config
1215
+
1216
+ if self.config.quantization_bit:
1217
+ self.quantize(self.config.quantization_bit, empty_init=True)
1218
+
1219
+ def forward(
1220
+ self,
1221
+ input_ids: Optional[torch.LongTensor] = None,
1222
+ position_ids: Optional[torch.LongTensor] = None,
1223
+ attention_mask: Optional[torch.Tensor] = None,
1224
+ full_attention_mask: Optional[torch.Tensor] = None,
1225
+ past_key_values: Optional[Tuple[Tuple[torch.Tensor, torch.Tensor], ...]] = None,
1226
+ inputs_embeds: Optional[torch.LongTensor] = None,
1227
+ labels: Optional[torch.LongTensor] = None,
1228
+ use_cache: Optional[bool] = None,
1229
+ output_hidden_states: Optional[bool] = None,
1230
+ return_dict: Optional[bool] = None,
1231
+ ) -> Union[Tuple[torch.Tensor, ...], SequenceClassifierOutputWithPast]:
1232
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
1233
+
1234
+ transformer_outputs = self.transformer(
1235
+ input_ids=input_ids,
1236
+ position_ids=position_ids,
1237
+ attention_mask=attention_mask,
1238
+ full_attention_mask=full_attention_mask,
1239
+ past_key_values=past_key_values,
1240
+ inputs_embeds=inputs_embeds,
1241
+ use_cache=use_cache,
1242
+ output_hidden_states=output_hidden_states,
1243
+ return_dict=return_dict,
1244
+ )
1245
+
1246
+ hidden_states = transformer_outputs[0]
1247
+ pooled_hidden_states = hidden_states[-1]
1248
+ if self.dropout is not None:
1249
+ pooled_hidden_states = self.dropout(pooled_hidden_states)
1250
+ logits = self.classifier_head(pooled_hidden_states)
1251
+
1252
+ loss = None
1253
+ if labels is not None:
1254
+ if self.config.problem_type is None:
1255
+ if self.num_labels == 1:
1256
+ self.config.problem_type = "regression"
1257
+ elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
1258
+ self.config.problem_type = "single_label_classification"
1259
+ else:
1260
+ self.config.problem_type = "multi_label_classification"
1261
+
1262
+ if self.config.problem_type == "regression":
1263
+ loss_fct = MSELoss()
1264
+ if self.num_labels == 1:
1265
+ loss = loss_fct(logits.squeeze().float(), labels.squeeze())
1266
+ else:
1267
+ loss = loss_fct(logits.float(), labels)
1268
+ elif self.config.problem_type == "single_label_classification":
1269
+ loss_fct = CrossEntropyLoss()
1270
+ loss = loss_fct(logits.view(-1, self.num_labels).float(), labels.view(-1))
1271
+ elif self.config.problem_type == "multi_label_classification":
1272
+ loss_fct = BCEWithLogitsLoss()
1273
+ loss = loss_fct(logits.float(), labels.view(-1, self.num_labels))
1274
+
1275
+ if not return_dict:
1276
+ output = (logits,) + transformer_outputs[1:]
1277
+ return ((loss,) + output) if loss is not None else output
1278
+
1279
+ return SequenceClassifierOutputWithPast(
1280
+ loss=loss,
1281
+ logits=logits,
1282
+ past_key_values=transformer_outputs.past_key_values,
1283
+ hidden_states=transformer_outputs.hidden_states,
1284
+ attentions=transformer_outputs.attentions,
1285
+ )
pytorch_model-00001-of-00002.bin ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:0a455d80479b771e8579dd9628bbafbfa573198e70698c0ed9fdaf9eae22c9b7
3
+ size 9986264643
pytorch_model-00002-of-00002.bin ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:603c5280e429908980c7ef6534657f2dfe9514a8385e73556e8b5566ff5844f7
3
+ size 2500975567
pytorch_model.bin.index.json ADDED
@@ -0,0 +1,207 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "metadata": {
3
+ "total_size": 12487168064
4
+ },
5
+ "weight_map": {
6
+ "transformer.embedding.word_embeddings.weight": "pytorch_model-00001-of-00002.bin",
7
+ "transformer.encoder.final_layernorm.weight": "pytorch_model-00002-of-00002.bin",
8
+ "transformer.encoder.layers.0.input_layernorm.weight": "pytorch_model-00001-of-00002.bin",
9
+ "transformer.encoder.layers.0.mlp.dense_4h_to_h.weight": "pytorch_model-00001-of-00002.bin",
10
+ "transformer.encoder.layers.0.mlp.dense_h_to_4h.weight": "pytorch_model-00001-of-00002.bin",
11
+ "transformer.encoder.layers.0.post_attention_layernorm.weight": "pytorch_model-00001-of-00002.bin",
12
+ "transformer.encoder.layers.0.self_attention.dense.weight": "pytorch_model-00001-of-00002.bin",
13
+ "transformer.encoder.layers.0.self_attention.query_key_value.bias": "pytorch_model-00001-of-00002.bin",
14
+ "transformer.encoder.layers.0.self_attention.query_key_value.weight": "pytorch_model-00001-of-00002.bin",
15
+ "transformer.encoder.layers.1.input_layernorm.weight": "pytorch_model-00001-of-00002.bin",
16
+ "transformer.encoder.layers.1.mlp.dense_4h_to_h.weight": "pytorch_model-00001-of-00002.bin",
17
+ "transformer.encoder.layers.1.mlp.dense_h_to_4h.weight": "pytorch_model-00001-of-00002.bin",
18
+ "transformer.encoder.layers.1.post_attention_layernorm.weight": "pytorch_model-00001-of-00002.bin",
19
+ "transformer.encoder.layers.1.self_attention.dense.weight": "pytorch_model-00001-of-00002.bin",
20
+ "transformer.encoder.layers.1.self_attention.query_key_value.bias": "pytorch_model-00001-of-00002.bin",
21
+ "transformer.encoder.layers.1.self_attention.query_key_value.weight": "pytorch_model-00001-of-00002.bin",
22
+ "transformer.encoder.layers.10.input_layernorm.weight": "pytorch_model-00001-of-00002.bin",
23
+ "transformer.encoder.layers.10.mlp.dense_4h_to_h.weight": "pytorch_model-00001-of-00002.bin",
24
+ "transformer.encoder.layers.10.mlp.dense_h_to_4h.weight": "pytorch_model-00001-of-00002.bin",
25
+ "transformer.encoder.layers.10.post_attention_layernorm.weight": "pytorch_model-00001-of-00002.bin",
26
+ "transformer.encoder.layers.10.self_attention.dense.weight": "pytorch_model-00001-of-00002.bin",
27
+ "transformer.encoder.layers.10.self_attention.query_key_value.bias": "pytorch_model-00001-of-00002.bin",
28
+ "transformer.encoder.layers.10.self_attention.query_key_value.weight": "pytorch_model-00001-of-00002.bin",
29
+ "transformer.encoder.layers.11.input_layernorm.weight": "pytorch_model-00001-of-00002.bin",
30
+ "transformer.encoder.layers.11.mlp.dense_4h_to_h.weight": "pytorch_model-00001-of-00002.bin",
31
+ "transformer.encoder.layers.11.mlp.dense_h_to_4h.weight": "pytorch_model-00001-of-00002.bin",
32
+ "transformer.encoder.layers.11.post_attention_layernorm.weight": "pytorch_model-00001-of-00002.bin",
33
+ "transformer.encoder.layers.11.self_attention.dense.weight": "pytorch_model-00001-of-00002.bin",
34
+ "transformer.encoder.layers.11.self_attention.query_key_value.bias": "pytorch_model-00001-of-00002.bin",
35
+ "transformer.encoder.layers.11.self_attention.query_key_value.weight": "pytorch_model-00001-of-00002.bin",
36
+ "transformer.encoder.layers.12.input_layernorm.weight": "pytorch_model-00001-of-00002.bin",
37
+ "transformer.encoder.layers.12.mlp.dense_4h_to_h.weight": "pytorch_model-00001-of-00002.bin",
38
+ "transformer.encoder.layers.12.mlp.dense_h_to_4h.weight": "pytorch_model-00001-of-00002.bin",
39
+ "transformer.encoder.layers.12.post_attention_layernorm.weight": "pytorch_model-00001-of-00002.bin",
40
+ "transformer.encoder.layers.12.self_attention.dense.weight": "pytorch_model-00001-of-00002.bin",
41
+ "transformer.encoder.layers.12.self_attention.query_key_value.bias": "pytorch_model-00001-of-00002.bin",
42
+ "transformer.encoder.layers.12.self_attention.query_key_value.weight": "pytorch_model-00001-of-00002.bin",
43
+ "transformer.encoder.layers.13.input_layernorm.weight": "pytorch_model-00001-of-00002.bin",
44
+ "transformer.encoder.layers.13.mlp.dense_4h_to_h.weight": "pytorch_model-00001-of-00002.bin",
45
+ "transformer.encoder.layers.13.mlp.dense_h_to_4h.weight": "pytorch_model-00001-of-00002.bin",
46
+ "transformer.encoder.layers.13.post_attention_layernorm.weight": "pytorch_model-00001-of-00002.bin",
47
+ "transformer.encoder.layers.13.self_attention.dense.weight": "pytorch_model-00001-of-00002.bin",
48
+ "transformer.encoder.layers.13.self_attention.query_key_value.bias": "pytorch_model-00001-of-00002.bin",
49
+ "transformer.encoder.layers.13.self_attention.query_key_value.weight": "pytorch_model-00001-of-00002.bin",
50
+ "transformer.encoder.layers.14.input_layernorm.weight": "pytorch_model-00001-of-00002.bin",
51
+ "transformer.encoder.layers.14.mlp.dense_4h_to_h.weight": "pytorch_model-00001-of-00002.bin",
52
+ "transformer.encoder.layers.14.mlp.dense_h_to_4h.weight": "pytorch_model-00001-of-00002.bin",
53
+ "transformer.encoder.layers.14.post_attention_layernorm.weight": "pytorch_model-00001-of-00002.bin",
54
+ "transformer.encoder.layers.14.self_attention.dense.weight": "pytorch_model-00001-of-00002.bin",
55
+ "transformer.encoder.layers.14.self_attention.query_key_value.bias": "pytorch_model-00001-of-00002.bin",
56
+ "transformer.encoder.layers.14.self_attention.query_key_value.weight": "pytorch_model-00001-of-00002.bin",
57
+ "transformer.encoder.layers.15.input_layernorm.weight": "pytorch_model-00001-of-00002.bin",
58
+ "transformer.encoder.layers.15.mlp.dense_4h_to_h.weight": "pytorch_model-00001-of-00002.bin",
59
+ "transformer.encoder.layers.15.mlp.dense_h_to_4h.weight": "pytorch_model-00001-of-00002.bin",
60
+ "transformer.encoder.layers.15.post_attention_layernorm.weight": "pytorch_model-00001-of-00002.bin",
61
+ "transformer.encoder.layers.15.self_attention.dense.weight": "pytorch_model-00001-of-00002.bin",
62
+ "transformer.encoder.layers.15.self_attention.query_key_value.bias": "pytorch_model-00001-of-00002.bin",
63
+ "transformer.encoder.layers.15.self_attention.query_key_value.weight": "pytorch_model-00001-of-00002.bin",
64
+ "transformer.encoder.layers.16.input_layernorm.weight": "pytorch_model-00001-of-00002.bin",
65
+ "transformer.encoder.layers.16.mlp.dense_4h_to_h.weight": "pytorch_model-00001-of-00002.bin",
66
+ "transformer.encoder.layers.16.mlp.dense_h_to_4h.weight": "pytorch_model-00001-of-00002.bin",
67
+ "transformer.encoder.layers.16.post_attention_layernorm.weight": "pytorch_model-00001-of-00002.bin",
68
+ "transformer.encoder.layers.16.self_attention.dense.weight": "pytorch_model-00001-of-00002.bin",
69
+ "transformer.encoder.layers.16.self_attention.query_key_value.bias": "pytorch_model-00001-of-00002.bin",
70
+ "transformer.encoder.layers.16.self_attention.query_key_value.weight": "pytorch_model-00001-of-00002.bin",
71
+ "transformer.encoder.layers.17.input_layernorm.weight": "pytorch_model-00001-of-00002.bin",
72
+ "transformer.encoder.layers.17.mlp.dense_4h_to_h.weight": "pytorch_model-00001-of-00002.bin",
73
+ "transformer.encoder.layers.17.mlp.dense_h_to_4h.weight": "pytorch_model-00001-of-00002.bin",
74
+ "transformer.encoder.layers.17.post_attention_layernorm.weight": "pytorch_model-00001-of-00002.bin",
75
+ "transformer.encoder.layers.17.self_attention.dense.weight": "pytorch_model-00001-of-00002.bin",
76
+ "transformer.encoder.layers.17.self_attention.query_key_value.bias": "pytorch_model-00001-of-00002.bin",
77
+ "transformer.encoder.layers.17.self_attention.query_key_value.weight": "pytorch_model-00001-of-00002.bin",
78
+ "transformer.encoder.layers.18.input_layernorm.weight": "pytorch_model-00001-of-00002.bin",
79
+ "transformer.encoder.layers.18.mlp.dense_4h_to_h.weight": "pytorch_model-00001-of-00002.bin",
80
+ "transformer.encoder.layers.18.mlp.dense_h_to_4h.weight": "pytorch_model-00001-of-00002.bin",
81
+ "transformer.encoder.layers.18.post_attention_layernorm.weight": "pytorch_model-00001-of-00002.bin",
82
+ "transformer.encoder.layers.18.self_attention.dense.weight": "pytorch_model-00001-of-00002.bin",
83
+ "transformer.encoder.layers.18.self_attention.query_key_value.bias": "pytorch_model-00001-of-00002.bin",
84
+ "transformer.encoder.layers.18.self_attention.query_key_value.weight": "pytorch_model-00001-of-00002.bin",
85
+ "transformer.encoder.layers.19.input_layernorm.weight": "pytorch_model-00001-of-00002.bin",
86
+ "transformer.encoder.layers.19.mlp.dense_4h_to_h.weight": "pytorch_model-00001-of-00002.bin",
87
+ "transformer.encoder.layers.19.mlp.dense_h_to_4h.weight": "pytorch_model-00001-of-00002.bin",
88
+ "transformer.encoder.layers.19.post_attention_layernorm.weight": "pytorch_model-00001-of-00002.bin",
89
+ "transformer.encoder.layers.19.self_attention.dense.weight": "pytorch_model-00001-of-00002.bin",
90
+ "transformer.encoder.layers.19.self_attention.query_key_value.bias": "pytorch_model-00001-of-00002.bin",
91
+ "transformer.encoder.layers.19.self_attention.query_key_value.weight": "pytorch_model-00001-of-00002.bin",
92
+ "transformer.encoder.layers.2.input_layernorm.weight": "pytorch_model-00001-of-00002.bin",
93
+ "transformer.encoder.layers.2.mlp.dense_4h_to_h.weight": "pytorch_model-00001-of-00002.bin",
94
+ "transformer.encoder.layers.2.mlp.dense_h_to_4h.weight": "pytorch_model-00001-of-00002.bin",
95
+ "transformer.encoder.layers.2.post_attention_layernorm.weight": "pytorch_model-00001-of-00002.bin",
96
+ "transformer.encoder.layers.2.self_attention.dense.weight": "pytorch_model-00001-of-00002.bin",
97
+ "transformer.encoder.layers.2.self_attention.query_key_value.bias": "pytorch_model-00001-of-00002.bin",
98
+ "transformer.encoder.layers.2.self_attention.query_key_value.weight": "pytorch_model-00001-of-00002.bin",
99
+ "transformer.encoder.layers.20.input_layernorm.weight": "pytorch_model-00001-of-00002.bin",
100
+ "transformer.encoder.layers.20.mlp.dense_4h_to_h.weight": "pytorch_model-00001-of-00002.bin",
101
+ "transformer.encoder.layers.20.mlp.dense_h_to_4h.weight": "pytorch_model-00001-of-00002.bin",
102
+ "transformer.encoder.layers.20.post_attention_layernorm.weight": "pytorch_model-00001-of-00002.bin",
103
+ "transformer.encoder.layers.20.self_attention.dense.weight": "pytorch_model-00001-of-00002.bin",
104
+ "transformer.encoder.layers.20.self_attention.query_key_value.bias": "pytorch_model-00001-of-00002.bin",
105
+ "transformer.encoder.layers.20.self_attention.query_key_value.weight": "pytorch_model-00001-of-00002.bin",
106
+ "transformer.encoder.layers.21.input_layernorm.weight": "pytorch_model-00001-of-00002.bin",
107
+ "transformer.encoder.layers.21.mlp.dense_4h_to_h.weight": "pytorch_model-00001-of-00002.bin",
108
+ "transformer.encoder.layers.21.mlp.dense_h_to_4h.weight": "pytorch_model-00001-of-00002.bin",
109
+ "transformer.encoder.layers.21.post_attention_layernorm.weight": "pytorch_model-00001-of-00002.bin",
110
+ "transformer.encoder.layers.21.self_attention.dense.weight": "pytorch_model-00001-of-00002.bin",
111
+ "transformer.encoder.layers.21.self_attention.query_key_value.bias": "pytorch_model-00001-of-00002.bin",
112
+ "transformer.encoder.layers.21.self_attention.query_key_value.weight": "pytorch_model-00001-of-00002.bin",
113
+ "transformer.encoder.layers.22.input_layernorm.weight": "pytorch_model-00001-of-00002.bin",
114
+ "transformer.encoder.layers.22.mlp.dense_4h_to_h.weight": "pytorch_model-00001-of-00002.bin",
115
+ "transformer.encoder.layers.22.mlp.dense_h_to_4h.weight": "pytorch_model-00001-of-00002.bin",
116
+ "transformer.encoder.layers.22.post_attention_layernorm.weight": "pytorch_model-00001-of-00002.bin",
117
+ "transformer.encoder.layers.22.self_attention.dense.weight": "pytorch_model-00001-of-00002.bin",
118
+ "transformer.encoder.layers.22.self_attention.query_key_value.bias": "pytorch_model-00001-of-00002.bin",
119
+ "transformer.encoder.layers.22.self_attention.query_key_value.weight": "pytorch_model-00001-of-00002.bin",
120
+ "transformer.encoder.layers.23.input_layernorm.weight": "pytorch_model-00001-of-00002.bin",
121
+ "transformer.encoder.layers.23.mlp.dense_4h_to_h.weight": "pytorch_model-00002-of-00002.bin",
122
+ "transformer.encoder.layers.23.mlp.dense_h_to_4h.weight": "pytorch_model-00002-of-00002.bin",
123
+ "transformer.encoder.layers.23.post_attention_layernorm.weight": "pytorch_model-00001-of-00002.bin",
124
+ "transformer.encoder.layers.23.self_attention.dense.weight": "pytorch_model-00001-of-00002.bin",
125
+ "transformer.encoder.layers.23.self_attention.query_key_value.bias": "pytorch_model-00001-of-00002.bin",
126
+ "transformer.encoder.layers.23.self_attention.query_key_value.weight": "pytorch_model-00001-of-00002.bin",
127
+ "transformer.encoder.layers.24.input_layernorm.weight": "pytorch_model-00002-of-00002.bin",
128
+ "transformer.encoder.layers.24.mlp.dense_4h_to_h.weight": "pytorch_model-00002-of-00002.bin",
129
+ "transformer.encoder.layers.24.mlp.dense_h_to_4h.weight": "pytorch_model-00002-of-00002.bin",
130
+ "transformer.encoder.layers.24.post_attention_layernorm.weight": "pytorch_model-00002-of-00002.bin",
131
+ "transformer.encoder.layers.24.self_attention.dense.weight": "pytorch_model-00002-of-00002.bin",
132
+ "transformer.encoder.layers.24.self_attention.query_key_value.bias": "pytorch_model-00002-of-00002.bin",
133
+ "transformer.encoder.layers.24.self_attention.query_key_value.weight": "pytorch_model-00002-of-00002.bin",
134
+ "transformer.encoder.layers.25.input_layernorm.weight": "pytorch_model-00002-of-00002.bin",
135
+ "transformer.encoder.layers.25.mlp.dense_4h_to_h.weight": "pytorch_model-00002-of-00002.bin",
136
+ "transformer.encoder.layers.25.mlp.dense_h_to_4h.weight": "pytorch_model-00002-of-00002.bin",
137
+ "transformer.encoder.layers.25.post_attention_layernorm.weight": "pytorch_model-00002-of-00002.bin",
138
+ "transformer.encoder.layers.25.self_attention.dense.weight": "pytorch_model-00002-of-00002.bin",
139
+ "transformer.encoder.layers.25.self_attention.query_key_value.bias": "pytorch_model-00002-of-00002.bin",
140
+ "transformer.encoder.layers.25.self_attention.query_key_value.weight": "pytorch_model-00002-of-00002.bin",
141
+ "transformer.encoder.layers.26.input_layernorm.weight": "pytorch_model-00002-of-00002.bin",
142
+ "transformer.encoder.layers.26.mlp.dense_4h_to_h.weight": "pytorch_model-00002-of-00002.bin",
143
+ "transformer.encoder.layers.26.mlp.dense_h_to_4h.weight": "pytorch_model-00002-of-00002.bin",
144
+ "transformer.encoder.layers.26.post_attention_layernorm.weight": "pytorch_model-00002-of-00002.bin",
145
+ "transformer.encoder.layers.26.self_attention.dense.weight": "pytorch_model-00002-of-00002.bin",
146
+ "transformer.encoder.layers.26.self_attention.query_key_value.bias": "pytorch_model-00002-of-00002.bin",
147
+ "transformer.encoder.layers.26.self_attention.query_key_value.weight": "pytorch_model-00002-of-00002.bin",
148
+ "transformer.encoder.layers.27.input_layernorm.weight": "pytorch_model-00002-of-00002.bin",
149
+ "transformer.encoder.layers.27.mlp.dense_4h_to_h.weight": "pytorch_model-00002-of-00002.bin",
150
+ "transformer.encoder.layers.27.mlp.dense_h_to_4h.weight": "pytorch_model-00002-of-00002.bin",
151
+ "transformer.encoder.layers.27.post_attention_layernorm.weight": "pytorch_model-00002-of-00002.bin",
152
+ "transformer.encoder.layers.27.self_attention.dense.weight": "pytorch_model-00002-of-00002.bin",
153
+ "transformer.encoder.layers.27.self_attention.query_key_value.bias": "pytorch_model-00002-of-00002.bin",
154
+ "transformer.encoder.layers.27.self_attention.query_key_value.weight": "pytorch_model-00002-of-00002.bin",
155
+ "transformer.encoder.layers.3.input_layernorm.weight": "pytorch_model-00001-of-00002.bin",
156
+ "transformer.encoder.layers.3.mlp.dense_4h_to_h.weight": "pytorch_model-00001-of-00002.bin",
157
+ "transformer.encoder.layers.3.mlp.dense_h_to_4h.weight": "pytorch_model-00001-of-00002.bin",
158
+ "transformer.encoder.layers.3.post_attention_layernorm.weight": "pytorch_model-00001-of-00002.bin",
159
+ "transformer.encoder.layers.3.self_attention.dense.weight": "pytorch_model-00001-of-00002.bin",
160
+ "transformer.encoder.layers.3.self_attention.query_key_value.bias": "pytorch_model-00001-of-00002.bin",
161
+ "transformer.encoder.layers.3.self_attention.query_key_value.weight": "pytorch_model-00001-of-00002.bin",
162
+ "transformer.encoder.layers.4.input_layernorm.weight": "pytorch_model-00001-of-00002.bin",
163
+ "transformer.encoder.layers.4.mlp.dense_4h_to_h.weight": "pytorch_model-00001-of-00002.bin",
164
+ "transformer.encoder.layers.4.mlp.dense_h_to_4h.weight": "pytorch_model-00001-of-00002.bin",
165
+ "transformer.encoder.layers.4.post_attention_layernorm.weight": "pytorch_model-00001-of-00002.bin",
166
+ "transformer.encoder.layers.4.self_attention.dense.weight": "pytorch_model-00001-of-00002.bin",
167
+ "transformer.encoder.layers.4.self_attention.query_key_value.bias": "pytorch_model-00001-of-00002.bin",
168
+ "transformer.encoder.layers.4.self_attention.query_key_value.weight": "pytorch_model-00001-of-00002.bin",
169
+ "transformer.encoder.layers.5.input_layernorm.weight": "pytorch_model-00001-of-00002.bin",
170
+ "transformer.encoder.layers.5.mlp.dense_4h_to_h.weight": "pytorch_model-00001-of-00002.bin",
171
+ "transformer.encoder.layers.5.mlp.dense_h_to_4h.weight": "pytorch_model-00001-of-00002.bin",
172
+ "transformer.encoder.layers.5.post_attention_layernorm.weight": "pytorch_model-00001-of-00002.bin",
173
+ "transformer.encoder.layers.5.self_attention.dense.weight": "pytorch_model-00001-of-00002.bin",
174
+ "transformer.encoder.layers.5.self_attention.query_key_value.bias": "pytorch_model-00001-of-00002.bin",
175
+ "transformer.encoder.layers.5.self_attention.query_key_value.weight": "pytorch_model-00001-of-00002.bin",
176
+ "transformer.encoder.layers.6.input_layernorm.weight": "pytorch_model-00001-of-00002.bin",
177
+ "transformer.encoder.layers.6.mlp.dense_4h_to_h.weight": "pytorch_model-00001-of-00002.bin",
178
+ "transformer.encoder.layers.6.mlp.dense_h_to_4h.weight": "pytorch_model-00001-of-00002.bin",
179
+ "transformer.encoder.layers.6.post_attention_layernorm.weight": "pytorch_model-00001-of-00002.bin",
180
+ "transformer.encoder.layers.6.self_attention.dense.weight": "pytorch_model-00001-of-00002.bin",
181
+ "transformer.encoder.layers.6.self_attention.query_key_value.bias": "pytorch_model-00001-of-00002.bin",
182
+ "transformer.encoder.layers.6.self_attention.query_key_value.weight": "pytorch_model-00001-of-00002.bin",
183
+ "transformer.encoder.layers.7.input_layernorm.weight": "pytorch_model-00001-of-00002.bin",
184
+ "transformer.encoder.layers.7.mlp.dense_4h_to_h.weight": "pytorch_model-00001-of-00002.bin",
185
+ "transformer.encoder.layers.7.mlp.dense_h_to_4h.weight": "pytorch_model-00001-of-00002.bin",
186
+ "transformer.encoder.layers.7.post_attention_layernorm.weight": "pytorch_model-00001-of-00002.bin",
187
+ "transformer.encoder.layers.7.self_attention.dense.weight": "pytorch_model-00001-of-00002.bin",
188
+ "transformer.encoder.layers.7.self_attention.query_key_value.bias": "pytorch_model-00001-of-00002.bin",
189
+ "transformer.encoder.layers.7.self_attention.query_key_value.weight": "pytorch_model-00001-of-00002.bin",
190
+ "transformer.encoder.layers.8.input_layernorm.weight": "pytorch_model-00001-of-00002.bin",
191
+ "transformer.encoder.layers.8.mlp.dense_4h_to_h.weight": "pytorch_model-00001-of-00002.bin",
192
+ "transformer.encoder.layers.8.mlp.dense_h_to_4h.weight": "pytorch_model-00001-of-00002.bin",
193
+ "transformer.encoder.layers.8.post_attention_layernorm.weight": "pytorch_model-00001-of-00002.bin",
194
+ "transformer.encoder.layers.8.self_attention.dense.weight": "pytorch_model-00001-of-00002.bin",
195
+ "transformer.encoder.layers.8.self_attention.query_key_value.bias": "pytorch_model-00001-of-00002.bin",
196
+ "transformer.encoder.layers.8.self_attention.query_key_value.weight": "pytorch_model-00001-of-00002.bin",
197
+ "transformer.encoder.layers.9.input_layernorm.weight": "pytorch_model-00001-of-00002.bin",
198
+ "transformer.encoder.layers.9.mlp.dense_4h_to_h.weight": "pytorch_model-00001-of-00002.bin",
199
+ "transformer.encoder.layers.9.mlp.dense_h_to_4h.weight": "pytorch_model-00001-of-00002.bin",
200
+ "transformer.encoder.layers.9.post_attention_layernorm.weight": "pytorch_model-00001-of-00002.bin",
201
+ "transformer.encoder.layers.9.self_attention.dense.weight": "pytorch_model-00001-of-00002.bin",
202
+ "transformer.encoder.layers.9.self_attention.query_key_value.bias": "pytorch_model-00001-of-00002.bin",
203
+ "transformer.encoder.layers.9.self_attention.query_key_value.weight": "pytorch_model-00001-of-00002.bin",
204
+ "transformer.output_layer.weight": "pytorch_model-00002-of-00002.bin",
205
+ "transformer.rotary_pos_emb.inv_freq": "pytorch_model-00001-of-00002.bin"
206
+ }
207
+ }
quantization.py ADDED
@@ -0,0 +1,188 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from torch.nn import Linear
2
+ from torch.nn.parameter import Parameter
3
+
4
+ import bz2
5
+ import torch
6
+ import base64
7
+ import ctypes
8
+ from transformers.utils import logging
9
+
10
+ from typing import List
11
+ from functools import partial
12
+
13
+ logger = logging.get_logger(__name__)
14
+
15
+ try:
16
+ from cpm_kernels.kernels.base import LazyKernelCModule, KernelFunction, round_up
17
+
18
+ class Kernel:
19
+ def __init__(self, code: bytes, function_names: List[str]):
20
+ self.code = code
21
+ self._function_names = function_names
22
+ self._cmodule = LazyKernelCModule(self.code)
23
+
24
+ for name in self._function_names:
25
+ setattr(self, name, KernelFunction(self._cmodule, name))
26
+
27
+ quantization_code = "$QlpoOTFBWSZTWU9yuJUAQHN//////////f/n/8/n///n//bt4dTidcVx8X3V9FV/92/v4B7/AD5FBQFAAAChSgKpFCFAFVSigUAAAEKhSgUUqgFBKigqVREQAABQBQIANDTTIGI00BkZBkNGE0A0BkBkGQGRkaNAaAGQNBoGgDIAAYIGTI0DQAQAaGmmQMRpoDIyDIaMJoBoDIDIMgMjI0aA0AMgaDQNAGQAAwQMmRoGgAgA0NNMgYjTQGRkGQ0YTQDQGQGQZAZGRo0BoAZA0GgaAMgABggZMjQNABABoaaZAxGmgMjIMhowmgGgMgMgyAyMjRoDQAyBoNA0AZAADBAyZGgaAAmqU1NEgJqnptU/Sn4jRR6J6epk2pqb1Q/SgAPUGgyNNGjQ2SBpoAZAAGg0NB6mgDIAAAAA2oaApSREBNAARhGiYEaEwU8pvImlP0k2aam1GaGqbFNM1MHpTwmkepmyU9R6nqPKekHqNNPUxNGhp6n6p6QaZ6o9TG1GMqcoV9ly6nRanHlq6zPNbnGZNi6HSug+2nPiZ13XcnFYZW+45W11CumhzYhchOJ2GLLV1OBjBjGf4TptOddTSOcVxhqYZMYwZXZZY00zI1paX5X9J+b+f4e+x43RXSxXPOdquiGpduatGyXneN696M9t4HU2eR5XX/kPhP261NTx3JO1Ow7LyuDmeo9a7d351T1ZxnvnrvYnrXv/hXxPCeuYx2XsNmO003eg9J3Z6U7b23meJ4ri01OdzTk9BNO96brz+qT5nuvvH3ds/G+m/JcG/F2XYuhXlvO+jP7U3XgrzPN/lr8Sf1n6j4j7jZs+s/T0tNaNNYzTs12rxjwztHlnire3Nzc3N1wuBwOBwXBvZfoHpD7rFmR99V5vj3aXza3xdBbXMalubTg/jIv5dfAi54Pdc75j4z412n3Npj3Ld/ENm7a3b/Cod6h/ret1/5vn/C+l+gdslMvgPSLJ8d8q+U66fevYn/tW1chleEtNTGlcHCbLRlq0tHzF5tsbbZZfHjjLgZu42XCuC3NrdjTasZGNzgxPIrGqp7r3p7L2p5XjnpPSmTd5XtzqnB6U87zzg1Ol0zd0zsLszxR6lkxp35u6/teL0L0W922cR7Lu1lpL9CsHirzuM2T+BgsyViT6LHcm0/Vr6U/7LGGyJeqTEjt0PHWhF5mCT7R9mtlDwriYv0Tyr/OxYt6qp5r0mPVT0608TqnqMZaarU2nFwrTzzlrs1ed7z1ux60wyr4ydCaTi3enW8x68x0zU7tXSlcmPSW1mGpWJMg4zmPC2lK96tp0OE80y4MfEvnZj8zGluR6b22ki1Ou9V2nCd9xovcPvcYMZYy0lvN60ScZ45vN6yeCeeXFb1lVjnnCar5fwXwE2bzJ4HI1XVPXfXZMm44GUsMpYsmLB65TuVdm0cl0b+i/wGNN66XjeV7zuPpHcnK/juhhjdfId5jMdE5nN0dGmmm2zZs2cexD5n9p/dY352XsvXHaZNWWsmmS1atjR452nYudzvqv2HMRyvNNnlMcDl3R2+yx2uVrBubTW9icHDVtbNXlZm7jma1rM4VurZZd2y6nUau7ZXZ7bVU+mnoOVxZGMrVmvX60605JwmzGZhhhjTWtaaaMaaGTGmNMZasY0iX8VMUl8eepaIrzGSpemWOQyZORk2bNpjUybMmxqYmknCGCFynutfksaZpjTNMaaatM0xsxcGR0sociNqxNSmhhR1ZJPbsn8qyF0t2qH6iYBclclalbtTTcHTDsPaX6rlnElph2Jyumumtynv2Kk8GI7rsvXbIcJgHJOSaSXnnGaI3m87RtVXJOZ/YtgdTE6Wpha6ZlE8ayXkef1fh602r2WwvfMXtMdLlkfnLFdYYwYso+bWqm7yJqHXZGw2nrS5ZanSYnWlxBxMF1V940K2wdrI7R6OYf7DGGamMmTSbRhlS45xmVOumF1EyPCmHrrN8wwZOOrdNtLeMtzFzDlWnfTBxMk2NaXIZHBYxYLD4w8yju0ao65Vz1OIXoS9dLanwCe1PWrYuWMqf1if1z2k2yYfKJ741PDgno1ZQ8DRqvUny3mNoWTzGO6m1DkrJI8JiR5cSd+vZdGOO8nrMoc5+NDUFsMSXaZJeNlMmGLtJsovOsUp7I9S5VojKxF6bTVEelXqlfJobQr3LozSh2Jk7VcrVMfhXqszGWMzNqGhqZY0OadxkyyMssKugZR0KNFXBHlqwmJgTE/BNVMk6ItJXZMR0H47GpXv/DMOvNkmVuaV1PRfEdxuqc7Hcd+ZV/zTLaRxWk0nl9CdCeM6mn5rstHIBcpiuwmUZXeq81DacHI2rmrZ5SuE5mOZd6LQrZg9mx32TprA8BMo5jKN6yLTCi3WzQaZSuhzTtM1fUTGVpG8Tw+KXI0tjEpiWxtLYynOlktSbVlaI5kxP8TDH8kx50xoxi5KcA4pcja8KWLRlO/Ks6q06ergnvm1ca3Tq8Uw7LTUsmWyctXPWmpitl/uvGcWTGXGuAXDfhqazGmjkxcJW5hMMMMpYsXl2TZYtVOddG3XCarUt6Ptq9CZXSNzyuRzqRZOjsxdBbFVz6OA5HI43r1jityVlVpVkxmOsyaYWE1NTGq1sOVh36mHMcxtSvcy70edG0ZGR3I1Go1GRlV7mWWo1G0ZGRqlvH40l7o4m5xMWLLLYyNjnqc8556mdPqLJ31n/1nWOncxzG1tizrHs/Z+d2vP/B/l8wdJ6rHUn2nbbDq4p6htFtYzMMMTaZis1K5GKzGNmxhmUx2DDlZ/qNnIx41xnaMfCZWYaZWtNLTNW8ND4Fw1MyZOCdM428suKG1ehW8TesOydg7J+YYcD4cYR+8dFK6M4E3HM9ZfRNNL+Sn6rsl4DsrDl2HpPCnfxjGXtbZtYys1ttlyJ4T+BvexjGWRjMszK4Jpc77D3GyuVD7q0+G8m9G+2+rGm7cOR2y7FdtY2XUYx/oNlfRYxhMYyYZkyyg55enna9Kt/FFi6GMMwYwdwxWgxGMLKYmUyGExTKMZkMFhkymKuh0NOBNnBu+23LdwDoZYYzGGMxtORaTU1pjTGWTTGGtMrNWUsyyTTLLG1qy2ZjbK2DBllWqxMtBMaYZQmcE7zvvRcTkclUwdkxTaSdyySt/7fpL+T1v516Ji97fwr5JbLu305zMn5+GMTTZ9F+y7ExwmGVfG44yxn3dLv6l5i+Wth1jCrDq21nW9LqvvDzz3Vf3LLH/O/32TJ/erx3bXftO4eF+G956D952K/An4NfvOpjFjExjevP/UmE0fIoZXx6/w6lX/no3D0bLt+ixjieBM6ksRd0yB4Lt2SwYNE+gd1detlZWUnpiZfGfFaK+4PyCa/v18V8X75pe9fLXzp7l3VjF76vWZmHwGz1IZNWT7b8yddJ4q5kyrVdfru6atWc7bVYztL9Jf4GXvT+Y8m9/YsXP6H018a8D4XVOqvfzqeR+6yZOD8dPv0+U7/q5Pl+2dNb0MjzGVH5p6MNQ7cOWvw62U9aHE8DprDek+McLyvDz+te+9Zhq5+YTruufMcWMabqysTmZVWjKPfnK0wyVcrsuhjZRdLkHNvD72b9abriOSGIxiLixMOoalNPXzy+wT/tf+U6HHONfsz+xe8ufHBdQWWGWLA9if0rsnmrxK5LvRZQeWsTCsrmOYy8VteVfuRfcVTtDLItLIsMYxZLdU/DbtSemxF6Z6Zo5WBXE4tFdCyVMMXMTEMZXVlS6Xec2T4e0tHsRcEuWshcJ2YsNF5rUx1E8ifCq6Z+ZP7qdCeu/aTwFd53l16/o0NOw6O3dLavP4Hbi4RdmuDk6DoYaninC0+o4uZjbJ7Rxeu0/FbuFg+q7DVS6fQe0rZ6NDGUNNU6DEqOaLTicKnYZMnBWruljQxoaS3dZhocDge0bSTyOvdAbG5hxe2xji7E/L55xX13wWNDi6HCekcFxfCPGxY0MXC+s7afWaMdDyjyr+o8Rudm/NabOZvdl274zH4f5XK9z6On1Pe/K5TdPAslg77BjuO6Y3eO7GqvOPG/stknp1leyvLL0Z7bl9I4noMvLkzytLhWYzrOZzLXCORe028rORzOg4N/L0HlMOQ3Pgmnbb6KczlabORpu980q37TBqRu0/p3PO6234Bl03Ynuz+9W7gnsEcmvYaYY3aMYY0wx3pYd+ujsXauWdaY5Xkbtl23fPzFHiDB/QMo0yFjBllYxTQYYyxkrwn7JufwJ/PfgJ+C83X69ni6zvXcnyXabv0ncbLwsceS+RNlyN2mnneJtX0ngYO0+e+0+UnA+Wch3ji8hj5an4h+i6XBySU4n+R0roVcbw5yvHrmr4Yw8Y7x6c+9POPYHI5HI5HI5HI5HGXGww4nE4nrVyOR8XeqPEO7PLOiukYa3Novk5hV4cdtYZLI93e+uxff2jRo0aNGjRo0aNG1bVtW1dy3m83m8+tQ5ZzHw3nObwOu8La9Rc1dtkdS8A3eTk823tnktXWlxN6Oixe06zrN70Isd9jiOgZFq9yfkPqP/SLhN2Myl8jDM43bl1nbcb4cO57jlh8Jow6pzXZdL4dyODTuuhu77FyO27DdwdRxmvO+O+3N2+BdqyTwLHVczDVY4UPE4O66/ZO2cx1LFzVdSXtF7G4HMbrauOHRw6c8FdZ5m9fHZHYZXfTlZquyynSyTTKke6vcffSD9pzPA/G7n7jxPmuhc1DHMynPMrGL6AdewYmwu5ko+UUyTwrMv27rPH1v1nGqd87+p6N6LU8k3NEng53xXyHS97+44OSg/sy/hn+Se6yfYNjW0/uTgP+PvWYzLMmjhcLB/gGpri6H83/84eUXWT6T9Hsv7785z/7z4icpW+zfXypuR7rx/gMdZb1/wC678pcs8/2a3mDitGHxl9mfPlll5MafWWqxk/eYuTDgcNMzDGWLWvsuglNxs53GtN6uWpktlW1tZZYcuinMMWmnNnJydze3b2Y1McBxrBkXw799izLMZZYyy0TkbsGM4p03S2uVu5s/XXUdSdec6smVxZYYGpVmT8A+8ajuEyV5FatkvVru2x6uxGXXbH4A+jvgP4GMYy3iPLXzq/6z65+E005ey+cwMZD3fZcqc6xpjTFjQ0P3U+e++cPYmTIwj0nrK5NPTfl3WvpfLtXDcb2HQMudYOxFXQBor4L4T6vrOauFctYXJQ++NUWmJe5bmx1jDiZS1dTqWxo4GR8jm3fttpmPHppk9PEyv4/y8/sO07XacOmcqc0x2Vi9BvNJvN5oW8x4mOsydpidRxMYJPx06m1bqPzq9KtK8sxXNXFodD/+MYYaJTLwOhc9brCsV18oOR1i4tXChyTkq4lf4y1Ke+9axjDHqs1mfBbMXuP4Hzi+X7t8vzv7bHerrUPgPCxhjre4fXdfLNtNM+Jd+Zdh8xd8wP87uNPoPgv4W7/5P2BuxfsMabNnMnza+54Pdi5U671GPZY8CehX8Voeoo7FHpkeEc6715FwHZrIrUrHaviPUbPZHND+IhczrP6FcYvhOZ0Di/ETt0OI+YwNWR9r7tpf6WDeZKZDB1+z2IthOl1mPyb5FluvEx9h9d0NnM0Y1XPFkWIsk1WotJ0PBMmkvjvQTd0e71tfeV+8r8lQ/tpzpsmxJ+InrI/dj2UajUajVTUajatRqNRtGo1Go1Go4wjeMpZFMVV9CHbofPraLsJ3JpWV2XOoanCuFky4y3PPNxucK2uKC1Lbdb1eo+m5XomN6HfeZsabHLHRX/K+offtNGGmHWctcVcG44MdSqsOLY9VzX+Zxfxn2HPdWTpzWvkrtJ8M5zorrKcquRytJ5N5DZmcaW02l76nWO+BqPXm1A2Ry/0q71dH/mqrqeFjkYxjEXtsX8qubTk67rGycyqsdm4tZx5D6D5hhi0waaWmiaMP81Yjii5qxPlPuU/GfTL1Y5E6Jyfiq63qTa39A4J0sOGDgO9WF9bOXl0XfPRbsY2bPNKPy1YrFYrFYmRhhlTIyMjJWJYZHXuCXI8OoXsvfljGLFicNifpp2XunoPiG1wtx3p1Tah+/DD66OnVtVXP9rKbVxOnL0tR/rHtqB5UDErUVcl11D4qqvjpOcxX7armUNJB3LpW6bxVvD08e8h3odKKvyCFZBdSh2FVcST9xV3n3T8t1j7Kr9qgrqXg+13Pt5U7JCvFXVIV1YG5lRhkVYZJYYDDD4KOIMoHCp26WS8GB7uBh2zIdgq/PKyInjV2STShuoapUdCpX1yTwqq/z1VvET7Kh5nVPkO8YyxjLt2MaaMmWTLQvx3qnzltnXW0p2jxgbEtSny/Osv8Y9pLMXYoHVPAhkVdWVeODhR6q9/Sxe2liwwZWMVvFXfRkeIDxAePUPIrdJ4ey6yquzH+PD/bUOWAu05qVHtFd8rrKHSoeNIOUqrYr3FXyToqfYJgwmJdKpXXOwYYegNNGMzfZPp/t3t/DVs4zjNTN61rRqaWaa4NYbRjTa0tWwy2Y2tGN8ZO8ofNKq4j9SL7I+cSm4/6ovLV5HNXLI0jJidwrtk6ynCaP6Z++GjRlWS3tLeW129Mi9evxU9mtz6s5J3Z7M2ngTgnKvmpomxpaLCzPfmx0JWE+m3NLDDGOX47RctdYYNK5jakdqLkRlI39n590T5zctGSwwZZDJj6kW8XSi6ot2MmWWJ0DUT3nuvebBudScjZ79g8cWJ8av0k+/bE5WKd5MdbFpbDVMxu1DVMmtNZGJvq1mtRbn6M+g/kP0FwDwr7quZs7xosNGpbscyxhhd9TyJyFwbLcxlTasg75vW7TsV5K7ji44XPMMrdoj+Y3rT0Hie62nlYV/pwczzOmdLqLhYkzGMzCZWGMQzGMSsZYY6Di1t4nlJ+Em63mJxrVLxPbYxNEdgc1dU2iOKyoYYWjNrEeHTYybVk0atSa7ehuwsWMWTqn1TrnS6hYsi71d1+s+k+ic70e20fzE/VaTdxT9ZtU4GIXdeNx3X77guYYfpHeTQjaMX6brOu4OY4K7Y2d9mbHarI5ox3p4GpJ2Vd/Tst60f7j999pppjR+Q/Qf8J/VaORs3cji7FfFuN61+ui9s8hix1OCh5KGVV23BPXvZfz3CLyHpix+exi8z/KnCnosY2eunor+cxyPO/xJ0vKey9OvE9VjqaYu0x3Z3jd6o2b1T12D+F8l232lwaaacD5LE8LBxu7WTlbWraWpew8Xexjel3E+wWD4APITdNqR8F3R3T0lunCQ4GaE9R37DxeCYfcHi4xci5ovKfxVs55y2hf+65E/Xdp6jR5nrebTmi5incpkyOjs50JvrZwstbbW6kfuuQw+2mykf/EXNFzxfKTrxew929TR6bWnGL//F3JFOFCQT3K4lQ"
28
+
29
+ kernels = Kernel(
30
+ bz2.decompress(base64.b64decode(quantization_code)),
31
+ [
32
+ "int4WeightCompression",
33
+ "int4WeightExtractionFloat",
34
+ "int4WeightExtractionHalf",
35
+ "int8WeightExtractionFloat",
36
+ "int8WeightExtractionHalf",
37
+ ],
38
+ )
39
+ except Exception as exception:
40
+ kernels = None
41
+ logger.warning("Failed to load cpm_kernels:" + str(exception))
42
+
43
+
44
+ class W8A16Linear(torch.autograd.Function):
45
+ @staticmethod
46
+ def forward(ctx, inp: torch.Tensor, quant_w: torch.Tensor, scale_w: torch.Tensor, weight_bit_width):
47
+ ctx.inp_shape = inp.size()
48
+ ctx.weight_bit_width = weight_bit_width
49
+ out_features = quant_w.size(0)
50
+ inp = inp.contiguous().view(-1, inp.size(-1))
51
+ weight = extract_weight_to_half(quant_w, scale_w, weight_bit_width)
52
+ ctx.weight_shape = weight.size()
53
+ output = inp.mm(weight.t())
54
+ ctx.save_for_backward(inp, quant_w, scale_w)
55
+ return output.view(*(ctx.inp_shape[:-1] + (out_features,)))
56
+
57
+ @staticmethod
58
+ def backward(ctx, grad_output: torch.Tensor):
59
+ inp, quant_w, scale_w = ctx.saved_tensors
60
+ weight = extract_weight_to_half(quant_w, scale_w, ctx.weight_bit_width)
61
+ grad_output = grad_output.contiguous().view(-1, weight.size(0))
62
+ grad_input = grad_output.mm(weight)
63
+ grad_weight = grad_output.t().mm(inp)
64
+ return grad_input.view(ctx.inp_shape), grad_weight.view(ctx.weight_shape), None, None
65
+
66
+
67
+ def compress_int4_weight(weight: torch.Tensor): # (n, m)
68
+ with torch.cuda.device(weight.device):
69
+ n, m = weight.size(0), weight.size(1)
70
+ assert m % 2 == 0
71
+ m = m // 2
72
+ out = torch.empty(n, m, dtype=torch.int8, device="cuda")
73
+ stream = torch.cuda.current_stream()
74
+
75
+ gridDim = (n, 1, 1)
76
+ blockDim = (min(round_up(m, 32), 1024), 1, 1)
77
+
78
+ kernels.int4WeightCompression(
79
+ gridDim,
80
+ blockDim,
81
+ 0,
82
+ stream,
83
+ [ctypes.c_void_p(weight.data_ptr()), ctypes.c_void_p(out.data_ptr()), ctypes.c_int32(n), ctypes.c_int32(m)],
84
+ )
85
+ return out
86
+
87
+
88
+ def extract_weight_to_half(weight: torch.Tensor, scale_list: torch.Tensor, source_bit_width: int):
89
+ assert scale_list.dtype in [torch.half, torch.bfloat16]
90
+ assert weight.dtype in [torch.int8]
91
+ if source_bit_width == 8:
92
+ return weight.to(scale_list.dtype) * scale_list[:, None]
93
+ elif source_bit_width == 4:
94
+ func = (
95
+ kernels.int4WeightExtractionHalf if scale_list.dtype == torch.half else kernels.int4WeightExtractionBFloat16
96
+ )
97
+ else:
98
+ assert False, "Unsupported bit-width"
99
+
100
+ with torch.cuda.device(weight.device):
101
+ n, m = weight.size(0), weight.size(1)
102
+ out = torch.empty(n, m * (8 // source_bit_width), dtype=scale_list.dtype, device="cuda")
103
+ stream = torch.cuda.current_stream()
104
+
105
+ gridDim = (n, 1, 1)
106
+ blockDim = (min(round_up(m, 32), 1024), 1, 1)
107
+
108
+ func(
109
+ gridDim,
110
+ blockDim,
111
+ 0,
112
+ stream,
113
+ [
114
+ ctypes.c_void_p(weight.data_ptr()),
115
+ ctypes.c_void_p(scale_list.data_ptr()),
116
+ ctypes.c_void_p(out.data_ptr()),
117
+ ctypes.c_int32(n),
118
+ ctypes.c_int32(m),
119
+ ],
120
+ )
121
+ return out
122
+
123
+
124
+ class QuantizedLinear(torch.nn.Module):
125
+ def __init__(self, weight_bit_width: int, weight, bias=None, device="cpu", dtype=None, empty_init=False, *args,
126
+ **kwargs):
127
+ super().__init__()
128
+ self.weight_bit_width = weight_bit_width
129
+
130
+ shape = weight.shape
131
+
132
+ if weight is None or empty_init:
133
+ self.weight = torch.empty(shape[0], shape[1] * weight_bit_width // 8, dtype=torch.int8, device=device)
134
+ self.weight_scale = torch.empty(shape[0], dtype=dtype, device=device)
135
+ else:
136
+ self.weight_scale = weight.abs().max(dim=-1).values / ((2 ** (weight_bit_width - 1)) - 1)
137
+ self.weight = torch.round(weight / self.weight_scale[:, None]).to(torch.int8)
138
+ if weight_bit_width == 4:
139
+ self.weight = compress_int4_weight(self.weight)
140
+
141
+ self.weight = Parameter(self.weight.to(device), requires_grad=False)
142
+ self.weight_scale = Parameter(self.weight_scale.to(device), requires_grad=False)
143
+ self.bias = Parameter(bias.to(device), requires_grad=False) if bias is not None else None
144
+
145
+ def forward(self, input):
146
+ output = W8A16Linear.apply(input, self.weight, self.weight_scale, self.weight_bit_width)
147
+ if self.bias is not None:
148
+ output = output + self.bias
149
+ return output
150
+
151
+
152
+ def quantize(model, weight_bit_width, empty_init=False, device=None):
153
+ """Replace fp16 linear with quantized linear"""
154
+ for layer in model.layers:
155
+ layer.self_attention.query_key_value = QuantizedLinear(
156
+ weight_bit_width=weight_bit_width,
157
+ weight=layer.self_attention.query_key_value.weight.to(torch.cuda.current_device()),
158
+ bias=layer.self_attention.query_key_value.bias,
159
+ dtype=layer.self_attention.query_key_value.weight.dtype,
160
+ device=layer.self_attention.query_key_value.weight.device if device is None else device,
161
+ empty_init=empty_init
162
+ )
163
+ layer.self_attention.dense = QuantizedLinear(
164
+ weight_bit_width=weight_bit_width,
165
+ weight=layer.self_attention.dense.weight.to(torch.cuda.current_device()),
166
+ bias=layer.self_attention.dense.bias,
167
+ dtype=layer.self_attention.dense.weight.dtype,
168
+ device=layer.self_attention.dense.weight.device if device is None else device,
169
+ empty_init=empty_init
170
+ )
171
+ layer.mlp.dense_h_to_4h = QuantizedLinear(
172
+ weight_bit_width=weight_bit_width,
173
+ weight=layer.mlp.dense_h_to_4h.weight.to(torch.cuda.current_device()),
174
+ bias=layer.mlp.dense_h_to_4h.bias,
175
+ dtype=layer.mlp.dense_h_to_4h.weight.dtype,
176
+ device=layer.mlp.dense_h_to_4h.weight.device if device is None else device,
177
+ empty_init=empty_init
178
+ )
179
+ layer.mlp.dense_4h_to_h = QuantizedLinear(
180
+ weight_bit_width=weight_bit_width,
181
+ weight=layer.mlp.dense_4h_to_h.weight.to(torch.cuda.current_device()),
182
+ bias=layer.mlp.dense_4h_to_h.bias,
183
+ dtype=layer.mlp.dense_4h_to_h.weight.dtype,
184
+ device=layer.mlp.dense_4h_to_h.weight.device if device is None else device,
185
+ empty_init=empty_init
186
+ )
187
+
188
+ return model
special_tokens_map.json ADDED
@@ -0,0 +1 @@
 
 
1
+ {}
tokenization_chatglm.py ADDED
@@ -0,0 +1,253 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ import torch
3
+ from typing import List, Optional, Union, Dict
4
+ from sentencepiece import SentencePieceProcessor
5
+ from transformers import PreTrainedTokenizer
6
+ from transformers.utils import logging, PaddingStrategy
7
+ from transformers.tokenization_utils_base import EncodedInput, BatchEncoding
8
+
9
+
10
+ class SPTokenizer:
11
+ def __init__(self, model_path: str):
12
+ # reload tokenizer
13
+ assert os.path.isfile(model_path), model_path
14
+ self.sp_model = SentencePieceProcessor(model_file=model_path)
15
+
16
+ # BOS / EOS token IDs
17
+ self.n_words: int = self.sp_model.vocab_size()
18
+ self.bos_id: int = self.sp_model.bos_id()
19
+ self.eos_id: int = self.sp_model.eos_id()
20
+ self.pad_id: int = self.sp_model.unk_id()
21
+ assert self.sp_model.vocab_size() == self.sp_model.get_piece_size()
22
+
23
+ special_tokens = ["[MASK]", "[gMASK]", "[sMASK]", "sop", "eop"]
24
+ self.special_tokens = {}
25
+ self.index_special_tokens = {}
26
+ for token in special_tokens:
27
+ self.special_tokens[token] = self.n_words
28
+ self.index_special_tokens[self.n_words] = token
29
+ self.n_words += 1
30
+
31
+ def tokenize(self, s: str):
32
+ return self.sp_model.EncodeAsPieces(s)
33
+
34
+ def encode(self, s: str, bos: bool = False, eos: bool = False) -> List[int]:
35
+ assert type(s) is str
36
+ t = self.sp_model.encode(s)
37
+ if bos:
38
+ t = [self.bos_id] + t
39
+ if eos:
40
+ t = t + [self.eos_id]
41
+ return t
42
+
43
+ def decode(self, t: List[int]) -> str:
44
+ return self.sp_model.decode(t)
45
+
46
+ def decode_tokens(self, tokens: List[str]) -> str:
47
+ text = self.sp_model.DecodePieces(tokens)
48
+ return text
49
+
50
+ def convert_token_to_id(self, token):
51
+ """ Converts a token (str) in an id using the vocab. """
52
+ if token in self.special_tokens:
53
+ return self.special_tokens[token]
54
+ return self.sp_model.PieceToId(token)
55
+
56
+ def convert_id_to_token(self, index):
57
+ """Converts an index (integer) in a token (str) using the vocab."""
58
+ if index in self.index_special_tokens or index in [self.eos_id, self.bos_id, self.pad_id] or index < 0:
59
+ return ""
60
+ return self.sp_model.IdToPiece(index)
61
+
62
+
63
+ class ChatGLMTokenizer(PreTrainedTokenizer):
64
+ vocab_files_names = {"vocab_file": "tokenizer.model"}
65
+
66
+ model_input_names = ["input_ids", "attention_mask", "position_ids"]
67
+
68
+ def __init__(self, vocab_file, padding_side="left", clean_up_tokenization_spaces=False, **kwargs):
69
+ self.name = "GLMTokenizer"
70
+
71
+ self.vocab_file = vocab_file
72
+ self.tokenizer = SPTokenizer(vocab_file)
73
+ self.special_tokens = {
74
+ "<bos>": self.tokenizer.bos_id,
75
+ "<eos>": self.tokenizer.eos_id,
76
+ "<pad>": self.tokenizer.pad_id
77
+ }
78
+ super().__init__(padding_side=padding_side, clean_up_tokenization_spaces=clean_up_tokenization_spaces, **kwargs)
79
+
80
+ def get_command(self, token):
81
+ if token in self.special_tokens:
82
+ return self.special_tokens[token]
83
+ assert token in self.tokenizer.special_tokens, f"{token} is not a special token for {self.name}"
84
+ return self.tokenizer.special_tokens[token]
85
+
86
+ @property
87
+ def pad_token(self) -> str:
88
+ return "<unk>"
89
+
90
+ @property
91
+ def pad_token_id(self):
92
+ return self.get_command("<pad>")
93
+
94
+ @property
95
+ def eos_token(self) -> str:
96
+ return "</s>"
97
+
98
+ @property
99
+ def eos_token_id(self):
100
+ return self.get_command("<eos>")
101
+
102
+ @property
103
+ def vocab_size(self):
104
+ return self.tokenizer.n_words
105
+
106
+ def get_vocab(self):
107
+ """ Returns vocab as a dict """
108
+ vocab = {self._convert_id_to_token(i): i for i in range(self.vocab_size)}
109
+ vocab.update(self.added_tokens_encoder)
110
+ return vocab
111
+
112
+ def _tokenize(self, text, **kwargs):
113
+ return self.tokenizer.tokenize(text)
114
+
115
+ def _convert_token_to_id(self, token):
116
+ """ Converts a token (str) in an id using the vocab. """
117
+ return self.tokenizer.convert_token_to_id(token)
118
+
119
+ def _convert_id_to_token(self, index):
120
+ """Converts an index (integer) in a token (str) using the vocab."""
121
+ return self.tokenizer.convert_id_to_token(index)
122
+
123
+ def convert_tokens_to_string(self, tokens: List[str]) -> str:
124
+ return self.tokenizer.decode_tokens(tokens)
125
+
126
+ def save_vocabulary(self, save_directory, filename_prefix=None):
127
+ """
128
+ Save the vocabulary and special tokens file to a directory.
129
+
130
+ Args:
131
+ save_directory (`str`):
132
+ The directory in which to save the vocabulary.
133
+ filename_prefix (`str`, *optional*):
134
+ An optional prefix to add to the named of the saved files.
135
+
136
+ Returns:
137
+ `Tuple(str)`: Paths to the files saved.
138
+ """
139
+ if os.path.isdir(save_directory):
140
+ vocab_file = os.path.join(
141
+ save_directory, self.vocab_files_names["vocab_file"]
142
+ )
143
+ else:
144
+ vocab_file = save_directory
145
+
146
+ with open(self.vocab_file, 'rb') as fin:
147
+ proto_str = fin.read()
148
+
149
+ with open(vocab_file, "wb") as writer:
150
+ writer.write(proto_str)
151
+
152
+ return (vocab_file,)
153
+
154
+ def get_prefix_tokens(self):
155
+ prefix_tokens = [self.get_command("[gMASK]"), self.get_command("sop")]
156
+ return prefix_tokens
157
+
158
+ def build_prompt(self, query, history=None):
159
+ if history is None:
160
+ history = []
161
+ prompt = ""
162
+ for i, (old_query, response) in enumerate(history):
163
+ prompt += "[Round {}]\n\n问:{}\n\n答:{}\n\n".format(i + 1, old_query, response)
164
+ prompt += "[Round {}]\n\n问:{}\n\n答:".format(len(history) + 1, query)
165
+ return prompt
166
+
167
+ def build_inputs_with_special_tokens(
168
+ self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None
169
+ ) -> List[int]:
170
+ """
171
+ Build model inputs from a sequence or a pair of sequence for sequence classification tasks by concatenating and
172
+ adding special tokens. A BERT sequence has the following format:
173
+
174
+ - single sequence: `[CLS] X [SEP]`
175
+ - pair of sequences: `[CLS] A [SEP] B [SEP]`
176
+
177
+ Args:
178
+ token_ids_0 (`List[int]`):
179
+ List of IDs to which the special tokens will be added.
180
+ token_ids_1 (`List[int]`, *optional*):
181
+ Optional second list of IDs for sequence pairs.
182
+
183
+ Returns:
184
+ `List[int]`: List of [input IDs](../glossary#input-ids) with the appropriate special tokens.
185
+ """
186
+ prefix_tokens = self.get_prefix_tokens()
187
+ token_ids_0 = prefix_tokens + token_ids_0
188
+ if token_ids_1 is not None:
189
+ token_ids_0 = token_ids_0 + token_ids_1 + [self.get_command("<eos>")]
190
+ return token_ids_0
191
+
192
+ def _pad(
193
+ self,
194
+ encoded_inputs: Union[Dict[str, EncodedInput], BatchEncoding],
195
+ max_length: Optional[int] = None,
196
+ padding_strategy: PaddingStrategy = PaddingStrategy.DO_NOT_PAD,
197
+ pad_to_multiple_of: Optional[int] = None,
198
+ return_attention_mask: Optional[bool] = None,
199
+ ) -> dict:
200
+ """
201
+ Pad encoded inputs (on left/right and up to predefined length or max length in the batch)
202
+
203
+ Args:
204
+ encoded_inputs:
205
+ Dictionary of tokenized inputs (`List[int]`) or batch of tokenized inputs (`List[List[int]]`).
206
+ max_length: maximum length of the returned list and optionally padding length (see below).
207
+ Will truncate by taking into account the special tokens.
208
+ padding_strategy: PaddingStrategy to use for padding.
209
+
210
+ - PaddingStrategy.LONGEST Pad to the longest sequence in the batch
211
+ - PaddingStrategy.MAX_LENGTH: Pad to the max length (default)
212
+ - PaddingStrategy.DO_NOT_PAD: Do not pad
213
+ The tokenizer padding sides are defined in self.padding_side:
214
+
215
+ - 'left': pads on the left of the sequences
216
+ - 'right': pads on the right of the sequences
217
+ pad_to_multiple_of: (optional) Integer if set will pad the sequence to a multiple of the provided value.
218
+ This is especially useful to enable the use of Tensor Core on NVIDIA hardware with compute capability
219
+ `>= 7.5` (Volta).
220
+ return_attention_mask:
221
+ (optional) Set to False to avoid returning attention mask (default: set to model specifics)
222
+ """
223
+ # Load from model defaults
224
+ assert self.padding_side == "left"
225
+
226
+ required_input = encoded_inputs[self.model_input_names[0]]
227
+ seq_length = len(required_input)
228
+
229
+ if padding_strategy == PaddingStrategy.LONGEST:
230
+ max_length = len(required_input)
231
+
232
+ if max_length is not None and pad_to_multiple_of is not None and (max_length % pad_to_multiple_of != 0):
233
+ max_length = ((max_length // pad_to_multiple_of) + 1) * pad_to_multiple_of
234
+
235
+ needs_to_be_padded = padding_strategy != PaddingStrategy.DO_NOT_PAD and len(required_input) != max_length
236
+
237
+ # Initialize attention mask if not present.
238
+ if "attention_mask" not in encoded_inputs:
239
+ encoded_inputs["attention_mask"] = [1] * seq_length
240
+
241
+ if "position_ids" not in encoded_inputs:
242
+ encoded_inputs["position_ids"] = list(range(seq_length))
243
+
244
+ if needs_to_be_padded:
245
+ difference = max_length - len(required_input)
246
+
247
+ if "attention_mask" in encoded_inputs:
248
+ encoded_inputs["attention_mask"] = [0] * difference + encoded_inputs["attention_mask"]
249
+ if "position_ids" in encoded_inputs:
250
+ encoded_inputs["position_ids"] = [0] * difference + encoded_inputs["position_ids"]
251
+ encoded_inputs[self.model_input_names[0]] = [self.pad_token_id] * difference + required_input
252
+
253
+ return encoded_inputs
tokenizer.model ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:e7dc4c393423b76e4373e5157ddc34803a0189ba96b21ddbb40269d31468a6f2
3
+ size 1018370
tokenizer_config.json ADDED
@@ -0,0 +1,14 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "auto_map": {
3
+ "AutoTokenizer": [
4
+ "tokenization_chatglm.ChatGLMTokenizer",
5
+ null
6
+ ]
7
+ },
8
+ "clean_up_tokenization_spaces": true,
9
+ "do_lower_case": false,
10
+ "model_max_length": 1000000000000000019884624838656,
11
+ "padding_side": "left",
12
+ "remove_space": false,
13
+ "tokenizer_class": "ChatGLMTokenizer"
14
+ }