zhaicunqi commited on
Commit
96e6701
1 Parent(s): 9457882

Upload folder using huggingface_hub

Browse files
LICENSE.txt ADDED
@@ -0,0 +1,51 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ Apache License
2
+ Version 2.0, January 2004
3
+ http://www.apache.org/licenses/
4
+
5
+ TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
6
+
7
+ 1. Definitions.
8
+
9
+ "License" shall mean the terms and conditions for use, reproduction, and distribution as defined by Sections 1 through 9 of this document.
10
+
11
+ "Licensor" shall mean the copyright owner or entity authorized by the copyright owner that is granting the License.
12
+
13
+ "Legal Entity" shall mean the union of the acting entity and all other entities that control, are controlled by, or are under common control with that entity. For the purposes of this definition, "control" means (i) the power, direct or indirect, to cause the direction or management of such entity, whether by contract or otherwise, or (ii) ownership of fifty percent (50%) or more of the outstanding shares, or (iii) beneficial ownership of such entity.
14
+
15
+ "You" (or "Your") shall mean an individual or Legal Entity exercising permissions granted by this License.
16
+
17
+ "Source" form shall mean the preferred form for making modifications, including but not limited to software source code, documentation source, and configuration files.
18
+
19
+ "Object" form shall mean any form resulting from mechanical transformation or translation of a Source form, including but not limited to compiled object code, generated documentation, and conversions to other media types.
20
+
21
+ "Work" shall mean the work of authorship, whether in Source or Object form, made available under the License, as indicated by a copyright notice that is included in or attached to the work (an example is provided in the Appendix below).
22
+
23
+ "Derivative Works" shall mean any work, whether in Source or Object form, that is based on (or derived from) the Work and for which the editorial revisions, annotations, elaborations, or other modifications represent, as a whole, an original work of authorship. For the purposes of this License, Derivative Works shall not include works that remain separable from, or merely link (or bind by name) to the interfaces of, the Work and Derivative Works thereof.
24
+
25
+ "Contribution" shall mean any work of authorship, including the original version of the Work and any modifications or additions to that Work or Derivative Works thereof, that is intentionally submitted to Licensor for inclusion in the Work by the copyright owner or by an individual or Legal Entity authorized to submit on behalf of the copyright owner. For the purposes of this definition, "submitted" means any form of electronic, verbal, or written communication sent to the Licensor or its representatives, including but not limited to communication on electronic mailing lists, source code control systems, and issue tracking systems that are managed by, or on behalf of, the Licensor for the purpose of discussing and improving the Work, but excluding communication that is conspicuously marked or otherwise designated in writing by the copyright owner as "Not a Contribution."
26
+
27
+ "Contributor" shall mean Licensor and any individual or Legal Entity on behalf of whom a Contribution has been received by Licensor and subsequently incorporated within the Work.
28
+
29
+ 2. Grant of Copyright License. Subject to the terms and conditions of this License, each Contributor hereby grants to You a perpetual, worldwide, non-exclusive, no-charge, royalty-free, irrevocable copyright license to reproduce, prepare Derivative Works of, publicly display, publicly perform, sublicense, and distribute the Work and such Derivative Works in Source or Object form.
30
+
31
+ 3. Grant of Patent License. Subject to the terms and conditions of this License, each Contributor hereby grants to You a perpetual, worldwide, non-exclusive, no-charge, royalty-free, irrevocable (except as stated in this section) patent license to make, have made, use, offer to sell, sell, import, and otherwise transfer the Work, where such license applies only to those patent claims licensable by such Contributor that are necessarily infringed by their Contribution(s) alone or by combination of their Contribution(s) with the Work to which such Contribution(s) was submitted. If You institute patent litigation against any entity (including a cross-claim or counterclaim in a lawsuit) alleging that the Work or a Contribution incorporated within the Work constitutes direct or contributory patent infringement, then any patent licenses granted to You under this License for that Work shall terminate as of the date such litigation is filed.
32
+
33
+ 4. Redistribution. You may reproduce and distribute copies of the Work or Derivative Works thereof in any medium, with or without modifications, and in Source or Object form, provided that You meet the following conditions:
34
+
35
+ You must give any other recipients of the Work or Derivative Works a copy of this License; and
36
+ You must cause any modified files to carry prominent notices stating that You changed the files; and
37
+ You must retain, in the Source form of any Derivative Works that You distribute, all copyright, patent, trademark, and attribution notices from the Source form of the Work, excluding those notices that do not pertain to any part of the Derivative Works; and
38
+ If the Work includes a "NOTICE" text file as part of its distribution, then any Derivative Works that You distribute must include a readable copy of the attribution notices contained within such NOTICE file, excluding those notices that do not pertain to any part of the Derivative Works, in at least one of the following places: within a NOTICE text file distributed as part of the Derivative Works; within the Source form or documentation, if provided along with the Derivative Works; or, within a display generated by the Derivative Works, if and wherever such third-party notices normally appear. The contents of the NOTICE file are for informational purposes only and do not modify the License. You may add Your own attribution notices within Derivative Works that You distribute, alongside or as an addendum to the NOTICE text from the Work, provided that such additional attribution notices cannot be construed as modifying the License.
39
+ You may add Your own copyright statement to Your modifications and may provide additional or different license terms and conditions for use, reproduction, or distribution of Your modifications, or for any such Derivative Works as a whole, provided Your use, reproduction, and distribution of the Work otherwise complies with the conditions stated in this License.
40
+
41
+ 5. Submission of Contributions. Unless You explicitly state otherwise, any Contribution intentionally submitted for inclusion in the Work by You to the Licensor shall be under the terms and conditions of this License, without any additional terms or conditions. Notwithstanding the above, nothing herein shall supersede or modify the terms of any separate license agreement you may have executed with Licensor regarding such Contributions.
42
+
43
+ 6. Trademarks. This License does not grant permission to use the trade names, trademarks, service marks, or product names of the Licensor, except as required for reasonable and customary use in describing the origin of the Work and reproducing the content of the NOTICE file.
44
+
45
+ 7. Disclaimer of Warranty. Unless required by applicable law or agreed to in writing, Licensor provides the Work (and each Contributor provides its Contributions) on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied, including, without limitation, any warranties or conditions of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A PARTICULAR PURPOSE. You are solely responsible for determining the appropriateness of using or redistributing the Work and assume any risks associated with Your exercise of permissions under this License.
46
+
47
+ 8. Limitation of Liability. In no event and under no legal theory, whether in tort (including negligence), contract, or otherwise, unless required by applicable law (such as deliberate and grossly negligent acts) or agreed to in writing, shall any Contributor be liable to You for damages, including any direct, indirect, special, incidental, or consequential damages of any character arising as a result of this License or out of the use or inability to use the Work (including but not limited to damages for loss of goodwill, work stoppage, computer failure or malfunction, or any and all other commercial damages or losses), even if such Contributor has been advised of the possibility of such damages.
48
+
49
+ 9. Accepting Warranty or Additional Liability. While redistributing the Work or Derivative Works thereof, You may choose to offer, and charge a fee for, acceptance of support, warranty, indemnity, or other liability obligations and/or rights consistent with this License. However, in accepting such obligations, You may act only on Your own behalf and on Your sole responsibility, not on behalf of any other Contributor, and only if You agree to indemnify, defend, and hold each Contributor harmless for any liability incurred by, or claims asserted against, such Contributor by reason of your accepting any such warranty or additional liability.
50
+
51
+ END OF TERMS AND CONDITIONS
config.json ADDED
@@ -0,0 +1,30 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "architectures": [
3
+ "ZhinaoForCausalLM"
4
+ ],
5
+ "auto_map": {
6
+ "AutoConfig": "configuration_zhinao.ZhinaoConfig",
7
+ "AutoModelForCausalLM": "modeling_zhinao.ZhinaoForCausalLM"
8
+ },
9
+ "bf16": true,
10
+ "fp16": false,
11
+ "hidden_act": "silu",
12
+ "hidden_size": 4096,
13
+ "initializer_range": 0.01,
14
+ "intermediate_size": 11008,
15
+ "max_position_embeddings": 32768,
16
+ "model_max_length": 32768,
17
+ "model_type": "zhinao",
18
+ "num_attention_heads": 32,
19
+ "num_hidden_layers": 32,
20
+ "num_key_value_heads": 32,
21
+ "rms_norm_eps": 1e-05,
22
+ "rope_scaling": null,
23
+ "rope_theta": 50000000.0,
24
+ "tie_word_embeddings": false,
25
+ "torch_dtype": "bfloat16",
26
+ "transformers_version": "4.38.2",
27
+ "use_cache": false,
28
+ "use_flash_attn": "auto",
29
+ "vocab_size": 158464
30
+ }
configuration_zhinao.py ADDED
@@ -0,0 +1,92 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright (c) 360zhinao and the HuggingFace Inc. team. All rights reserved.
2
+ # This code is built upon Huggingface's transformers repository.
3
+
4
+
5
+ from transformers.configuration_utils import PretrainedConfig
6
+ from transformers.utils import logging
7
+
8
+
9
+ logger = logging.get_logger(__name__)
10
+
11
+
12
+ class ZhinaoConfig(PretrainedConfig):
13
+
14
+ model_type = "zhinao"
15
+ keys_to_ignore_at_inference = ["past_key_values"]
16
+
17
+ def __init__(
18
+ self,
19
+ vocab_size=32000,
20
+ hidden_size=4096,
21
+ intermediate_size=11008,
22
+ num_hidden_layers=32,
23
+ num_attention_heads=32,
24
+ num_key_value_heads=None,
25
+ hidden_act="silu",
26
+ max_position_embeddings=2048,
27
+ initializer_range=0.02,
28
+ rms_norm_eps=1e-6,
29
+ use_cache=True,
30
+ pad_token_id=None,
31
+ bos_token_id=None,
32
+ eos_token_id=None,
33
+ tie_word_embeddings=False,
34
+ rope_theta=10000.0,
35
+ rope_scaling=None,
36
+ bf16 = False,
37
+ fp16 = False,
38
+ use_flash_attn="auto",
39
+ **kwargs,
40
+ ):
41
+ self.vocab_size = vocab_size
42
+ self.max_position_embeddings = max_position_embeddings
43
+ self.hidden_size = hidden_size
44
+ self.intermediate_size = intermediate_size
45
+ self.num_hidden_layers = num_hidden_layers
46
+ self.num_attention_heads = num_attention_heads
47
+
48
+ # for backward compatibility
49
+ if num_key_value_heads is None:
50
+ num_key_value_heads = num_attention_heads
51
+
52
+ self.num_key_value_heads = num_key_value_heads
53
+ self.hidden_act = hidden_act
54
+ self.initializer_range = initializer_range
55
+ self.rms_norm_eps = rms_norm_eps
56
+ self.use_cache = use_cache
57
+ self.rope_theta = rope_theta
58
+ self.rope_scaling = rope_scaling
59
+ self._rope_scaling_validation()
60
+
61
+ self.bf16 = bf16
62
+ self.fp16 = fp16
63
+ self.use_flash_attn = use_flash_attn
64
+
65
+ super().__init__(
66
+ pad_token_id=pad_token_id,
67
+ bos_token_id=bos_token_id,
68
+ eos_token_id=eos_token_id,
69
+ tie_word_embeddings=tie_word_embeddings,
70
+ **kwargs,
71
+ )
72
+
73
+ def _rope_scaling_validation(self):
74
+ """
75
+ Validate the `rope_scaling` configuration.
76
+ """
77
+ if self.rope_scaling is None:
78
+ return
79
+
80
+ if not isinstance(self.rope_scaling, dict) or len(self.rope_scaling) != 2:
81
+ raise ValueError(
82
+ "`rope_scaling` must be a dictionary with with two fields, `type` and `factor`, "
83
+ f"got {self.rope_scaling}"
84
+ )
85
+ rope_scaling_type = self.rope_scaling.get("type", None)
86
+ rope_scaling_factor = self.rope_scaling.get("factor", None)
87
+ if rope_scaling_type is None or rope_scaling_type not in ["linear", "dynamic", "ntk"]:
88
+ raise ValueError(
89
+ f"`rope_scaling`'s type field must be one of ['linear', 'dynamic'], got {rope_scaling_type}"
90
+ )
91
+ if rope_scaling_factor is None or not isinstance(rope_scaling_factor, float) or rope_scaling_factor <= 1.0:
92
+ raise ValueError(f"`rope_scaling`'s factor field must be an float > 1, got {rope_scaling_factor}")
generation_config.json ADDED
@@ -0,0 +1,14 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "_from_model_config": true,
3
+ "do_sample": true,
4
+ "eos_token_id": [
5
+ 158326,
6
+ 158332,
7
+ 158333
8
+ ],
9
+ "max_new_tokens": 1024,
10
+ "pad_token_id": 158326,
11
+ "top_k": 0,
12
+ "top_p": 0.8,
13
+ "transformers_version": "4.38.2"
14
+ }
generation_utils.py ADDED
@@ -0,0 +1,250 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import re
2
+ import torch
3
+ import numpy as np
4
+ from queue import Queue
5
+ from typing import Tuple, List, Union, Iterable
6
+ from transformers.utils import logging, add_start_docstrings
7
+ from transformers.generation.logits_process import LogitsProcessor, LOGITS_PROCESSOR_INPUTS_DOCSTRING, LogitsProcessorList
8
+
9
+
10
+ def make_context(model, tokenizer,
11
+ messages: List[dict],
12
+ system: str = "You are a helpful assistant.",
13
+ max_new_tokens: int=0,
14
+ ):
15
+
16
+ max_new_tokens = max_new_tokens or model.generation_config.max_new_tokens
17
+ max_input_length = model.config.model_max_length - max_new_tokens
18
+
19
+ im_start_id = [tokenizer.im_start_id]
20
+ im_end_id = [tokenizer.im_end_id]
21
+ nl_tokens = tokenizer.encode("\n")
22
+
23
+ def _tokenize_str(role, content):
24
+ return tokenizer.encode(role, allowed_special=set()) + nl_tokens + tokenizer.encode(content, allowed_special=set())
25
+
26
+ def _parse_messages(messages):
27
+ system, query, history = "", "", []
28
+ ## system
29
+ if messages[0]["role"] == "system":
30
+ system = messages[0]["content"]
31
+ messages = messages[1:]
32
+ ## query
33
+ assert messages[-1]["role"] == "user"
34
+ query = messages[-1]["content"]
35
+ messages = messages[:-1]
36
+ ## history
37
+ assert len(messages) % 2 == 0
38
+ for i in range(0, len(messages), 2):
39
+ assert messages[i]["role"] == "user" and messages[i+1]["role"] == "assistant"
40
+ history.append([messages[i]["content"], messages[i+1]["content"]])
41
+
42
+ return system, query, history
43
+
44
+ _system, query, history = _parse_messages(messages)
45
+
46
+ ## system
47
+ system_text = _system if _system != "" else system
48
+ system_tokens = []
49
+ if system_text:
50
+ system_tokens = im_start_id + _tokenize_str("system", system_text) + im_end_id + nl_tokens
51
+
52
+ ## query
53
+ query_tokens = im_start_id + _tokenize_str("user", query) + im_end_id + nl_tokens
54
+ ## final assistant
55
+ final_tokens = im_start_id + tokenizer.encode("assistant", allowed_special=set()) + nl_tokens
56
+
57
+ ## max_history_tokens
58
+ max_history_length = max_input_length - len(system_tokens) - len(query_tokens) - len(final_tokens)
59
+
60
+ ## history
61
+ context_tokens = []
62
+ for turn_query, turn_response in reversed(history):
63
+ ## query tokens
64
+ history_query_tokens = im_start_id + _tokenize_str("user", turn_query) + im_end_id + nl_tokens
65
+ ## answer tokens
66
+ histroy_response_tokens = im_start_id + _tokenize_str("assistant", turn_response) + im_end_id + nl_tokens
67
+ ## this round tokens
68
+ next_context_tokens = history_query_tokens + histroy_response_tokens
69
+ ## concat
70
+ current_context_size = len(next_context_tokens) + len(context_tokens)
71
+ if current_context_size < max_history_length:
72
+ context_tokens = next_context_tokens + context_tokens
73
+ else:
74
+ break
75
+ input_tokens = system_tokens + context_tokens + query_tokens + final_tokens
76
+
77
+ return torch.LongTensor([input_tokens]).to(model.device)
78
+
79
+
80
+ def parse_pot_no_stream(inputs):
81
+ try:
82
+ s = re.findall(r'<<(.*?)>>', inputs, re.DOTALL)
83
+ if not s:
84
+ #print("err inputs: ", origin_inputs, flush=True)
85
+ return inputs
86
+
87
+ index = 0
88
+ for k in s:
89
+ try:
90
+ if "func" in k:
91
+ var = k.split("=", 1)
92
+ try:
93
+ var[1] = var[1].strip(" ")
94
+ exec(var[1], globals())
95
+ ans = func()
96
+ except:
97
+ if 'sympy' in var[1]:
98
+ var[1] = var[1].replace('res[x]', 'res[0][0]').replace('res[y]', 'res[0][1]')
99
+ exec(var[1], globals())
100
+ ans = func()
101
+ pass
102
+ var_list = [c.strip(" ") for c in var[0].split(",")]
103
+ if len(var_list) == 1:
104
+ ans = [ans]
105
+
106
+ for i in range(len(ans)):
107
+ try:
108
+ ans[i] = float(ans[i])
109
+ if abs(ans[i] - int(ans[i])) < 1e-10:
110
+ ans[i] = str(int(ans[i]))
111
+ except:
112
+ pass
113
+
114
+ inputs = inputs.replace("<<"+k+">>", "")
115
+ for i in range(len(var_list)):
116
+ inputs = inputs.replace(var_list[i], str(ans[i]))
117
+ index += 1
118
+ for c in range(index, len(s)):
119
+ for i in range(len(var_list)):
120
+ s[c] = s[c].replace(var_list[i], str(ans[i]))
121
+ else:
122
+ var = k.replace(" ", "").split("=")
123
+ var[1] = var[1].replace("eval", "")
124
+ ans = round(eval(var[1]), 10)
125
+ ans = float(ans)
126
+ if abs(ans - int(ans)) < 1e-10:
127
+ ans = str(int(ans))
128
+ inputs = inputs.replace("<<"+k+">>", "").replace(var[0], str(ans))
129
+ index += 1
130
+ for c in range(index, len(s)):
131
+ s[c] = s[c].replace(var[0], str(ans))
132
+ except:
133
+ return inputs
134
+ except Exception as e:
135
+ return inputs
136
+
137
+ return inputs
138
+
139
+
140
+ class TextIterStreamer:
141
+ def __init__(self, tokenizer, skip_prompt=False, skip_special_tokens=False, use_pot=True):
142
+ self.tokenizer = tokenizer
143
+ self.skip_prompt = skip_prompt
144
+ self.skip_special_tokens = skip_special_tokens
145
+ self.tokens = []
146
+ self.text_queue = Queue()
147
+ self.next_tokens_are_prompt = True
148
+ self.use_pot = use_pot
149
+
150
+ def put(self, value):
151
+ if self.skip_prompt and self.next_tokens_are_prompt:
152
+ self.next_tokens_are_prompt = False
153
+ else:
154
+ if len(value.shape) > 1:
155
+ value = value[0]
156
+ self.tokens.extend(value.tolist())
157
+ tokens_str = self.tokenizer.decode(self.tokens, skip_special_tokens=self.skip_special_tokens, errors='ignore')
158
+ if self.use_pot:
159
+ tokens_str = parse_pot_no_stream(tokens_str)
160
+ self.text_queue.put(tokens_str)
161
+
162
+ def end(self):
163
+ self.text_queue.put(None)
164
+
165
+ def __iter__(self):
166
+ return self
167
+
168
+ def __next__(self):
169
+ value = self.text_queue.get()
170
+ if value is None:
171
+ raise StopIteration()
172
+ else:
173
+ return value
174
+
175
+
176
+ class OutputRepetitionPenaltyLogitsProcessor(LogitsProcessor):
177
+ r"""
178
+ [`OutputLogitsProcessor`] that prevents the repetition of previous tokens through a penalty. This penalty is applied at
179
+ most once per token. Note that, for decoder-only models like most LLMs, the considered tokens include the prompt.
180
+
181
+ In the original [paper](https://arxiv.org/pdf/1909.05858.pdf), the authors suggest the use of a penalty of around
182
+ 1.2 to achieve a good balance between truthful generation and lack of repetition. To penalize and reduce
183
+ repetition, use `penalty` values above 1.0, where a higher value penalizes more strongly. To reward and encourage
184
+ repetition, use `penalty` values between 0.0 and 1.0, where a lower value rewards more strongly.
185
+
186
+ Args:
187
+ penalty (`float`):
188
+ The parameter for repetition penalty. 1.0 means no penalty. Above 1.0 penalizes previously generated
189
+ tokens. Between 0.0 and 1.0 rewards previously generated tokens.
190
+ """
191
+
192
+ def __init__(self, input_length: int,
193
+ presence_penalties: float = 1.0,
194
+ frequency_penalties: float = 0,
195
+ repetition_penalties: float = 0):
196
+ if not (repetition_penalties > 0):
197
+ raise ValueError(f"`repetition_penalties` has to be a strictly positive float, but is {repetition_penalties}")
198
+ if not ( (frequency_penalties >= -2) and (frequency_penalties <= 2) ):
199
+ raise ValueError(f"`frequency_penalties` has to be [-2, 2], but is {frequency_penalties}")
200
+ if not ( (presence_penalties >= -2) and (presence_penalties <= 2) ):
201
+ raise ValueError(f"`presence_penalties` has to be [-2, 2], but is {presence_penalties}")
202
+
203
+ self.repetition_penalties = repetition_penalties
204
+ self.frequency_penalties = frequency_penalties
205
+ self.presence_penalties = presence_penalties
206
+ self.input_length = input_length
207
+
208
+ def _get_bin_counts_and_mask(
209
+ self,
210
+ tokens: torch.Tensor,
211
+ vocab_size: int,
212
+ num_seqs: int,
213
+ ) -> Tuple[torch.Tensor, torch.Tensor]:
214
+ # Compute the bin counts for the tokens.
215
+ # vocab_size + 1 for padding.
216
+ bin_counts = torch.zeros((num_seqs, vocab_size + 1),
217
+ dtype=torch.long,
218
+ device=tokens.device)
219
+ bin_counts.scatter_add_(1, tokens, torch.ones_like(tokens))
220
+ bin_counts = bin_counts[:, :vocab_size]
221
+ mask = bin_counts > 0
222
+
223
+ return bin_counts, mask
224
+
225
+ @add_start_docstrings(LOGITS_PROCESSOR_INPUTS_DOCSTRING)
226
+ def __call__(self, input_ids: torch.LongTensor, logits: torch.FloatTensor) -> torch.FloatTensor:
227
+ prompt_tokens_tensor = input_ids[:, :self.input_length+1]
228
+ output_tokens_tensor = input_ids[:, self.input_length+1:]
229
+
230
+ num_seqs, vocab_size = logits.shape
231
+ _, prompt_mask = self._get_bin_counts_and_mask(
232
+ prompt_tokens_tensor, vocab_size, num_seqs)
233
+ output_bin_counts, output_mask = self._get_bin_counts_and_mask(
234
+ output_tokens_tensor, vocab_size, num_seqs)
235
+
236
+ repetition_penalties = torch.Tensor([self.repetition_penalties]).to(logits.device)
237
+ frequency_penalties = torch.Tensor([self.frequency_penalties]).to(logits.device)
238
+ presence_penalties = torch.Tensor([self.presence_penalties]).to(logits.device)
239
+
240
+ repetition_penalties = repetition_penalties[:, None].repeat(1, vocab_size)
241
+ repetition_penalties[~(prompt_mask | output_mask)] = 1.0
242
+ logits = torch.where(logits > 0, logits / repetition_penalties,
243
+ logits * repetition_penalties)
244
+
245
+ # We follow the definition in OpenAI API.
246
+ # Refer to https://platform.openai.com/docs/api-reference/parameter-details
247
+ logits -= frequency_penalties.unsqueeze_(dim=1) * output_bin_counts
248
+ logits -= presence_penalties.unsqueeze_(dim=1) * output_mask
249
+
250
+ return logits
model-00001-of-00008.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:856d462402916887c1048790a8ecc32175fca3065ad7035eb03a310295c2b1c7
3
+ size 1927349752
model-00002-of-00008.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:a7c2b1740bade360b6e17108bdf2e4b260751b5e0300ef0cc481e79952c53d35
3
+ size 1933783432
model-00003-of-00008.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:1d1d4bc41ceb88399374b89c5af26da531ee043683cd356c3b48bdd7e339773d
3
+ size 1990406544
model-00004-of-00008.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:2133b8691122b917f81a4892185572a98db7a0242f8182233cd884da3184cd7e
3
+ size 1923273016
model-00005-of-00008.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:2e8509f2facb86cdbde2b5ab0628a9e5f2dbfc374bd7ddcc15778015567c1e78
3
+ size 1933766856
model-00006-of-00008.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:4444d457c0138d14621208d4152d469aef0c20158eaf980e03b54e2ab49f458c
3
+ size 1933783464
model-00007-of-00008.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:8f487464447192d4efb4d306741609b65338a3a91fcf35ab467f08031fe20a0f
3
+ size 1933783472
model-00008-of-00008.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:2da863911d97309a5abd1d08271496323c50ca742502d3caf92d7daab07b8e99
3
+ size 1973487296
model.safetensors.index.json ADDED
@@ -0,0 +1,266 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "metadata": {
3
+ "total_size": 15549603840
4
+ },
5
+ "weight_map": {
6
+ "lm_head.weight": "model-00008-of-00008.safetensors",
7
+ "model.embed_tokens.weight": "model-00001-of-00008.safetensors",
8
+ "model.layers.0.input_layernorm.weight": "model-00001-of-00008.safetensors",
9
+ "model.layers.0.mlp.down_proj.weight": "model-00001-of-00008.safetensors",
10
+ "model.layers.0.mlp.gate_proj.weight": "model-00001-of-00008.safetensors",
11
+ "model.layers.0.mlp.up_proj.weight": "model-00001-of-00008.safetensors",
12
+ "model.layers.0.post_attention_layernorm.weight": "model-00001-of-00008.safetensors",
13
+ "model.layers.0.self_attn.o_proj.weight": "model-00001-of-00008.safetensors",
14
+ "model.layers.0.self_attn.qkv_proj.bias": "model-00001-of-00008.safetensors",
15
+ "model.layers.0.self_attn.qkv_proj.weight": "model-00001-of-00008.safetensors",
16
+ "model.layers.1.input_layernorm.weight": "model-00002-of-00008.safetensors",
17
+ "model.layers.1.mlp.down_proj.weight": "model-00002-of-00008.safetensors",
18
+ "model.layers.1.mlp.gate_proj.weight": "model-00001-of-00008.safetensors",
19
+ "model.layers.1.mlp.up_proj.weight": "model-00002-of-00008.safetensors",
20
+ "model.layers.1.post_attention_layernorm.weight": "model-00002-of-00008.safetensors",
21
+ "model.layers.1.self_attn.o_proj.weight": "model-00001-of-00008.safetensors",
22
+ "model.layers.1.self_attn.qkv_proj.bias": "model-00001-of-00008.safetensors",
23
+ "model.layers.1.self_attn.qkv_proj.weight": "model-00001-of-00008.safetensors",
24
+ "model.layers.10.input_layernorm.weight": "model-00003-of-00008.safetensors",
25
+ "model.layers.10.mlp.down_proj.weight": "model-00003-of-00008.safetensors",
26
+ "model.layers.10.mlp.gate_proj.weight": "model-00003-of-00008.safetensors",
27
+ "model.layers.10.mlp.up_proj.weight": "model-00003-of-00008.safetensors",
28
+ "model.layers.10.post_attention_layernorm.weight": "model-00003-of-00008.safetensors",
29
+ "model.layers.10.self_attn.o_proj.weight": "model-00003-of-00008.safetensors",
30
+ "model.layers.10.self_attn.qkv_proj.bias": "model-00003-of-00008.safetensors",
31
+ "model.layers.10.self_attn.qkv_proj.weight": "model-00003-of-00008.safetensors",
32
+ "model.layers.11.input_layernorm.weight": "model-00004-of-00008.safetensors",
33
+ "model.layers.11.mlp.down_proj.weight": "model-00004-of-00008.safetensors",
34
+ "model.layers.11.mlp.gate_proj.weight": "model-00004-of-00008.safetensors",
35
+ "model.layers.11.mlp.up_proj.weight": "model-00004-of-00008.safetensors",
36
+ "model.layers.11.post_attention_layernorm.weight": "model-00004-of-00008.safetensors",
37
+ "model.layers.11.self_attn.o_proj.weight": "model-00004-of-00008.safetensors",
38
+ "model.layers.11.self_attn.qkv_proj.bias": "model-00003-of-00008.safetensors",
39
+ "model.layers.11.self_attn.qkv_proj.weight": "model-00003-of-00008.safetensors",
40
+ "model.layers.12.input_layernorm.weight": "model-00004-of-00008.safetensors",
41
+ "model.layers.12.mlp.down_proj.weight": "model-00004-of-00008.safetensors",
42
+ "model.layers.12.mlp.gate_proj.weight": "model-00004-of-00008.safetensors",
43
+ "model.layers.12.mlp.up_proj.weight": "model-00004-of-00008.safetensors",
44
+ "model.layers.12.post_attention_layernorm.weight": "model-00004-of-00008.safetensors",
45
+ "model.layers.12.self_attn.o_proj.weight": "model-00004-of-00008.safetensors",
46
+ "model.layers.12.self_attn.qkv_proj.bias": "model-00004-of-00008.safetensors",
47
+ "model.layers.12.self_attn.qkv_proj.weight": "model-00004-of-00008.safetensors",
48
+ "model.layers.13.input_layernorm.weight": "model-00004-of-00008.safetensors",
49
+ "model.layers.13.mlp.down_proj.weight": "model-00004-of-00008.safetensors",
50
+ "model.layers.13.mlp.gate_proj.weight": "model-00004-of-00008.safetensors",
51
+ "model.layers.13.mlp.up_proj.weight": "model-00004-of-00008.safetensors",
52
+ "model.layers.13.post_attention_layernorm.weight": "model-00004-of-00008.safetensors",
53
+ "model.layers.13.self_attn.o_proj.weight": "model-00004-of-00008.safetensors",
54
+ "model.layers.13.self_attn.qkv_proj.bias": "model-00004-of-00008.safetensors",
55
+ "model.layers.13.self_attn.qkv_proj.weight": "model-00004-of-00008.safetensors",
56
+ "model.layers.14.input_layernorm.weight": "model-00004-of-00008.safetensors",
57
+ "model.layers.14.mlp.down_proj.weight": "model-00004-of-00008.safetensors",
58
+ "model.layers.14.mlp.gate_proj.weight": "model-00004-of-00008.safetensors",
59
+ "model.layers.14.mlp.up_proj.weight": "model-00004-of-00008.safetensors",
60
+ "model.layers.14.post_attention_layernorm.weight": "model-00004-of-00008.safetensors",
61
+ "model.layers.14.self_attn.o_proj.weight": "model-00004-of-00008.safetensors",
62
+ "model.layers.14.self_attn.qkv_proj.bias": "model-00004-of-00008.safetensors",
63
+ "model.layers.14.self_attn.qkv_proj.weight": "model-00004-of-00008.safetensors",
64
+ "model.layers.15.input_layernorm.weight": "model-00004-of-00008.safetensors",
65
+ "model.layers.15.mlp.down_proj.weight": "model-00004-of-00008.safetensors",
66
+ "model.layers.15.mlp.gate_proj.weight": "model-00004-of-00008.safetensors",
67
+ "model.layers.15.mlp.up_proj.weight": "model-00004-of-00008.safetensors",
68
+ "model.layers.15.post_attention_layernorm.weight": "model-00004-of-00008.safetensors",
69
+ "model.layers.15.self_attn.o_proj.weight": "model-00004-of-00008.safetensors",
70
+ "model.layers.15.self_attn.qkv_proj.bias": "model-00004-of-00008.safetensors",
71
+ "model.layers.15.self_attn.qkv_proj.weight": "model-00004-of-00008.safetensors",
72
+ "model.layers.16.input_layernorm.weight": "model-00005-of-00008.safetensors",
73
+ "model.layers.16.mlp.down_proj.weight": "model-00005-of-00008.safetensors",
74
+ "model.layers.16.mlp.gate_proj.weight": "model-00005-of-00008.safetensors",
75
+ "model.layers.16.mlp.up_proj.weight": "model-00005-of-00008.safetensors",
76
+ "model.layers.16.post_attention_layernorm.weight": "model-00005-of-00008.safetensors",
77
+ "model.layers.16.self_attn.o_proj.weight": "model-00005-of-00008.safetensors",
78
+ "model.layers.16.self_attn.qkv_proj.bias": "model-00005-of-00008.safetensors",
79
+ "model.layers.16.self_attn.qkv_proj.weight": "model-00005-of-00008.safetensors",
80
+ "model.layers.17.input_layernorm.weight": "model-00005-of-00008.safetensors",
81
+ "model.layers.17.mlp.down_proj.weight": "model-00005-of-00008.safetensors",
82
+ "model.layers.17.mlp.gate_proj.weight": "model-00005-of-00008.safetensors",
83
+ "model.layers.17.mlp.up_proj.weight": "model-00005-of-00008.safetensors",
84
+ "model.layers.17.post_attention_layernorm.weight": "model-00005-of-00008.safetensors",
85
+ "model.layers.17.self_attn.o_proj.weight": "model-00005-of-00008.safetensors",
86
+ "model.layers.17.self_attn.qkv_proj.bias": "model-00005-of-00008.safetensors",
87
+ "model.layers.17.self_attn.qkv_proj.weight": "model-00005-of-00008.safetensors",
88
+ "model.layers.18.input_layernorm.weight": "model-00005-of-00008.safetensors",
89
+ "model.layers.18.mlp.down_proj.weight": "model-00005-of-00008.safetensors",
90
+ "model.layers.18.mlp.gate_proj.weight": "model-00005-of-00008.safetensors",
91
+ "model.layers.18.mlp.up_proj.weight": "model-00005-of-00008.safetensors",
92
+ "model.layers.18.post_attention_layernorm.weight": "model-00005-of-00008.safetensors",
93
+ "model.layers.18.self_attn.o_proj.weight": "model-00005-of-00008.safetensors",
94
+ "model.layers.18.self_attn.qkv_proj.bias": "model-00005-of-00008.safetensors",
95
+ "model.layers.18.self_attn.qkv_proj.weight": "model-00005-of-00008.safetensors",
96
+ "model.layers.19.input_layernorm.weight": "model-00005-of-00008.safetensors",
97
+ "model.layers.19.mlp.down_proj.weight": "model-00005-of-00008.safetensors",
98
+ "model.layers.19.mlp.gate_proj.weight": "model-00005-of-00008.safetensors",
99
+ "model.layers.19.mlp.up_proj.weight": "model-00005-of-00008.safetensors",
100
+ "model.layers.19.post_attention_layernorm.weight": "model-00005-of-00008.safetensors",
101
+ "model.layers.19.self_attn.o_proj.weight": "model-00005-of-00008.safetensors",
102
+ "model.layers.19.self_attn.qkv_proj.bias": "model-00005-of-00008.safetensors",
103
+ "model.layers.19.self_attn.qkv_proj.weight": "model-00005-of-00008.safetensors",
104
+ "model.layers.2.input_layernorm.weight": "model-00002-of-00008.safetensors",
105
+ "model.layers.2.mlp.down_proj.weight": "model-00002-of-00008.safetensors",
106
+ "model.layers.2.mlp.gate_proj.weight": "model-00002-of-00008.safetensors",
107
+ "model.layers.2.mlp.up_proj.weight": "model-00002-of-00008.safetensors",
108
+ "model.layers.2.post_attention_layernorm.weight": "model-00002-of-00008.safetensors",
109
+ "model.layers.2.self_attn.o_proj.weight": "model-00002-of-00008.safetensors",
110
+ "model.layers.2.self_attn.qkv_proj.bias": "model-00002-of-00008.safetensors",
111
+ "model.layers.2.self_attn.qkv_proj.weight": "model-00002-of-00008.safetensors",
112
+ "model.layers.20.input_layernorm.weight": "model-00006-of-00008.safetensors",
113
+ "model.layers.20.mlp.down_proj.weight": "model-00006-of-00008.safetensors",
114
+ "model.layers.20.mlp.gate_proj.weight": "model-00005-of-00008.safetensors",
115
+ "model.layers.20.mlp.up_proj.weight": "model-00005-of-00008.safetensors",
116
+ "model.layers.20.post_attention_layernorm.weight": "model-00006-of-00008.safetensors",
117
+ "model.layers.20.self_attn.o_proj.weight": "model-00005-of-00008.safetensors",
118
+ "model.layers.20.self_attn.qkv_proj.bias": "model-00005-of-00008.safetensors",
119
+ "model.layers.20.self_attn.qkv_proj.weight": "model-00005-of-00008.safetensors",
120
+ "model.layers.21.input_layernorm.weight": "model-00006-of-00008.safetensors",
121
+ "model.layers.21.mlp.down_proj.weight": "model-00006-of-00008.safetensors",
122
+ "model.layers.21.mlp.gate_proj.weight": "model-00006-of-00008.safetensors",
123
+ "model.layers.21.mlp.up_proj.weight": "model-00006-of-00008.safetensors",
124
+ "model.layers.21.post_attention_layernorm.weight": "model-00006-of-00008.safetensors",
125
+ "model.layers.21.self_attn.o_proj.weight": "model-00006-of-00008.safetensors",
126
+ "model.layers.21.self_attn.qkv_proj.bias": "model-00006-of-00008.safetensors",
127
+ "model.layers.21.self_attn.qkv_proj.weight": "model-00006-of-00008.safetensors",
128
+ "model.layers.22.input_layernorm.weight": "model-00006-of-00008.safetensors",
129
+ "model.layers.22.mlp.down_proj.weight": "model-00006-of-00008.safetensors",
130
+ "model.layers.22.mlp.gate_proj.weight": "model-00006-of-00008.safetensors",
131
+ "model.layers.22.mlp.up_proj.weight": "model-00006-of-00008.safetensors",
132
+ "model.layers.22.post_attention_layernorm.weight": "model-00006-of-00008.safetensors",
133
+ "model.layers.22.self_attn.o_proj.weight": "model-00006-of-00008.safetensors",
134
+ "model.layers.22.self_attn.qkv_proj.bias": "model-00006-of-00008.safetensors",
135
+ "model.layers.22.self_attn.qkv_proj.weight": "model-00006-of-00008.safetensors",
136
+ "model.layers.23.input_layernorm.weight": "model-00006-of-00008.safetensors",
137
+ "model.layers.23.mlp.down_proj.weight": "model-00006-of-00008.safetensors",
138
+ "model.layers.23.mlp.gate_proj.weight": "model-00006-of-00008.safetensors",
139
+ "model.layers.23.mlp.up_proj.weight": "model-00006-of-00008.safetensors",
140
+ "model.layers.23.post_attention_layernorm.weight": "model-00006-of-00008.safetensors",
141
+ "model.layers.23.self_attn.o_proj.weight": "model-00006-of-00008.safetensors",
142
+ "model.layers.23.self_attn.qkv_proj.bias": "model-00006-of-00008.safetensors",
143
+ "model.layers.23.self_attn.qkv_proj.weight": "model-00006-of-00008.safetensors",
144
+ "model.layers.24.input_layernorm.weight": "model-00006-of-00008.safetensors",
145
+ "model.layers.24.mlp.down_proj.weight": "model-00006-of-00008.safetensors",
146
+ "model.layers.24.mlp.gate_proj.weight": "model-00006-of-00008.safetensors",
147
+ "model.layers.24.mlp.up_proj.weight": "model-00006-of-00008.safetensors",
148
+ "model.layers.24.post_attention_layernorm.weight": "model-00006-of-00008.safetensors",
149
+ "model.layers.24.self_attn.o_proj.weight": "model-00006-of-00008.safetensors",
150
+ "model.layers.24.self_attn.qkv_proj.bias": "model-00006-of-00008.safetensors",
151
+ "model.layers.24.self_attn.qkv_proj.weight": "model-00006-of-00008.safetensors",
152
+ "model.layers.25.input_layernorm.weight": "model-00007-of-00008.safetensors",
153
+ "model.layers.25.mlp.down_proj.weight": "model-00007-of-00008.safetensors",
154
+ "model.layers.25.mlp.gate_proj.weight": "model-00006-of-00008.safetensors",
155
+ "model.layers.25.mlp.up_proj.weight": "model-00007-of-00008.safetensors",
156
+ "model.layers.25.post_attention_layernorm.weight": "model-00007-of-00008.safetensors",
157
+ "model.layers.25.self_attn.o_proj.weight": "model-00006-of-00008.safetensors",
158
+ "model.layers.25.self_attn.qkv_proj.bias": "model-00006-of-00008.safetensors",
159
+ "model.layers.25.self_attn.qkv_proj.weight": "model-00006-of-00008.safetensors",
160
+ "model.layers.26.input_layernorm.weight": "model-00007-of-00008.safetensors",
161
+ "model.layers.26.mlp.down_proj.weight": "model-00007-of-00008.safetensors",
162
+ "model.layers.26.mlp.gate_proj.weight": "model-00007-of-00008.safetensors",
163
+ "model.layers.26.mlp.up_proj.weight": "model-00007-of-00008.safetensors",
164
+ "model.layers.26.post_attention_layernorm.weight": "model-00007-of-00008.safetensors",
165
+ "model.layers.26.self_attn.o_proj.weight": "model-00007-of-00008.safetensors",
166
+ "model.layers.26.self_attn.qkv_proj.bias": "model-00007-of-00008.safetensors",
167
+ "model.layers.26.self_attn.qkv_proj.weight": "model-00007-of-00008.safetensors",
168
+ "model.layers.27.input_layernorm.weight": "model-00007-of-00008.safetensors",
169
+ "model.layers.27.mlp.down_proj.weight": "model-00007-of-00008.safetensors",
170
+ "model.layers.27.mlp.gate_proj.weight": "model-00007-of-00008.safetensors",
171
+ "model.layers.27.mlp.up_proj.weight": "model-00007-of-00008.safetensors",
172
+ "model.layers.27.post_attention_layernorm.weight": "model-00007-of-00008.safetensors",
173
+ "model.layers.27.self_attn.o_proj.weight": "model-00007-of-00008.safetensors",
174
+ "model.layers.27.self_attn.qkv_proj.bias": "model-00007-of-00008.safetensors",
175
+ "model.layers.27.self_attn.qkv_proj.weight": "model-00007-of-00008.safetensors",
176
+ "model.layers.28.input_layernorm.weight": "model-00007-of-00008.safetensors",
177
+ "model.layers.28.mlp.down_proj.weight": "model-00007-of-00008.safetensors",
178
+ "model.layers.28.mlp.gate_proj.weight": "model-00007-of-00008.safetensors",
179
+ "model.layers.28.mlp.up_proj.weight": "model-00007-of-00008.safetensors",
180
+ "model.layers.28.post_attention_layernorm.weight": "model-00007-of-00008.safetensors",
181
+ "model.layers.28.self_attn.o_proj.weight": "model-00007-of-00008.safetensors",
182
+ "model.layers.28.self_attn.qkv_proj.bias": "model-00007-of-00008.safetensors",
183
+ "model.layers.28.self_attn.qkv_proj.weight": "model-00007-of-00008.safetensors",
184
+ "model.layers.29.input_layernorm.weight": "model-00007-of-00008.safetensors",
185
+ "model.layers.29.mlp.down_proj.weight": "model-00007-of-00008.safetensors",
186
+ "model.layers.29.mlp.gate_proj.weight": "model-00007-of-00008.safetensors",
187
+ "model.layers.29.mlp.up_proj.weight": "model-00007-of-00008.safetensors",
188
+ "model.layers.29.post_attention_layernorm.weight": "model-00007-of-00008.safetensors",
189
+ "model.layers.29.self_attn.o_proj.weight": "model-00007-of-00008.safetensors",
190
+ "model.layers.29.self_attn.qkv_proj.bias": "model-00007-of-00008.safetensors",
191
+ "model.layers.29.self_attn.qkv_proj.weight": "model-00007-of-00008.safetensors",
192
+ "model.layers.3.input_layernorm.weight": "model-00002-of-00008.safetensors",
193
+ "model.layers.3.mlp.down_proj.weight": "model-00002-of-00008.safetensors",
194
+ "model.layers.3.mlp.gate_proj.weight": "model-00002-of-00008.safetensors",
195
+ "model.layers.3.mlp.up_proj.weight": "model-00002-of-00008.safetensors",
196
+ "model.layers.3.post_attention_layernorm.weight": "model-00002-of-00008.safetensors",
197
+ "model.layers.3.self_attn.o_proj.weight": "model-00002-of-00008.safetensors",
198
+ "model.layers.3.self_attn.qkv_proj.bias": "model-00002-of-00008.safetensors",
199
+ "model.layers.3.self_attn.qkv_proj.weight": "model-00002-of-00008.safetensors",
200
+ "model.layers.30.input_layernorm.weight": "model-00008-of-00008.safetensors",
201
+ "model.layers.30.mlp.down_proj.weight": "model-00008-of-00008.safetensors",
202
+ "model.layers.30.mlp.gate_proj.weight": "model-00008-of-00008.safetensors",
203
+ "model.layers.30.mlp.up_proj.weight": "model-00008-of-00008.safetensors",
204
+ "model.layers.30.post_attention_layernorm.weight": "model-00008-of-00008.safetensors",
205
+ "model.layers.30.self_attn.o_proj.weight": "model-00007-of-00008.safetensors",
206
+ "model.layers.30.self_attn.qkv_proj.bias": "model-00007-of-00008.safetensors",
207
+ "model.layers.30.self_attn.qkv_proj.weight": "model-00007-of-00008.safetensors",
208
+ "model.layers.31.input_layernorm.weight": "model-00008-of-00008.safetensors",
209
+ "model.layers.31.mlp.down_proj.weight": "model-00008-of-00008.safetensors",
210
+ "model.layers.31.mlp.gate_proj.weight": "model-00008-of-00008.safetensors",
211
+ "model.layers.31.mlp.up_proj.weight": "model-00008-of-00008.safetensors",
212
+ "model.layers.31.post_attention_layernorm.weight": "model-00008-of-00008.safetensors",
213
+ "model.layers.31.self_attn.o_proj.weight": "model-00008-of-00008.safetensors",
214
+ "model.layers.31.self_attn.qkv_proj.bias": "model-00008-of-00008.safetensors",
215
+ "model.layers.31.self_attn.qkv_proj.weight": "model-00008-of-00008.safetensors",
216
+ "model.layers.4.input_layernorm.weight": "model-00002-of-00008.safetensors",
217
+ "model.layers.4.mlp.down_proj.weight": "model-00002-of-00008.safetensors",
218
+ "model.layers.4.mlp.gate_proj.weight": "model-00002-of-00008.safetensors",
219
+ "model.layers.4.mlp.up_proj.weight": "model-00002-of-00008.safetensors",
220
+ "model.layers.4.post_attention_layernorm.weight": "model-00002-of-00008.safetensors",
221
+ "model.layers.4.self_attn.o_proj.weight": "model-00002-of-00008.safetensors",
222
+ "model.layers.4.self_attn.qkv_proj.bias": "model-00002-of-00008.safetensors",
223
+ "model.layers.4.self_attn.qkv_proj.weight": "model-00002-of-00008.safetensors",
224
+ "model.layers.5.input_layernorm.weight": "model-00002-of-00008.safetensors",
225
+ "model.layers.5.mlp.down_proj.weight": "model-00002-of-00008.safetensors",
226
+ "model.layers.5.mlp.gate_proj.weight": "model-00002-of-00008.safetensors",
227
+ "model.layers.5.mlp.up_proj.weight": "model-00002-of-00008.safetensors",
228
+ "model.layers.5.post_attention_layernorm.weight": "model-00002-of-00008.safetensors",
229
+ "model.layers.5.self_attn.o_proj.weight": "model-00002-of-00008.safetensors",
230
+ "model.layers.5.self_attn.qkv_proj.bias": "model-00002-of-00008.safetensors",
231
+ "model.layers.5.self_attn.qkv_proj.weight": "model-00002-of-00008.safetensors",
232
+ "model.layers.6.input_layernorm.weight": "model-00003-of-00008.safetensors",
233
+ "model.layers.6.mlp.down_proj.weight": "model-00003-of-00008.safetensors",
234
+ "model.layers.6.mlp.gate_proj.weight": "model-00003-of-00008.safetensors",
235
+ "model.layers.6.mlp.up_proj.weight": "model-00003-of-00008.safetensors",
236
+ "model.layers.6.post_attention_layernorm.weight": "model-00003-of-00008.safetensors",
237
+ "model.layers.6.self_attn.o_proj.weight": "model-00002-of-00008.safetensors",
238
+ "model.layers.6.self_attn.qkv_proj.bias": "model-00002-of-00008.safetensors",
239
+ "model.layers.6.self_attn.qkv_proj.weight": "model-00002-of-00008.safetensors",
240
+ "model.layers.7.input_layernorm.weight": "model-00003-of-00008.safetensors",
241
+ "model.layers.7.mlp.down_proj.weight": "model-00003-of-00008.safetensors",
242
+ "model.layers.7.mlp.gate_proj.weight": "model-00003-of-00008.safetensors",
243
+ "model.layers.7.mlp.up_proj.weight": "model-00003-of-00008.safetensors",
244
+ "model.layers.7.post_attention_layernorm.weight": "model-00003-of-00008.safetensors",
245
+ "model.layers.7.self_attn.o_proj.weight": "model-00003-of-00008.safetensors",
246
+ "model.layers.7.self_attn.qkv_proj.bias": "model-00003-of-00008.safetensors",
247
+ "model.layers.7.self_attn.qkv_proj.weight": "model-00003-of-00008.safetensors",
248
+ "model.layers.8.input_layernorm.weight": "model-00003-of-00008.safetensors",
249
+ "model.layers.8.mlp.down_proj.weight": "model-00003-of-00008.safetensors",
250
+ "model.layers.8.mlp.gate_proj.weight": "model-00003-of-00008.safetensors",
251
+ "model.layers.8.mlp.up_proj.weight": "model-00003-of-00008.safetensors",
252
+ "model.layers.8.post_attention_layernorm.weight": "model-00003-of-00008.safetensors",
253
+ "model.layers.8.self_attn.o_proj.weight": "model-00003-of-00008.safetensors",
254
+ "model.layers.8.self_attn.qkv_proj.bias": "model-00003-of-00008.safetensors",
255
+ "model.layers.8.self_attn.qkv_proj.weight": "model-00003-of-00008.safetensors",
256
+ "model.layers.9.input_layernorm.weight": "model-00003-of-00008.safetensors",
257
+ "model.layers.9.mlp.down_proj.weight": "model-00003-of-00008.safetensors",
258
+ "model.layers.9.mlp.gate_proj.weight": "model-00003-of-00008.safetensors",
259
+ "model.layers.9.mlp.up_proj.weight": "model-00003-of-00008.safetensors",
260
+ "model.layers.9.post_attention_layernorm.weight": "model-00003-of-00008.safetensors",
261
+ "model.layers.9.self_attn.o_proj.weight": "model-00003-of-00008.safetensors",
262
+ "model.layers.9.self_attn.qkv_proj.bias": "model-00003-of-00008.safetensors",
263
+ "model.layers.9.self_attn.qkv_proj.weight": "model-00003-of-00008.safetensors",
264
+ "model.norm.weight": "model-00008-of-00008.safetensors"
265
+ }
266
+ }
modeling_zhinao.py ADDED
@@ -0,0 +1,1058 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright (c) 360zhinao and the HuggingFace Inc. team. All rights reserved.
2
+ # This code is built upon Huggingface's transformers repository.
3
+
4
+ import math
5
+ import warnings
6
+ from threading import Thread
7
+ from typing import List, Optional, Tuple, Union
8
+
9
+ import torch
10
+ import torch.nn.functional as F
11
+ import torch.utils.checkpoint
12
+ from torch import nn
13
+ from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
14
+
15
+ from transformers.activations import ACT2FN
16
+ from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast, SequenceClassifierOutputWithPast
17
+ from transformers.modeling_utils import PreTrainedModel
18
+ from transformers.utils import logging
19
+ from transformers.generation.utils import GenerationConfig
20
+ from transformers.generation.logits_process import LogitsProcessorList
21
+ from .configuration_zhinao import ZhinaoConfig
22
+ from .generation_utils import TextIterStreamer, make_context, OutputRepetitionPenaltyLogitsProcessor, parse_pot_no_stream
23
+
24
+
25
+ try:
26
+ from flash_attn import flash_attn_varlen_func
27
+ from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input
28
+ except:
29
+ flash_attn_varlen_func = None
30
+ index_first_axis, pad_input, unpad_input = None, None, None
31
+
32
+
33
+ logger = logging.get_logger(__name__)
34
+
35
+ _CONFIG_FOR_DOC = "ZhinaoConfig"
36
+
37
+
38
+ def _get_unpad_data(attention_mask):
39
+ seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32)
40
+ indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()
41
+ max_seqlen_in_batch = seqlens_in_batch.max().item()
42
+ cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.torch.int32), (1, 0))
43
+ return (
44
+ indices,
45
+ cu_seqlens,
46
+ max_seqlen_in_batch,
47
+ )
48
+
49
+
50
+ # Copied from transformers.models.bart.modeling_bart._make_causal_mask
51
+ def _make_causal_mask(
52
+ input_ids_shape: torch.Size, dtype: torch.dtype, device: torch.device, past_key_values_length: int = 0
53
+ ):
54
+ """
55
+ Make causal mask used for bi-directional self-attention.
56
+ """
57
+ bsz, tgt_len = input_ids_shape
58
+ mask = torch.full((tgt_len, tgt_len), torch.finfo(dtype).min, device=device)
59
+ mask_cond = torch.arange(mask.size(-1), device=device)
60
+ mask.masked_fill_(mask_cond < (mask_cond + 1).view(mask.size(-1), 1), 0)
61
+ mask = mask.to(dtype)
62
+
63
+ if past_key_values_length > 0:
64
+ mask = torch.cat([torch.zeros(tgt_len, past_key_values_length, dtype=dtype, device=device), mask], dim=-1)
65
+ return mask[None, None, :, :].expand(bsz, 1, tgt_len, tgt_len + past_key_values_length)
66
+
67
+
68
+ # Copied from transformers.models.bart.modeling_bart._expand_mask
69
+ def _expand_mask(mask: torch.Tensor, dtype: torch.dtype, tgt_len: Optional[int] = None):
70
+ """
71
+ Expands attention_mask from `[bsz, seq_len]` to `[bsz, 1, tgt_seq_len, src_seq_len]`.
72
+ """
73
+ bsz, src_len = mask.size()
74
+ tgt_len = tgt_len if tgt_len is not None else src_len
75
+
76
+ expanded_mask = mask[:, None, None, :].expand(bsz, 1, tgt_len, src_len).to(dtype)
77
+
78
+ inverted_mask = 1.0 - expanded_mask
79
+
80
+ return inverted_mask.masked_fill(inverted_mask.to(torch.bool), torch.finfo(dtype).min)
81
+
82
+
83
+ class ZhinaoRMSNorm(nn.Module):
84
+ def __init__(self, hidden_size, eps=1e-6):
85
+ """
86
+ ZhinaoRMSNorm is equivalent to T5LayerNorm
87
+ """
88
+ super().__init__()
89
+ self.weight = nn.Parameter(torch.ones(hidden_size))
90
+ self.variance_epsilon = eps
91
+
92
+ def forward(self, hidden_states):
93
+ input_dtype = hidden_states.dtype
94
+ hidden_states = hidden_states.to(torch.float32)
95
+ variance = hidden_states.pow(2).mean(-1, keepdim=True)
96
+ hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
97
+ return self.weight * hidden_states.to(input_dtype)
98
+
99
+
100
+ class ZhinaoRotaryEmbedding(torch.nn.Module):
101
+ def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None):
102
+ super().__init__()
103
+
104
+ self.dim = dim
105
+ self.max_position_embeddings = max_position_embeddings
106
+ self.base = base
107
+ inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2).float().to(device) / self.dim))
108
+ self.register_buffer("inv_freq", inv_freq, persistent=False)
109
+
110
+ # Build here to make `torch.jit.trace` work.
111
+ self._set_cos_sin_cache(
112
+ seq_len=max_position_embeddings, device=self.inv_freq.device, dtype=torch.get_default_dtype()
113
+ )
114
+
115
+ def _set_cos_sin_cache(self, seq_len, device, dtype):
116
+ self.max_seq_len_cached = seq_len
117
+ t = torch.arange(self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype)
118
+
119
+ freqs = torch.einsum("i,j->ij", t, self.inv_freq)
120
+ # Different from paper, but it uses a different permutation in order to obtain the same calculation
121
+ emb = torch.cat((freqs, freqs), dim=-1)
122
+ self.register_buffer("cos_cached", emb.cos()[None, None, :, :].to(dtype), persistent=False)
123
+ self.register_buffer("sin_cached", emb.sin()[None, None, :, :].to(dtype), persistent=False)
124
+
125
+ def forward(self, x, seq_len=None):
126
+ # x: [bs, num_attention_heads, seq_len, head_size]
127
+ if seq_len > self.max_seq_len_cached:
128
+ self._set_cos_sin_cache(seq_len=seq_len, device=x.device, dtype=x.dtype)
129
+
130
+ return (
131
+ self.cos_cached[:, :, :seq_len, ...].to(dtype=x.dtype),
132
+ self.sin_cached[:, :, :seq_len, ...].to(dtype=x.dtype),
133
+ )
134
+
135
+
136
+ class ZhinaoLinearScalingRotaryEmbedding(ZhinaoRotaryEmbedding):
137
+ """ZhinaoRotaryEmbedding extended with linear scaling. Credits to the Reddit user /u/kaiokendev"""
138
+
139
+ def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None, scaling_factor=1.0):
140
+ self.scaling_factor = scaling_factor
141
+ super().__init__(dim, max_position_embeddings, base, device)
142
+
143
+ def _set_cos_sin_cache(self, seq_len, device, dtype):
144
+ self.max_seq_len_cached = seq_len
145
+ t = torch.arange(self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype)
146
+ t = t / self.scaling_factor
147
+
148
+ freqs = torch.einsum("i,j->ij", t, self.inv_freq)
149
+ # Different from paper, but it uses a different permutation in order to obtain the same calculation
150
+ emb = torch.cat((freqs, freqs), dim=-1)
151
+ self.register_buffer("cos_cached", emb.cos()[None, None, :, :].to(dtype), persistent=False)
152
+ self.register_buffer("sin_cached", emb.sin()[None, None, :, :].to(dtype), persistent=False)
153
+
154
+
155
+ class ZhinaoDynamicNTKScalingRotaryEmbedding(ZhinaoRotaryEmbedding):
156
+ """ZhinaoRotaryEmbedding extended with Dynamic NTK scaling. Credits to the Reddit users /u/bloc97 and /u/emozilla"""
157
+
158
+ def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None, scaling_factor=1.0):
159
+ self.scaling_factor = scaling_factor
160
+ super().__init__(dim, max_position_embeddings, base, device)
161
+
162
+ def _set_cos_sin_cache(self, seq_len, device, dtype):
163
+ self.max_seq_len_cached = seq_len
164
+
165
+ if seq_len > self.max_position_embeddings:
166
+ base = self.base * (
167
+ (self.scaling_factor * seq_len / self.max_position_embeddings) - (self.scaling_factor - 1)
168
+ ) ** (self.dim / (self.dim - 2))
169
+ inv_freq = 1.0 / (base ** (torch.arange(0, self.dim, 2).float().to(device) / self.dim))
170
+ self.register_buffer("inv_freq", inv_freq, persistent=False)
171
+
172
+ t = torch.arange(self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype)
173
+
174
+ freqs = torch.einsum("i,j->ij", t, self.inv_freq)
175
+ # Different from paper, but it uses a different permutation in order to obtain the same calculation
176
+ emb = torch.cat((freqs, freqs), dim=-1)
177
+ self.register_buffer("cos_cached", emb.cos()[None, None, :, :].to(dtype), persistent=False)
178
+ self.register_buffer("sin_cached", emb.sin()[None, None, :, :].to(dtype), persistent=False)
179
+
180
+
181
+ class ZhinaoNTKScalingRotaryEmbedding(torch.nn.Module):
182
+ def __init__(self, dim, max_position_embeddings=2048, base=10000, scaling_factor=100, device=None):
183
+ super().__init__()
184
+
185
+ self.dim = dim
186
+ self.max_position_embeddings = max_position_embeddings
187
+ self.base = base * scaling_factor
188
+ inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2).float().to(device) / self.dim))
189
+ self.register_buffer("inv_freq", inv_freq, persistent=False)
190
+
191
+ # Build here to make `torch.jit.trace` work.
192
+ self._set_cos_sin_cache(
193
+ seq_len=max_position_embeddings, device=self.inv_freq.device, dtype=torch.get_default_dtype()
194
+ )
195
+
196
+ def _set_cos_sin_cache(self, seq_len, device, dtype):
197
+ self.max_seq_len_cached = seq_len
198
+ t = torch.arange(self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype)
199
+ freqs = torch.einsum("i,j->ij", t, self.inv_freq)
200
+ emb = torch.cat((freqs, freqs), dim=-1)
201
+ self.register_buffer("cos_cached", emb.cos()[None, None, :, :].to(dtype), persistent=False)
202
+ self.register_buffer("sin_cached", emb.sin()[None, None, :, :].to(dtype), persistent=False)
203
+
204
+ def forward(self, x, seq_len=None):
205
+ if seq_len > self.max_seq_len_cached:
206
+ self._set_cos_sin_cache(seq_len=seq_len, device=x.device, dtype=x.dtype)
207
+
208
+ return (
209
+ self.cos_cached[:, :, :seq_len, ...].to(dtype=x.dtype),
210
+ self.sin_cached[:, :, :seq_len, ...].to(dtype=x.dtype),
211
+ )
212
+
213
+
214
+ def rotate_half(x):
215
+ """Rotates half the hidden dims of the input."""
216
+ x1 = x[..., : x.shape[-1] // 2]
217
+ x2 = x[..., x.shape[-1] // 2 :]
218
+ return torch.cat((-x2, x1), dim=-1)
219
+
220
+
221
+ def apply_rotary_pos_emb(q, k, cos, sin, position_ids):
222
+ # The first two dimensions of cos and sin are always 1, so we can `squeeze` them.
223
+ cos = cos.squeeze(1).squeeze(0) # [seq_len, dim]
224
+ sin = sin.squeeze(1).squeeze(0) # [seq_len, dim]
225
+ cos = cos[position_ids].unsqueeze(1) # [bs, 1, seq_len, dim]
226
+ sin = sin[position_ids].unsqueeze(1) # [bs, 1, seq_len, dim]
227
+ q_embed = (q * cos) + (rotate_half(q) * sin)
228
+ k_embed = (k * cos) + (rotate_half(k) * sin)
229
+ return q_embed, k_embed
230
+
231
+
232
+ class ZhinaoMLP(nn.Module):
233
+ def __init__(self, config):
234
+ super().__init__()
235
+ self.config = config
236
+ self.hidden_size = config.hidden_size
237
+ self.intermediate_size = config.intermediate_size
238
+ self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
239
+ self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
240
+ self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
241
+ self.act_fn = ACT2FN[config.hidden_act]
242
+
243
+ def forward(self, x):
244
+ intermediate = self.act_fn(self.gate_proj(x)) * self.up_proj(x)
245
+ down_proj = self.down_proj(intermediate)
246
+ return down_proj
247
+
248
+
249
+ def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
250
+ """
251
+ This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
252
+ num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
253
+ """
254
+ batch, num_key_value_heads, slen, head_dim = hidden_states.shape
255
+ if n_rep == 1:
256
+ return hidden_states
257
+ hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
258
+ return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
259
+
260
+
261
+ class ZhinaoAttention(nn.Module):
262
+ """Multi-headed attention from 'Attention Is All You Need' paper"""
263
+
264
+ def __init__(self, config: ZhinaoConfig):
265
+ super().__init__()
266
+ self.config = config
267
+ self.hidden_size = config.hidden_size
268
+ self.num_heads = config.num_attention_heads
269
+ self.head_dim = self.hidden_size // self.num_heads
270
+ self.num_key_value_heads = config.num_key_value_heads
271
+ self.num_key_value_groups = self.num_heads // self.num_key_value_heads
272
+ self.max_position_embeddings = config.max_position_embeddings
273
+ self.rope_theta = config.rope_theta
274
+ self.is_causal = True
275
+ self.dropout = 0.0
276
+ self.use_flash_attn = config.use_flash_attn
277
+
278
+ if (self.head_dim * self.num_heads) != self.hidden_size:
279
+ raise ValueError(
280
+ f"hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size}"
281
+ f" and `num_heads`: {self.num_heads})."
282
+ )
283
+
284
+ self.qkv_hidden_size = (self.num_heads + 2 * self.num_key_value_heads) * self.head_dim
285
+ self.qkv_proj = nn.Linear(self.hidden_size, self.qkv_hidden_size, bias=True)
286
+ self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False)
287
+ self._init_rope()
288
+
289
+ def _init_rope(self):
290
+ if self.config.rope_scaling is None:
291
+ self.rotary_emb = ZhinaoRotaryEmbedding(
292
+ self.head_dim,
293
+ max_position_embeddings=self.max_position_embeddings,
294
+ base=self.rope_theta,
295
+ )
296
+ else:
297
+ scaling_type = self.config.rope_scaling["type"]
298
+ scaling_factor = self.config.rope_scaling["factor"]
299
+ if scaling_type == "linear":
300
+ self.rotary_emb = ZhinaoLinearScalingRotaryEmbedding(
301
+ self.head_dim,
302
+ max_position_embeddings=self.max_position_embeddings,
303
+ scaling_factor=scaling_factor,
304
+ base=self.rope_theta,
305
+ )
306
+ elif scaling_type == "dynamic":
307
+ self.rotary_emb = ZhinaoDynamicNTKScalingRotaryEmbedding(
308
+ self.head_dim,
309
+ max_position_embeddings=self.max_position_embeddings,
310
+ scaling_factor=scaling_factor,
311
+ base=self.rope_theta,
312
+ )
313
+ elif scaling_type == "ntk":
314
+ self.rotary_emb = ZhinaoNTKScalingRotaryEmbedding(
315
+ self.head_dim,
316
+ max_position_embeddings=self.max_position_embeddings,
317
+ scaling_factor=scaling_factor,
318
+ base=self.rope_theta,
319
+ )
320
+ else:
321
+ raise ValueError(f"Unknown RoPE scaling type {scaling_type}")
322
+
323
+ def raw_attention(self, query_states, key_states, value_states, attention_mask):
324
+ attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim)
325
+
326
+ if attention_mask is not None:
327
+ attn_weights = attn_weights + attention_mask
328
+
329
+ # upcast attention to fp32
330
+ attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype)
331
+ attn_output = torch.matmul(attn_weights, value_states)
332
+
333
+ attn_output = attn_output.transpose(1, 2).contiguous()
334
+
335
+ return attn_output
336
+
337
+ def flash_attention(self, query_states, key_states, value_states, attention_mask):
338
+ # TODO: These transpose are quite inefficient but Flash Attention requires the layout [batch_size, sequence_length, num_heads, head_dim]. We would need to refactor the KV cache
339
+ # to be able to avoid many of these transpose/reshape/view.
340
+ query_states = query_states.transpose(1, 2)
341
+ key_states = key_states.transpose(1, 2)
342
+ value_states = value_states.transpose(1, 2)
343
+
344
+ batch_size, query_length = query_states.shape[0], query_states.shape[1]
345
+ query_states, key_states, value_states, indices_q, cu_seq_lens, max_seq_lens = self._upad_input(
346
+ query_states, key_states, value_states, attention_mask, query_length
347
+ )
348
+
349
+ cu_seqlens_q, cu_seqlens_k = cu_seq_lens
350
+ max_seqlen_in_batch_q, max_seqlen_in_batch_k = max_seq_lens
351
+
352
+ attn_output_unpad = flash_attn_varlen_func(
353
+ query_states,
354
+ key_states,
355
+ value_states,
356
+ cu_seqlens_q=cu_seqlens_q,
357
+ cu_seqlens_k=cu_seqlens_k,
358
+ max_seqlen_q=max_seqlen_in_batch_q,
359
+ max_seqlen_k=max_seqlen_in_batch_k,
360
+ dropout_p=self.dropout,
361
+ softmax_scale=None,
362
+ causal=self.is_causal,
363
+ )
364
+ attn_output = pad_input(attn_output_unpad, indices_q, batch_size, query_length)
365
+ return attn_output
366
+
367
+ def forward(
368
+ self,
369
+ hidden_states: torch.Tensor,
370
+ attention_mask: Optional[torch.Tensor] = None,
371
+ position_ids: Optional[torch.LongTensor] = None,
372
+ past_key_value: Optional[Tuple[torch.Tensor]] = None,
373
+ output_attentions: bool = False,
374
+ use_cache: bool = False,
375
+ ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
376
+ bsz, q_len, _ = hidden_states.size()
377
+
378
+ mixed_x_layer = self.qkv_proj(hidden_states)
379
+ new_tensor_shape = mixed_x_layer.size()[:-1] + \
380
+ (self.num_key_value_heads, ((self.num_heads // self.num_key_value_heads + 2) * self.head_dim))
381
+ mixed_x_layer = mixed_x_layer.view(*new_tensor_shape)
382
+ query, key_states, value_states = torch.split(
383
+ mixed_x_layer,
384
+ [self.num_heads // self.num_key_value_heads * self.head_dim, self.head_dim, self.head_dim],
385
+ dim=3
386
+ )
387
+ # [sq, b, ng, np/ng * hn] -> [sq, b, np, hn]
388
+ query_states = query.contiguous().view(query.size(0), query.size(1), -1, self.head_dim)
389
+
390
+ query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
391
+ key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
392
+ value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
393
+
394
+ kv_seq_len = key_states.shape[-2]
395
+ if past_key_value is not None:
396
+ kv_seq_len += past_key_value[0].shape[-2]
397
+ cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
398
+ query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)
399
+
400
+ if past_key_value is not None:
401
+ # reuse k, v, self_attention
402
+ key_states = torch.cat([past_key_value[0], key_states], dim=2)
403
+ value_states = torch.cat([past_key_value[1], value_states], dim=2)
404
+
405
+ past_key_value = (key_states, value_states) if use_cache else None
406
+
407
+ # repeat k/v heads if n_kv_heads < n_heads
408
+ key_states = repeat_kv(key_states, self.num_key_value_groups)
409
+ value_states = repeat_kv(value_states, self.num_key_value_groups)
410
+
411
+ # q, k, v: [b, n, s, h]
412
+ # check attention mask
413
+ if self.use_flash_attn:
414
+ if attention_mask is not None and attention_mask.size() != (bsz, kv_seq_len):
415
+ raise ValueError(f"Attention mask should be of size {(bsz, kv_seq_len)}, but is {attention_mask.size()}")
416
+ attn_output = self.flash_attention(query_states, key_states, value_states, attention_mask)
417
+ else:
418
+ if attention_mask is not None and attention_mask.size() != (bsz, 1, q_len, kv_seq_len):
419
+ raise ValueError(f"Attention mask should be of size {bsz, 1, q_len, kv_seq_len}, but is {attention_mask.size()}")
420
+ attn_output = self.raw_attention(query_states, key_states, value_states, attention_mask)
421
+
422
+ attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)
423
+ attn_output = self.o_proj(attn_output)
424
+
425
+ if not output_attentions:
426
+ attn_weights = None
427
+
428
+ return attn_output, attn_weights, past_key_value
429
+
430
+ # Copied from transformers.models.mistral.modeling_mistral.MistralFlashAttention2._upad_input
431
+ def _upad_input(self, query_layer, key_layer, value_layer, attention_mask, query_length):
432
+ batch_size, kv_seq_len, num_heads, head_dim = key_layer.shape
433
+
434
+ # On the first iteration we need to properly re-create the padding mask
435
+ # by slicing it on the proper place
436
+ if kv_seq_len != attention_mask.shape[-1]:
437
+ attention_mask_num_tokens = attention_mask.shape[-1]
438
+ attention_mask = attention_mask[:, attention_mask_num_tokens - kv_seq_len :]
439
+
440
+ indices_k, cu_seqlens_k, max_seqlen_in_batch_k = _get_unpad_data(attention_mask)
441
+
442
+ key_layer = index_first_axis(key_layer.reshape(batch_size * kv_seq_len, num_heads, head_dim), indices_k)
443
+ value_layer = index_first_axis(value_layer.reshape(batch_size * kv_seq_len, num_heads, head_dim), indices_k)
444
+
445
+ if query_length == kv_seq_len:
446
+ query_layer = index_first_axis(
447
+ query_layer.reshape(batch_size * kv_seq_len, num_heads, head_dim), indices_k
448
+ )
449
+ cu_seqlens_q = cu_seqlens_k
450
+ max_seqlen_in_batch_q = max_seqlen_in_batch_k
451
+ indices_q = indices_k
452
+ elif query_length == 1:
453
+ max_seqlen_in_batch_q = 1
454
+ cu_seqlens_q = torch.arange(
455
+ batch_size + 1, dtype=torch.int32, device=query_layer.device
456
+ ) # There is a memcpy here, that is very bad.
457
+ indices_q = cu_seqlens_q[:-1]
458
+ query_layer = query_layer.squeeze(1)
459
+ else:
460
+ # The -q_len: slice assumes left padding.
461
+ attention_mask = attention_mask[:, -query_length:]
462
+
463
+ query_layer, indices_q, cu_seqlens_q, max_seqlen_in_batch_q = unpad_input(query_layer, attention_mask)
464
+ return (
465
+ query_layer,
466
+ key_layer,
467
+ value_layer,
468
+ indices_q,
469
+ (cu_seqlens_q, cu_seqlens_k),
470
+ (max_seqlen_in_batch_q, max_seqlen_in_batch_k),
471
+ )
472
+
473
+
474
+ class ZhinaoDecoderLayer(nn.Module):
475
+ def __init__(self, config: ZhinaoConfig):
476
+ super().__init__()
477
+ self.hidden_size = config.hidden_size
478
+
479
+ self.self_attn = ZhinaoAttention(config=config)
480
+ self.mlp = ZhinaoMLP(config)
481
+ self.input_layernorm = ZhinaoRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
482
+ self.post_attention_layernorm = ZhinaoRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
483
+
484
+ def forward(
485
+ self,
486
+ hidden_states: torch.Tensor,
487
+ attention_mask: Optional[torch.Tensor] = None,
488
+ position_ids: Optional[torch.LongTensor] = None,
489
+ past_key_value: Optional[Tuple[torch.Tensor]] = None,
490
+ output_attentions: Optional[bool] = False,
491
+ use_cache: Optional[bool] = False,
492
+ ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]:
493
+ """
494
+ Args:
495
+ hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`
496
+ attention_mask (`torch.FloatTensor`, *optional*):
497
+ attention mask of size `(batch_size, sequence_length)` if flash attention is used or `(batch_size, 1,
498
+ query_sequence_length, key_sequence_length)` if default attention is used.
499
+ output_attentions (`bool`, *optional*):
500
+ Whether or not to return the attentions tensors of all attention layers. See `attentions` under
501
+ returned tensors for more detail.
502
+ use_cache (`bool`, *optional*):
503
+ If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding
504
+ (see `past_key_values`).
505
+ past_key_value (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states
506
+ """
507
+
508
+ residual = hidden_states
509
+
510
+ hidden_states = self.input_layernorm(hidden_states)
511
+
512
+ # Self Attention
513
+ hidden_states, self_attn_weights, present_key_value = self.self_attn(
514
+ hidden_states=hidden_states,
515
+ attention_mask=attention_mask,
516
+ position_ids=position_ids,
517
+ past_key_value=past_key_value,
518
+ output_attentions=output_attentions,
519
+ use_cache=use_cache,
520
+ )
521
+ hidden_states = residual + hidden_states
522
+
523
+ # Fully Connected
524
+ residual = hidden_states
525
+ hidden_states = self.post_attention_layernorm(hidden_states)
526
+ hidden_states = self.mlp(hidden_states)
527
+ hidden_states = residual + hidden_states
528
+
529
+ outputs = (hidden_states,)
530
+
531
+ if output_attentions:
532
+ outputs += (self_attn_weights,)
533
+
534
+ if use_cache:
535
+ outputs += (present_key_value,)
536
+
537
+ return outputs
538
+
539
+
540
+ class ZhinaoPreTrainedModel(PreTrainedModel):
541
+ config_class = ZhinaoConfig
542
+ base_model_prefix = "model"
543
+ supports_gradient_checkpointing = True
544
+ _no_split_modules = ["ZhinaoDecoderLayer"]
545
+ _skip_keys_device_placement = "past_key_values"
546
+
547
+ def _init_weights(self, module):
548
+ std = self.config.initializer_range
549
+ if isinstance(module, nn.Linear):
550
+ module.weight.data.normal_(mean=0.0, std=std)
551
+ if module.bias is not None:
552
+ module.bias.data.zero_()
553
+ elif isinstance(module, nn.Embedding):
554
+ module.weight.data.normal_(mean=0.0, std=std)
555
+ if module.padding_idx is not None:
556
+ module.weight.data[module.padding_idx].zero_()
557
+
558
+ def _set_gradient_checkpointing(self, module, value=False):
559
+ if isinstance(module, ZhinaoModel):
560
+ module.gradient_checkpointing = value
561
+
562
+
563
+ class ZhinaoModel(ZhinaoPreTrainedModel):
564
+ """
565
+ Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`ZhinaoDecoderLayer`]
566
+
567
+ Args:
568
+ config: ZhinaoConfig
569
+ """
570
+
571
+ def __init__(self, config: ZhinaoConfig):
572
+ super().__init__(config)
573
+ self.padding_idx = config.pad_token_id
574
+ self.vocab_size = config.vocab_size
575
+ self.use_flash_attn = config.use_flash_attn
576
+
577
+ self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
578
+ self.layers = nn.ModuleList([ZhinaoDecoderLayer(config) for _ in range(config.num_hidden_layers)])
579
+ self.norm = ZhinaoRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
580
+
581
+ self.gradient_checkpointing = False
582
+ # Initialize weights and apply final processing
583
+ self.post_init()
584
+
585
+ def get_input_embeddings(self):
586
+ return self.embed_tokens
587
+
588
+ def set_input_embeddings(self, value):
589
+ self.embed_tokens = value
590
+
591
+ # Copied from transformers.models.bart.modeling_bart.BartDecoder._prepare_decoder_attention_mask
592
+ def _prepare_decoder_attention_mask(self, attention_mask, input_shape, inputs_embeds, past_key_values_length):
593
+ # create causal mask
594
+ # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
595
+ combined_attention_mask = None
596
+ if input_shape[-1] > 1:
597
+ combined_attention_mask = _make_causal_mask(
598
+ input_shape,
599
+ inputs_embeds.dtype,
600
+ device=inputs_embeds.device,
601
+ past_key_values_length=past_key_values_length,
602
+ )
603
+
604
+ if attention_mask is not None:
605
+ # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
606
+ expanded_attn_mask = _expand_mask(attention_mask, inputs_embeds.dtype, tgt_len=input_shape[-1]).to(
607
+ inputs_embeds.device
608
+ )
609
+ combined_attention_mask = (
610
+ expanded_attn_mask if combined_attention_mask is None else expanded_attn_mask + combined_attention_mask
611
+ )
612
+
613
+ return combined_attention_mask
614
+
615
+ def forward(
616
+ self,
617
+ input_ids: torch.LongTensor = None,
618
+ attention_mask: Optional[torch.Tensor] = None,
619
+ position_ids: Optional[torch.LongTensor] = None,
620
+ past_key_values: Optional[List[torch.FloatTensor]] = None,
621
+ inputs_embeds: Optional[torch.FloatTensor] = None,
622
+ use_cache: Optional[bool] = None,
623
+ output_attentions: Optional[bool] = None,
624
+ output_hidden_states: Optional[bool] = None,
625
+ return_dict: Optional[bool] = None,
626
+ ) -> Union[Tuple, BaseModelOutputWithPast]:
627
+ output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
628
+ output_hidden_states = (
629
+ output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
630
+ )
631
+ use_cache = use_cache if use_cache is not None else self.config.use_cache
632
+
633
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
634
+
635
+ # retrieve input_ids and inputs_embeds
636
+ if input_ids is not None and inputs_embeds is not None:
637
+ raise ValueError("You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time")
638
+ elif input_ids is not None:
639
+ batch_size, seq_length = input_ids.shape
640
+ elif inputs_embeds is not None:
641
+ batch_size, seq_length, _ = inputs_embeds.shape
642
+ else:
643
+ raise ValueError("You have to specify either decoder_input_ids or decoder_inputs_embeds")
644
+
645
+ seq_length_with_past = seq_length
646
+ past_key_values_length = 0
647
+
648
+ if past_key_values is not None:
649
+ past_key_values_length = past_key_values[0][0].shape[2]
650
+ seq_length_with_past = seq_length_with_past + past_key_values_length
651
+
652
+ if position_ids is None:
653
+ device = input_ids.device if input_ids is not None else inputs_embeds.device
654
+ position_ids = torch.arange(
655
+ past_key_values_length, seq_length + past_key_values_length, dtype=torch.long, device=device
656
+ )
657
+ position_ids = position_ids.unsqueeze(0).view(-1, seq_length)
658
+ else:
659
+ position_ids = position_ids.view(-1, seq_length).long()
660
+
661
+ if inputs_embeds is None:
662
+ inputs_embeds = self.embed_tokens(input_ids)
663
+ # embed positions
664
+ if attention_mask is None:
665
+ attention_mask = torch.ones(
666
+ (batch_size, seq_length_with_past), dtype=torch.bool, device=inputs_embeds.device
667
+ )
668
+
669
+ # (batch_size, 1, seq_length, seq_length)` if default attention is used
670
+ if not self.use_flash_attn:
671
+ attention_mask = self._prepare_decoder_attention_mask(
672
+ attention_mask, (batch_size, seq_length), inputs_embeds, past_key_values_length
673
+ )
674
+
675
+ hidden_states = inputs_embeds
676
+
677
+ if self.gradient_checkpointing and self.training:
678
+ if use_cache:
679
+ logger.warning_once(
680
+ "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
681
+ )
682
+ use_cache = False
683
+
684
+ # decoder layers
685
+ all_hidden_states = () if output_hidden_states else None
686
+ all_self_attns = () if output_attentions else None
687
+ next_decoder_cache = () if use_cache else None
688
+
689
+ for idx, decoder_layer in enumerate(self.layers):
690
+ if output_hidden_states:
691
+ all_hidden_states += (hidden_states,)
692
+
693
+ past_key_value = past_key_values[idx] if past_key_values is not None else None
694
+
695
+ if self.gradient_checkpointing and self.training:
696
+
697
+ def create_custom_forward(module):
698
+ def custom_forward(*inputs):
699
+ # None for past_key_value
700
+ return module(*inputs, past_key_value, output_attentions)
701
+
702
+ return custom_forward
703
+
704
+ layer_outputs = torch.utils.checkpoint.checkpoint(
705
+ create_custom_forward(decoder_layer),
706
+ hidden_states,
707
+ attention_mask,
708
+ position_ids,
709
+ )
710
+ else:
711
+ layer_outputs = decoder_layer(
712
+ hidden_states,
713
+ attention_mask=attention_mask,
714
+ position_ids=position_ids,
715
+ past_key_value=past_key_value,
716
+ output_attentions=output_attentions,
717
+ use_cache=use_cache,
718
+ )
719
+
720
+ hidden_states = layer_outputs[0]
721
+
722
+ if use_cache:
723
+ next_decoder_cache += (layer_outputs[2 if output_attentions else 1],)
724
+
725
+ if output_attentions:
726
+ all_self_attns += (layer_outputs[1],)
727
+
728
+ hidden_states = self.norm(hidden_states)
729
+
730
+ # add hidden states from the last decoder layer
731
+ if output_hidden_states:
732
+ all_hidden_states += (hidden_states,)
733
+
734
+ next_cache = next_decoder_cache if use_cache else None
735
+ if not return_dict:
736
+ return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None)
737
+
738
+ return BaseModelOutputWithPast(
739
+ last_hidden_state=hidden_states,
740
+ past_key_values=next_cache,
741
+ hidden_states=all_hidden_states,
742
+ attentions=all_self_attns,
743
+ )
744
+
745
+
746
+ class ZhinaoForCausalLM(ZhinaoPreTrainedModel):
747
+ _tied_weights_keys = ["lm_head.weight"]
748
+
749
+ def __init__(self, config):
750
+ super().__init__(config)
751
+ self.model = ZhinaoModel(config)
752
+ self.vocab_size = config.vocab_size
753
+ self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
754
+
755
+ # Initialize weights and apply final processing
756
+ if config.bf16:
757
+ self.model.bfloat16()
758
+ self.lm_head.bfloat16()
759
+ if config.fp16:
760
+ self.model.half()
761
+ self.lm_head.half()
762
+
763
+ if config.use_flash_attn == "auto":
764
+ if flash_attn_varlen_func:
765
+ if config.bf16 or config.fp16:
766
+ logger.warn("Try importing flash-attention.")
767
+ config.use_flash_attn = True
768
+ else:
769
+ config.use_flash_attn = False
770
+ logger.warn("Flash attention will be disabled because it does NOT support fp32.")
771
+ else:
772
+ config.use_flash_attn = False
773
+ logger.warn("Please install FlashAttention first, " "e.g., with pip install flash-attn")
774
+
775
+ self.post_init()
776
+
777
+ def get_input_embeddings(self):
778
+ return self.model.embed_tokens
779
+
780
+ def set_input_embeddings(self, value):
781
+ self.model.embed_tokens = value
782
+
783
+ def get_output_embeddings(self):
784
+ return self.lm_head
785
+
786
+ def set_output_embeddings(self, new_embeddings):
787
+ self.lm_head = new_embeddings
788
+
789
+ def set_decoder(self, decoder):
790
+ self.model = decoder
791
+
792
+ def get_decoder(self):
793
+ return self.model
794
+
795
+ def forward(
796
+ self,
797
+ input_ids: torch.LongTensor = None,
798
+ attention_mask: Optional[torch.Tensor] = None,
799
+ position_ids: Optional[torch.LongTensor] = None,
800
+ past_key_values: Optional[List[torch.FloatTensor]] = None,
801
+ inputs_embeds: Optional[torch.FloatTensor] = None,
802
+ labels: Optional[torch.LongTensor] = None,
803
+ use_cache: Optional[bool] = None,
804
+ output_attentions: Optional[bool] = None,
805
+ output_hidden_states: Optional[bool] = None,
806
+ return_dict: Optional[bool] = None,
807
+ ) -> Union[Tuple, CausalLMOutputWithPast]:
808
+
809
+ output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
810
+ output_hidden_states = (
811
+ output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
812
+ )
813
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
814
+
815
+ # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
816
+ outputs = self.model(
817
+ input_ids=input_ids,
818
+ attention_mask=attention_mask,
819
+ position_ids=position_ids,
820
+ past_key_values=past_key_values,
821
+ inputs_embeds=inputs_embeds,
822
+ use_cache=use_cache,
823
+ output_attentions=output_attentions,
824
+ output_hidden_states=output_hidden_states,
825
+ return_dict=return_dict,
826
+ )
827
+
828
+ hidden_states = outputs[0]
829
+ logits = self.lm_head(hidden_states)
830
+
831
+ # warn:Huge gpu memory
832
+ logits = logits.float()
833
+
834
+ loss = None
835
+ if labels is not None:
836
+ # Shift so that tokens < n predict n
837
+ shift_logits = logits[..., :-1, :].contiguous()
838
+ shift_labels = labels[..., 1:].contiguous()
839
+ # Flatten the tokens
840
+ loss_fct = CrossEntropyLoss()
841
+ shift_logits = shift_logits.view(-1, self.config.vocab_size)
842
+ shift_labels = shift_labels.view(-1)
843
+ # Enable model parallelism
844
+ shift_labels = shift_labels.to(shift_logits.device)
845
+ loss = loss_fct(shift_logits, shift_labels)
846
+
847
+ if not return_dict:
848
+ output = (logits,) + outputs[1:]
849
+ return (loss,) + output if loss is not None else output
850
+
851
+ return CausalLMOutputWithPast(
852
+ loss=loss,
853
+ logits=logits,
854
+ past_key_values=outputs.past_key_values,
855
+ hidden_states=outputs.hidden_states,
856
+ attentions=outputs.attentions,
857
+ )
858
+
859
+ def prepare_inputs_for_generation(
860
+ self, input_ids, past_key_values=None, attention_mask=None, inputs_embeds=None, **kwargs
861
+ ):
862
+ if past_key_values:
863
+ input_ids = input_ids[:, -1:]
864
+
865
+ position_ids = kwargs.get("position_ids", None)
866
+ if attention_mask is not None and position_ids is None:
867
+ # create position_ids on the fly for batch generation
868
+ position_ids = attention_mask.long().cumsum(-1) - 1
869
+ position_ids.masked_fill_(attention_mask == 0, 1)
870
+ if past_key_values:
871
+ position_ids = position_ids[:, -1].unsqueeze(-1)
872
+
873
+ # if `inputs_embeds` are passed, we only want to use them in the 1st generation step
874
+ if inputs_embeds is not None and past_key_values is None:
875
+ model_inputs = {"inputs_embeds": inputs_embeds}
876
+ else:
877
+ model_inputs = {"input_ids": input_ids}
878
+
879
+ model_inputs.update(
880
+ {
881
+ "position_ids": position_ids,
882
+ "past_key_values": past_key_values,
883
+ "use_cache": kwargs.get("use_cache"),
884
+ "attention_mask": attention_mask,
885
+ }
886
+ )
887
+ return model_inputs
888
+
889
+ @staticmethod
890
+ def _reorder_cache(past_key_values, beam_idx):
891
+ reordered_past = ()
892
+ for layer_past in past_key_values:
893
+ reordered_past += (
894
+ tuple(past_state.index_select(0, beam_idx.to(past_state.device)) for past_state in layer_past),
895
+ )
896
+ return reordered_past
897
+
898
+
899
+ def generate(
900
+ self,
901
+ inputs: Optional[torch.Tensor] = None,
902
+ generation_config: Optional[GenerationConfig] = None,
903
+ streamer = None,
904
+ **kwargs,
905
+ ):
906
+ repetition_penalty = kwargs.pop("repetition_penalty", generation_config.repetition_penalty)
907
+ generation_config.repetition_penalty = 1.0
908
+
909
+ logits_processor = None
910
+ if repetition_penalty > 1.0:
911
+ warnings.warn("We highly recommend using OpenAI's frequency and presence penalty instead of the original repetition penalty. The original repetition penalty penalizes prompt tokens, which may lead to various potential issues. Therefore, your repetition penalty coefficient will be transformed into frequency penalty and presence penalty.", UserWarning)
912
+ presence_penalty = repetition_penalty - 1.0
913
+ frequency_penalty = repetition_penalty - 1.0
914
+ logits_processor = LogitsProcessorList(
915
+ [OutputRepetitionPenaltyLogitsProcessor(inputs.size(1), presence_penalty, frequency_penalty, 1.0)]
916
+ )
917
+
918
+ response = super().generate(
919
+ inputs,
920
+ generation_config=generation_config,
921
+ logits_processor=logits_processor,
922
+ streamer=streamer,
923
+ **kwargs,
924
+ )
925
+ generation_config.repetition_penalty = repetition_penalty
926
+ return response
927
+
928
+
929
+ def chat(
930
+ self,
931
+ tokenizer,
932
+ messages: List[dict],
933
+ system: str = "You are a helpful assistant.",
934
+ stream=False,
935
+ use_pot=True,
936
+ generation_config: Optional[GenerationConfig]=None):
937
+
938
+ generation_config = generation_config or self.generation_config
939
+ input_ids = make_context(
940
+ model=self, tokenizer=tokenizer, messages=messages,
941
+ system=system, max_new_tokens=generation_config.max_new_tokens
942
+ )
943
+
944
+ if stream:
945
+ streamer = TextIterStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True, use_pot=use_pot)
946
+ Thread(target=self.generate, kwargs=dict(
947
+ inputs=input_ids, streamer=streamer,
948
+ generation_config=generation_config,
949
+ )).start()
950
+ return streamer
951
+ else:
952
+ outputs = self.generate(input_ids, generation_config=generation_config)
953
+ response = tokenizer.decode(outputs[0][len(input_ids[0]):], skip_special_tokens=True)
954
+ if use_pot:
955
+ response = parse_pot_no_stream(response)
956
+ return response
957
+
958
+
959
+ class ZhinaoForSequenceClassification(ZhinaoPreTrainedModel):
960
+ def __init__(self, config):
961
+ super().__init__(config)
962
+ self.num_labels = config.num_labels
963
+ self.model = ZhinaoModel(config)
964
+ self.score = nn.Linear(config.hidden_size, self.num_labels, bias=False)
965
+
966
+ # Initialize weights and apply final processing
967
+ self.post_init()
968
+
969
+ def get_input_embeddings(self):
970
+ return self.model.embed_tokens
971
+
972
+ def set_input_embeddings(self, value):
973
+ self.model.embed_tokens = value
974
+
975
+ def forward(
976
+ self,
977
+ input_ids: torch.LongTensor = None,
978
+ attention_mask: Optional[torch.Tensor] = None,
979
+ position_ids: Optional[torch.LongTensor] = None,
980
+ past_key_values: Optional[List[torch.FloatTensor]] = None,
981
+ inputs_embeds: Optional[torch.FloatTensor] = None,
982
+ labels: Optional[torch.LongTensor] = None,
983
+ use_cache: Optional[bool] = None,
984
+ output_attentions: Optional[bool] = None,
985
+ output_hidden_states: Optional[bool] = None,
986
+ return_dict: Optional[bool] = None,
987
+ ) -> Union[Tuple, SequenceClassifierOutputWithPast]:
988
+
989
+
990
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
991
+
992
+ transformer_outputs = self.model(
993
+ input_ids,
994
+ attention_mask=attention_mask,
995
+ position_ids=position_ids,
996
+ past_key_values=past_key_values,
997
+ inputs_embeds=inputs_embeds,
998
+ use_cache=use_cache,
999
+ output_attentions=output_attentions,
1000
+ output_hidden_states=output_hidden_states,
1001
+ return_dict=return_dict,
1002
+ )
1003
+ hidden_states = transformer_outputs[0]
1004
+ logits = self.score(hidden_states)
1005
+
1006
+ if input_ids is not None:
1007
+ batch_size = input_ids.shape[0]
1008
+ else:
1009
+ batch_size = inputs_embeds.shape[0]
1010
+
1011
+ if self.config.pad_token_id is None and batch_size != 1:
1012
+ raise ValueError("Cannot handle batch sizes > 1 if no padding token is defined.")
1013
+ if self.config.pad_token_id is None:
1014
+ sequence_lengths = -1
1015
+ else:
1016
+ if input_ids is not None:
1017
+ sequence_lengths = (torch.eq(input_ids, self.config.pad_token_id).long().argmax(-1) - 1).to(
1018
+ logits.device
1019
+ )
1020
+ else:
1021
+ sequence_lengths = -1
1022
+
1023
+ pooled_logits = logits[torch.arange(batch_size, device=logits.device), sequence_lengths]
1024
+
1025
+ loss = None
1026
+ if labels is not None:
1027
+ labels = labels.to(logits.device)
1028
+ if self.config.problem_type is None:
1029
+ if self.num_labels == 1:
1030
+ self.config.problem_type = "regression"
1031
+ elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
1032
+ self.config.problem_type = "single_label_classification"
1033
+ else:
1034
+ self.config.problem_type = "multi_label_classification"
1035
+
1036
+ if self.config.problem_type == "regression":
1037
+ loss_fct = MSELoss()
1038
+ if self.num_labels == 1:
1039
+ loss = loss_fct(pooled_logits.squeeze(), labels.squeeze())
1040
+ else:
1041
+ loss = loss_fct(pooled_logits, labels)
1042
+ elif self.config.problem_type == "single_label_classification":
1043
+ loss_fct = CrossEntropyLoss()
1044
+ loss = loss_fct(pooled_logits.view(-1, self.num_labels), labels.view(-1))
1045
+ elif self.config.problem_type == "multi_label_classification":
1046
+ loss_fct = BCEWithLogitsLoss()
1047
+ loss = loss_fct(pooled_logits, labels)
1048
+ if not return_dict:
1049
+ output = (pooled_logits,) + transformer_outputs[1:]
1050
+ return ((loss,) + output) if loss is not None else output
1051
+
1052
+ return SequenceClassifierOutputWithPast(
1053
+ loss=loss,
1054
+ logits=pooled_logits,
1055
+ past_key_values=transformer_outputs.past_key_values,
1056
+ hidden_states=transformer_outputs.hidden_states,
1057
+ attentions=transformer_outputs.attentions,
1058
+ )
special_tokens_map.json ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ {
2
+ "pad_token": "<pad>"
3
+ }
tokenization_zhinao.py ADDED
@@ -0,0 +1,279 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ import torch
3
+ import base64
4
+ import tiktoken
5
+ from typing import Collection, Optional, Dict, List, Set, Tuple, Union
6
+ from transformers import PreTrainedTokenizer
7
+ from transformers.utils import PaddingStrategy
8
+ from transformers.tokenization_utils import PreTrainedTokenizer
9
+
10
+
11
+ PAT_STR = r"""(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+"""
12
+
13
+
14
+ class SPTokenizer:
15
+ def __init__(self, model_path):
16
+ self.vocab_file = model_path
17
+ self.pad_token = '<pad>'
18
+ self.unk_token = '<unk>'
19
+ self.mask_token = '<mask>'
20
+ self.eod_token = '<eod>'
21
+ self.eop_token = '<eop>'
22
+ self.im_start_token = '<|im_start|>'
23
+ self.im_end_token = '<|im_end|>'
24
+
25
+ ## special_tokens
26
+ self.SPECIAL_TOKENS = (
27
+ self.pad_token,
28
+ self.unk_token,
29
+ self.mask_token,
30
+ self.eod_token,
31
+ self.eop_token,
32
+ '[space2]', '[space3]', '[space4]', '[space8]',
33
+ self.im_start_token, self.im_end_token
34
+ )
35
+ self.bulid_tokenizer()
36
+ self.out = self.output_core_token()
37
+
38
+ self.token2strs = {
39
+ "[space2]": " ",
40
+ "[space3]": " ",
41
+ "[space4]": " ",
42
+ "[space8]": " ",
43
+ }
44
+ self.str2tokens = {v: k for k, v in self.token2strs.items()}
45
+ self.sorted_strs = sorted(list(self.str2tokens.keys()),
46
+ key=lambda x: len(x), reverse=True)
47
+
48
+ ## skip_special_tokens
49
+ self.decode_skip_special_tokens = [
50
+ self.pad_token,
51
+ self.unk_token,
52
+ self.mask_token,
53
+ self.eod_token,
54
+ self.eop_token,
55
+ self.im_start_token,
56
+ self.im_end_token]
57
+ self.decode_skip_special_tokens_ids = [self.convert_token_to_id(token) for token in self.decode_skip_special_tokens]
58
+
59
+ def _load_tiktoken_bpe(self, tiktoken_bpe_file: str):
60
+ with open(tiktoken_bpe_file, "rb") as f:
61
+ contents = f.read()
62
+ return {
63
+ base64.b64decode(token): int(rank)
64
+ for token, rank in (line.split() for line in contents.splitlines() if line)
65
+ }
66
+
67
+ def bulid_tokenizer(self):
68
+ mergeable_ranks = self._load_tiktoken_bpe(self.vocab_file)
69
+ special_tokens = {
70
+ token: index
71
+ for index, token in enumerate(
72
+ self.SPECIAL_TOKENS, start=len(mergeable_ranks)
73
+ )
74
+ }
75
+ encode = tiktoken.Encoding(
76
+ "zhinao",
77
+ pat_str=PAT_STR,
78
+ mergeable_ranks=mergeable_ranks,
79
+ special_tokens=special_tokens
80
+ )
81
+ decoder = {v: k for k, v in mergeable_ranks.items()}
82
+ decoder.update({v: k for k, v in special_tokens.items()})
83
+ decoder_token2id = {v: k for k, v in decoder.items()}
84
+
85
+ self.tokenizer = encode
86
+ self.decoder = decoder
87
+ self.decoder_token2id = decoder_token2id
88
+ self.num_tokens = len(mergeable_ranks) + len(self.SPECIAL_TOKENS)
89
+
90
+ def output_core_token(self):
91
+ """output special tokens"""
92
+ out = {}
93
+ for t in self.SPECIAL_TOKENS:
94
+ out[t] = self.convert_token_to_id(t)
95
+ return out
96
+
97
+ def tokenize(
98
+ self,
99
+ text,
100
+ allowed_special: Union[Set, str] = "all",
101
+ disallowed_special: Union[Collection, str] = ()):
102
+ tokens = []
103
+ text = self.convert(text)
104
+ for idx in self.tokenizer.encode(text, allowed_special=allowed_special, disallowed_special=disallowed_special):
105
+ tokens.append(self.decoder[idx])
106
+ return tokens
107
+
108
+ def encode(self, text, allowed_special="all", disallowed_special=()):
109
+ """text to id"""
110
+ text = self.convert(text)
111
+ return self.tokenizer.encode(text, allowed_special=allowed_special, disallowed_special=disallowed_special)
112
+
113
+ def decode(self, ids, errors="replace"):
114
+ """id to text"""
115
+ text = self.tokenizer.decode(ids, errors=errors)
116
+ return self.deconvert(text)
117
+
118
+ def decode_tokens(self, tokens: List[str]) -> str:
119
+ """
120
+ Converts a sequence of tokens in a single string.
121
+ """
122
+ text = ""
123
+ temp = b""
124
+ for t in tokens:
125
+ if isinstance(t, str):
126
+ if temp:
127
+ text += temp.decode("utf-8", errors="replace")
128
+ temp = b""
129
+ text += t
130
+ elif isinstance(t, bytes):
131
+ temp += t
132
+ else:
133
+ raise TypeError("token should only be of type bytes or str")
134
+ if temp:
135
+ text += temp.decode("utf-8", errors="replace")
136
+ return self.deconvert(text)
137
+
138
+ def convert_id_to_token(self, idx):
139
+ return self.decoder[idx]
140
+
141
+ def convert_token_to_id(self, token):
142
+ return self.decoder_token2id[token]
143
+
144
+ def convert(self, text):
145
+ """将文本的特殊字符转换成特殊token"""
146
+ for k in ["[br]", "<br>"]:
147
+ text = text.replace(k, "\n")
148
+ for k in self.sorted_strs:
149
+ if k in text:
150
+ text = text.replace(k, self.str2tokens[k])
151
+ return text
152
+
153
+ def deconvert(self, text):
154
+ """将解码文本恢复原始字符"""
155
+ for t in self.token2strs:
156
+ if t in text:
157
+ text = text.replace(t, self.token2strs[t])
158
+ return text
159
+
160
+
161
+ class ZhinaoTokenizer(PreTrainedTokenizer):
162
+ vocab_files_names = {"vocab_file": "vocab/360.tiktoken"}
163
+ model_input_names = ["input_ids", "attention_mask"]
164
+
165
+ def __init__(self, vocab_file, padding_side="left", clean_up_tokenization_spaces=False, **kwargs):
166
+ self.name = "ZhinaoTokenizer"
167
+ self.errors = "replace"
168
+ self.vocab_file = vocab_file
169
+ self.tokenizer = SPTokenizer(model_path=vocab_file)
170
+ try:
171
+ kwargs.pop('eos_token')
172
+ kwargs.pop('pad_token')
173
+ kwargs.pop('unk_token')
174
+ except:
175
+ pass
176
+ super().__init__(padding_side=padding_side, clean_up_tokenization_spaces=clean_up_tokenization_spaces, **kwargs)
177
+ self.pad_token_id = self.tokenizer.convert_token_to_id(self.tokenizer.pad_token)
178
+ self.eod_id = self.tokenizer.convert_token_to_id(self.tokenizer.eod_token)
179
+ self.im_start_id = self.tokenizer.convert_token_to_id(self.tokenizer.im_start_token)
180
+ self.im_end_id = self.tokenizer.convert_token_to_id(self.tokenizer.im_end_token)
181
+ from icecream import ic
182
+ ic(
183
+ self.eos_token_id,
184
+ self.pad_token_id,
185
+ self.im_start_id,
186
+ self.im_end_id)
187
+
188
+ @property
189
+ def unk_token(self) -> str:
190
+ return self.tokenizer.unk_token
191
+
192
+ @property
193
+ def pad_token(self) -> str:
194
+ return self.tokenizer.pad_token
195
+
196
+ @property
197
+ def eos_token(self) -> str:
198
+ return self.tokenizer.eod_token
199
+
200
+ @property
201
+ def eos_token_id(self):
202
+ return self.tokenizer.convert_token_to_id(self.tokenizer.eod_token)
203
+
204
+ @property
205
+ def eop_token(self) -> str:
206
+ return self.tokenizer.eop_token
207
+
208
+ @property
209
+ def eop_token_id(self):
210
+ return self.tokenizer.convert_token_to_id(self.tokenizer.eop_token)
211
+
212
+ @property
213
+ def vocab_size(self):
214
+ return self.tokenizer.num_tokens
215
+
216
+ def get_vocab(self):
217
+ """ Returns vocab as a dict """
218
+ vocab = {self._convert_id_to_token(i): i for i in range(self.vocab_size)}
219
+ vocab.update(self.added_tokens_encoder)
220
+ return vocab
221
+
222
+ def tokenize(
223
+ self,
224
+ text: str,
225
+ allowed_special: Union[Set, str] = "all",
226
+ disallowed_special: Union[Collection, str] = (),
227
+ ) -> List[Union[bytes, str]]:
228
+ tokens = []
229
+ for t in self.tokenizer.encode(
230
+ text, allowed_special=allowed_special, disallowed_special=disallowed_special
231
+ ):
232
+ tokens.append(self.tokenizer.decoder[t])
233
+ return tokens
234
+
235
+ def _decode(
236
+ self,
237
+ token_ids: Union[int, List[int]],
238
+ skip_special_tokens: bool = False,
239
+ errors: str = None,
240
+ **kwargs,
241
+ ) -> str:
242
+ if isinstance(token_ids, int):
243
+ token_ids = [token_ids]
244
+ if skip_special_tokens:
245
+ token_ids = [i for i in token_ids if i not in self.tokenizer.decode_skip_special_tokens_ids]
246
+ return self.tokenizer.decode(token_ids, errors=errors or self.errors)
247
+
248
+ def _tokenize(self, text, **kwargs):
249
+ raise NotImplementedError
250
+
251
+ def _convert_token_to_id(self, token):
252
+ """ Converts a token (str) in an id using the vocab. """
253
+ return self.tokenizer.convert_token_to_id(token)
254
+
255
+ def _convert_id_to_token(self, index):
256
+ """Converts an index (integer) in a token (str) using the vocab. """
257
+ return self.tokenizer.convert_id_to_token(index)
258
+
259
+ def convert_tokens_to_string(self, tokens: List[str]) -> str:
260
+ """
261
+ Converts a sequence of tokens in a single string.
262
+ """
263
+ return self.tokenizer.decode_tokens(tokens)
264
+
265
+ def save_vocabulary(self, save_directory, filename_prefix=None):
266
+ """Save only the vocabulary of the tokenizer (vocabulary). """
267
+ if os.path.isdir(save_directory):
268
+ vocab_file = os.path.join(save_directory, self.vocab_files_names["vocab_file"])
269
+ else:
270
+ vocab_file = save_directory
271
+
272
+ with open(self.vocab_file, 'rb') as fin:
273
+ proto_str = fin.read()
274
+
275
+ os.makedirs(save_directory + "/vocab", exist_ok=True)
276
+ with open(vocab_file, "wb") as writer:
277
+ writer.write(proto_str)
278
+
279
+ return (vocab_file,)
tokenizer_config.json ADDED
@@ -0,0 +1,18 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "added_tokens_decoder": {},
3
+ "auto_map": {
4
+ "AutoTokenizer": [
5
+ "tokenization_zhinao.ZhinaoTokenizer",
6
+ null
7
+ ]
8
+ },
9
+ "clean_up_tokenization_spaces": false,
10
+ "do_lower_case": false,
11
+ "eos_token": "<eod>",
12
+ "model_max_length": 4096,
13
+ "pad_token": "<pad>",
14
+ "padding_side": "left",
15
+ "remove_space": false,
16
+ "tokenizer_class": "ZhinaoTokenizer",
17
+ "unk_token": "<unk>"
18
+ }
vocab/360.tiktoken ADDED
The diff for this file is too large to render. See raw diff