manaladan6 Stanislas commited on
Commit
db71889
·
0 Parent(s):

Duplicate from zai-org/codegeex4-all-9b

Browse files

Co-authored-by: Qinkai Zheng <Stanislas@users.noreply.huggingface.co>

.gitattributes ADDED
@@ -0,0 +1,36 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ *.7z filter=lfs diff=lfs merge=lfs -text
2
+ *.arrow filter=lfs diff=lfs merge=lfs -text
3
+ *.bin filter=lfs diff=lfs merge=lfs -text
4
+ *.bz2 filter=lfs diff=lfs merge=lfs -text
5
+ *.ckpt filter=lfs diff=lfs merge=lfs -text
6
+ *.ftz filter=lfs diff=lfs merge=lfs -text
7
+ *.gz filter=lfs diff=lfs merge=lfs -text
8
+ *.h5 filter=lfs diff=lfs merge=lfs -text
9
+ *.joblib filter=lfs diff=lfs merge=lfs -text
10
+ *.lfs.* filter=lfs diff=lfs merge=lfs -text
11
+ *.mlmodel filter=lfs diff=lfs merge=lfs -text
12
+ *.model filter=lfs diff=lfs merge=lfs -text
13
+ *.msgpack filter=lfs diff=lfs merge=lfs -text
14
+ *.npy filter=lfs diff=lfs merge=lfs -text
15
+ *.npz filter=lfs diff=lfs merge=lfs -text
16
+ *.onnx filter=lfs diff=lfs merge=lfs -text
17
+ *.ot filter=lfs diff=lfs merge=lfs -text
18
+ *.parquet filter=lfs diff=lfs merge=lfs -text
19
+ *.pb filter=lfs diff=lfs merge=lfs -text
20
+ *.pickle filter=lfs diff=lfs merge=lfs -text
21
+ *.pkl filter=lfs diff=lfs merge=lfs -text
22
+ *.pt filter=lfs diff=lfs merge=lfs -text
23
+ *.pth filter=lfs diff=lfs merge=lfs -text
24
+ *.rar filter=lfs diff=lfs merge=lfs -text
25
+ *.safetensors filter=lfs diff=lfs merge=lfs -text
26
+ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
27
+ *.tar.* filter=lfs diff=lfs merge=lfs -text
28
+ *.tar filter=lfs diff=lfs merge=lfs -text
29
+ *.tflite filter=lfs diff=lfs merge=lfs -text
30
+ *.tgz filter=lfs diff=lfs merge=lfs -text
31
+ *.wasm filter=lfs diff=lfs merge=lfs -text
32
+ *.xz filter=lfs diff=lfs merge=lfs -text
33
+ *.zip filter=lfs diff=lfs merge=lfs -text
34
+ *.zst filter=lfs diff=lfs merge=lfs -text
35
+ *tfevents* filter=lfs diff=lfs merge=lfs -text
36
+ tokenizer.json filter=lfs diff=lfs merge=lfs -text
LICENSE ADDED
@@ -0,0 +1,70 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ The CodeGeeX4 License
2
+
3
+ 1. 定义
4
+
5
+ “许可方”是指分发其软件的 CodeGeeX 团队。
6
+ “软件”是指根据本许可提供的 CodeGeeX4 模型参数。
7
+
8
+ 2. 许可授予
9
+
10
+ 根据本许可的条款和条件,许可方特此授予您非排他性、全球性、不可转让、不可再许可、可撤销、免版税的版权许可。
11
+ 本许可允许您免费使用本仓库中的所有开源模型进行学术研究,对于希望将模型用于商业目的的用户,需在[这里](https://open.bigmodel.cn/mla/form)完成登记。经过登记的用户可以免费使用本模型进行商业活动,但必须遵守本许可的所有条款和条件。
12
+ 上述版权声明和本许可声明应包含在本软件的所有副本或重要部分中。
13
+ 如果您分发或提供 THUDM / 智谱AI 关于 CodeGeeX4 开源模型的材料(或其任何衍生作品),或使用其中任何材料(包括 CodeGeeX4 系列的所有开源模型)的产品或服务,您应:
14
+
15
+ (A) 随任何此类 THUDM / 智谱AI 材料提供本协议的副本;
16
+ (B) 在相关网站、用户界面、博客文章、关于页面或产品文档上突出显示 “Built with CodeGeeX4”。
17
+ 如果您使用 THUDM / 智谱AI的 CodeGeeX4 开源模型的材料来创建、训练、微调或以其他方式改进已分发或可用的 AI 模型,您还应在任何此类 AI 模型名称的开头添加 “CodeGeeX4”。
18
+
19
+ 3. 限制
20
+
21
+ 您不得出于任何军事或非法目的使用、复制、修改、合并、发布、分发、复制或创建本软件的全部或部分衍生作品。
22
+ 您不得利用本软件从事任何危害国家安全和国家统一,危害社会公共利益及公序良俗,侵犯他人商业秘密、知识产权、名誉权、肖像权、财产权等权益的行为。
23
+ 您在使用中应遵循使用地所适用的法律法规政策、道德规范等要求。
24
+
25
+ 4. 免责声明
26
+
27
+ 本软件“按原样”提供,不提供任何明示或暗示的保证,包括但不限于对适销性、特定用途的适用性和非侵权性的保证。
28
+ 在任何情况下,作者或版权持有人均不对任何索赔、损害或其他责任负责,无论是在合同诉讼、侵权行为还是其他方面,由软件或软件的使用或其他交易引起、由软件引起或与之相关
29
+ 软件。
30
+
31
+ 5. 责任限制
32
+
33
+ 除适用法律禁止的范围外,在任何情况下且根据任何法律理论,无论是基于侵权行为、疏忽、合同、责任或其他原因,任何许可方均不对您承担任何直接、间接、特殊、偶然、示范性、
34
+ 或间接损害,或任何其他商业损失,即使许可人已被告知此类损害的可能性。
35
+
36
+ 6. 争议解决
37
+
38
+ 本许可受中华人民共和国法律管辖并按其解释。 因本许可引起的或与本许可有关的任何争议应提交北京市海淀区人民法院。
39
+ 请注意,许可证可能会更新到更全面的版本。 有关许可和版权的任何问题,请通过 license@zhipuai.cn 与我们联系。
40
+ 1. Definitions
41
+ “Licensor” means the CodeGeeX Team that distributes its Software.
42
+ “Software” means the CodeGeeX4 model parameters made available under this license.
43
+ 2. License
44
+ Under the terms and conditions of this license, the Licensor hereby grants you a non-exclusive, worldwide, non-transferable, non-sublicensable, revocable, royalty-free copyright license.
45
+ This license allows you to use all open source models in this repository for free for academic research. For users who wish to use the models for commercial purposes, please do so [here](https://open.bigmodel.cn/mla/form)
46
+ Complete registration. Registered users are free to use this model for commercial activities, but must comply with all terms and conditions of this license.
47
+ The copyright notice and this license notice shall be included in all copies or substantial portions of the Software.
48
+ If you distribute or provide THUDM / Zhipu AI materials on the CodeGeeX4 open source model (or any derivative works thereof), or products or services that use any materials therein (including all open source models of the CodeGeeX4 series), you should:
49
+ (A) Provide a copy of this Agreement with any such THUDM/Zhipu AI Materials;
50
+ (B) Prominently display "Built with CodeGeeX4" on the relevant website, user interface, blog post, related page or product documentation.
51
+ If you use materials from THUDM/Zhipu AI's CodeGeeX4 model to create, train, operate, or otherwise improve assigned or available AI models, you should also add "CodeGeeX4" to the beginning of any such AI model name.
52
+ 3. Restrictions
53
+ You are not allowed to use, copy, modify, merge, publish, distribute, copy or create all or part of the derivative works of this software for any military or illegal purposes.
54
+ You are not allowed to use this software to engage in any behavior that endangers national security and unity, endangers social public interests and public order, infringes on the rights and interests of others such as trade secrets, intellectual property rights, reputation rights, portrait rights, and property rights.
55
+ You should comply with the applicable laws, regulations, policies, ethical standards, and other requirements in the place of use during use.
56
+ 4. Disclaimer
57
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE
58
+ WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
59
+ COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR
60
+ OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
61
+ 5. Limitation of Liability
62
+ EXCEPT TO THE EXTENT PROHIBITED BY APPLICABLE LAW, IN NO EVENT AND UNDER NO LEGAL THEORY, WHETHER BASED IN TORT,
63
+ NEGLIGENCE, CONTRACT, LIABILITY, OR OTHERWISE WILL ANY LICENSOR BE LIABLE TO YOU FOR ANY DIRECT, INDIRECT, SPECIAL,
64
+ INCIDENTAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES, OR ANY OTHER COMMERCIAL LOSSES, EVEN IF THE LICENSOR HAS BEEN ADVISED
65
+ OF THE POSSIBILITY OF SUCH DAMAGES.
66
+ 6. Dispute Resolution
67
+ This license shall be governed and construed in accordance with the laws of People’s Republic of China. Any dispute
68
+ arising from or in connection with this License shall be submitted to Haidian District People's Court in Beijing.
69
+ Note that the license is subject to update to a more comprehensive version. For any questions related to the license and
70
+ copyright, please contact us at license@zhipuai.cn.
README.md ADDED
@@ -0,0 +1,102 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ ---
2
+ license: other
3
+ license_name: codegeex4
4
+ license_link: https://huggingface.co/THUDM/codegeex4-all-9b/blob/main/LICENSE
5
+ language:
6
+ - zh
7
+ - en
8
+ tags:
9
+ - glm
10
+ - codegeex
11
+ - thudm
12
+ inference: false
13
+ pipeline_tag: text-generation
14
+ ---
15
+
16
+ # CodeGeeX4: Open Multilingual Code Generation Model
17
+
18
+ <center>
19
+ <img src="https://raw.githubusercontent.com/THUDM/CodeGeeX4/main/resources/logo.jpeg" alt="CodeGeeX4">
20
+ </center>
21
+
22
+ [中文](./README_zh.md)
23
+
24
+ [GitHub](https://github.com/THUDM/CodeGeeX4)
25
+
26
+ We introduce CodeGeeX4-ALL-9B, the open-source version of the latest CodeGeeX4 model series. It is a multilingual code generation model continually trained on the [GLM-4-9B](https://github.com/THUDM/GLM-4), significantly enhancing its code generation capabilities. Using a single CodeGeeX4-ALL-9B model, it can support comprehensive functions such as code completion and generation, code interpreter, web search, function call, repository-level code Q&A, covering various scenarios of software development. CodeGeeX4-ALL-9B has achieved highly competitive performance on public benchmarks, such as [BigCodeBench](https://huggingface.co/spaces/bigcode/bigcodebench-leaderboard) and [NaturalCodeBench](https://github.com/THUDM/NaturalCodeBench). It is currently the most powerful code generation model with less than 10B parameters, even surpassing much larger general-purpose models, achieving the best balance in terms of inference speed and model performance.
27
+
28
+ ## Get Started
29
+
30
+ Use `4.39.0<=transformers<=4.40.2` to quickly launch [codegeex4-all-9b](https://huggingface.co/THUDM/codegeex4-all-9b):
31
+
32
+ ```python
33
+ import torch
34
+ from transformers import AutoTokenizer, AutoModelForCausalLM
35
+
36
+ device = "cuda" if torch.cuda.is_available() else "cpu"
37
+ tokenizer = AutoTokenizer.from_pretrained("THUDM/codegeex4-all-9b", trust_remote_code=True)
38
+ model = AutoModelForCausalLM.from_pretrained(
39
+ "THUDM/codegeex4-all-9b",
40
+ torch_dtype=torch.bfloat16,
41
+ low_cpu_mem_usage=True,
42
+ trust_remote_code=True
43
+ ).to(device).eval()
44
+ inputs = tokenizer.apply_chat_template([{"role": "user", "content": "write a quick sort"}], add_generation_prompt=True, tokenize=True, return_tensors="pt", return_dict=True ).to(device)
45
+ with torch.no_grad():
46
+ outputs = model.generate(**inputs, max_length=256)
47
+ outputs = outputs[:, inputs['input_ids'].shape[1]:]
48
+ print(tokenizer.decode(outputs[0], skip_special_tokens=True))
49
+ ```
50
+
51
+ If you want to build the **chat** prompt manually, please make sure it follows the following format:
52
+ ```
53
+ f"<|system|>\n{system_prompt}\n<|user|>\n{prompt}\n<|assistant|>\n"
54
+ ```
55
+ Default system_prompt:
56
+ ```
57
+ 你是一位智能编程助手,你叫CodeGeeX。你会为用户回答关于编程、代码、计算机方面的任何问题,并提供格式规范、可以执行、准确安全的代码,并在必要时提供详细的解释。
58
+ ```
59
+ The English version:
60
+ ```
61
+ You are an intelligent programming assistant named CodeGeeX. You will answer any questions users have about programming, coding, and computers, and provide code that is formatted correctly.
62
+ ```
63
+
64
+ For **infilling** ability, please use (without system prompt):
65
+ ```
66
+ f"<|user|>\n<|code_suffix|>{suffix}<|code_prefix|>{prefix}<|code_middle|><|assistant|>\n"
67
+ ```
68
+ Additional infos (like file path, programming language, mode) can be added. Example:
69
+ ```
70
+ <|user|>
71
+ ###PATH:src/example.py
72
+ ###LANGUAGE:Python
73
+ ###MODE:BLOCK
74
+ <|code_suffix|>{suffix}<|code_prefix|>{prefix}<|code_middle|><|assistant|>
75
+ ```
76
+
77
+ ## Evaluation
78
+
79
+ | **Model** | **Seq Length** | **HumanEval** | **MBPP** | **NCB** | **LCB** | **HumanEvalFIM** | **CRUXEval-O** |
80
+ |-----------------------------|----------------|---------------|----------|---------|---------|------------------|----------------|
81
+ | Llama3-70B-intruct | 8K | 77.4 | 82.3 | 37.0 | 27.4 | - | - |
82
+ | DeepSeek Coder 33B Instruct | 16K | 81.1 | 80.4 | 39.3 | 29.3 | 78.2 | 49.9 |
83
+ | Codestral-22B | 32K | 81.1 | 78.2 | 46.0 | 35.3 | 91.6 | 51.3 |
84
+ | CodeGeeX4-All-9B | 128K | 82.3 | 75.7 | 40.4 | 28.5 | 85.0 | 47.1 |
85
+
86
+ ## License
87
+
88
+ The model weights are licensed under the following [License](./LICENSE).
89
+
90
+ ## Citation
91
+
92
+ If you find our work helpful, please feel free to cite the following paper:
93
+
94
+ ```
95
+ @inproceedings{zheng2023codegeex,
96
+ title={CodeGeeX: A Pre-Trained Model for Code Generation with Multilingual Benchmarking on HumanEval-X},
97
+ author={Qinkai Zheng and Xiao Xia and Xu Zou and Yuxiao Dong and Shan Wang and Yufei Xue and Zihan Wang and Lei Shen and Andi Wang and Yang Li and Teng Su and Zhilin Yang and Jie Tang},
98
+ booktitle={Proceedings of the 29th ACM SIGKDD Conference on Knowledge Discovery and Data Mining},
99
+ pages={5673--5684},
100
+ year={2023}
101
+ }
102
+ ```
README_zh.md ADDED
@@ -0,0 +1,82 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # CodeGeeX4: 开源多语言代码生成模型
2
+
3
+ <center>
4
+ <img src="https://raw.githubusercontent.com/THUDM/CodeGeeX4/main/resources/logo.jpeg" alt="CodeGeeX4">
5
+ </center>
6
+
7
+ [CodeGeeX4 GitHub](https://github.com/THUDM/CodeGeeX4)
8
+
9
+ 我们推出了 CodeGeeX4-ALL-9B,这是最新的 CodeGeeX4 系列模型的开源版本。该模型是在 [GLM-4-9B](https://github.com/THUDM/GLM-4) 基础上持续训练的多语言代码生成模型,显著提升了代码生成能力。使用单个 CodeGeeX4-ALL-9B 模型,可以支持代码补全与生成、代码解释、联网搜索、函数调用、仓库级代码问答等多种功能,覆盖了软件开发的各个场景。CodeGeeX4-ALL-9B 在 [BigCodeBench](https://huggingface.co/spaces/bigcode/bigcodebench-leaderboard) 和 [NaturalCodeBench](https://github.com/THUDM/NaturalCodeBench) 等公开基准测试中取得了极具竞争力的表现。它是目前参数量少于 100 亿的最强代码生成模型,甚至超越了更大的通用模型,在推理速度和模型性能方面达到了最佳平衡。
10
+
11
+
12
+ ## 快速开始
13
+
14
+ 请使用 `4.39.0<=transformers<=4.40.2` 部署:
15
+
16
+ ```python
17
+ from transformers import AutoTokenizer, AutoModelForCausalLM
18
+
19
+ device = "cuda" if torch.cuda.is_available() else "cpu"
20
+ tokenizer = AutoTokenizer.from_pretrained("THUDM/codegeex4-all-9b", trust_remote_code=True)
21
+ model = AutoModelForCausalLM.from_pretrained(
22
+ "THUDM/codegeex4-all-9b",
23
+ torch_dtype=torch.bfloat16,
24
+ low_cpu_mem_usage=True,
25
+ trust_remote_code=True
26
+ ).to(device).eval()
27
+ inputs = tokenizer.apply_chat_template([{"role": "user", "content": "write a quick sort"}], add_generation_prompt=True, tokenize=True, return_tensors="pt", return_dict=True).to(device)
28
+ with torch.no_grad():
29
+ outputs = model.generate(**inputs, max_length=256)
30
+ outputs = outputs[:, inputs['input_ids'].shape[1]:]
31
+ print(tokenizer.decode(outputs[0], skip_special_tokens=True))
32
+ ```
33
+
34
+ 如果希望手动拼接**聊天**prompt,请使用以下格式:
35
+ ```
36
+ f"<|system|>\n{system_prompt}\n<|user|>\n{prompt}\n<|assistant|>\n"
37
+ ```
38
+ 默认系统指令:
39
+ ```
40
+ 你是一位智能编程助手,你叫CodeGeeX。你会为用户回答关于编程、代码、计算机方面的任何问题,并提供格式规范、可以执行、准确安全的代码,并在必要时提供详细的解释。
41
+ ```
42
+
43
+ 使用**上下文补全能力**能力,请使用以下格式(不需要系统指令):
44
+ ```
45
+ f"<|user|>\n<|code_suffix|>{suffix}<|code_prefix|>{prefix}<|code_middle|><|assistant|>\n"
46
+ ```
47
+ 可以添加额外信息(如文件名,编程语言,模式等),示例:
48
+ ```
49
+ <|user|>
50
+ ###PATH:src/example.py
51
+ ###LANGUAGE:Python
52
+ ###MODE:BLOCK
53
+ <|code_suffix|>{suffix}<|code_prefix|>{prefix}<|code_middle|><|assistant|>
54
+ ```
55
+
56
+ ## 评测指标
57
+
58
+ | **模型** | **序列长度** | **HumanEval** | **MBPP** | **NCB** | **LCB** | **HumanEvalFIM** | **CRUXEval-O** |
59
+ |-----------------------------|----------------|---------------|----------|---------|---------|------------------|----------------|
60
+ | Llama3-70B-intruct | 8K | 77.4 | 82.3 | 37.0 | 27.4 | - | - |
61
+ | DeepSeek Coder 33B Instruct | 16K | 81.1 | 80.4 | 39.3 | 29.3 | 78.2 | 49.9 |
62
+ | Codestral-22B | 32K | 81.1 | 78.2 | 46.0 | 35.3 | 91.6 | 51.3 |
63
+ | CodeGeeX4-All-9B | 128K | 82.3 | 75.7 | 40.4 | 28.5 | 85.0 | 47.1 |
64
+
65
+ ## License
66
+
67
+ CodeGeeX4-ALL-9B 模型的权重的使用则需要遵循 [License](./LICENSE).
68
+
69
+
70
+ ## 引用
71
+
72
+ 如果您觉得我们的工作对您有帮助,欢迎引用以下论文:
73
+
74
+ ```
75
+ @inproceedings{zheng2023codegeex,
76
+ title={CodeGeeX: A Pre-Trained Model for Code Generation with Multilingual Benchmarking on HumanEval-X},
77
+ author={Qinkai Zheng and Xiao Xia and Xu Zou and Yuxiao Dong and Shan Wang and Yufei Xue and Zihan Wang and Lei Shen and Andi Wang and Yang Li and Teng Su and Zhilin Yang and Jie Tang},
78
+ booktitle={Proceedings of the 29th ACM SIGKDD Conference on Knowledge Discovery and Data Mining},
79
+ pages={5673--5684},
80
+ year={2023}
81
+ }
82
+ ```
config.json ADDED
@@ -0,0 +1,50 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "_name_or_path": "THUDM/codegeex4-all-9b",
3
+ "add_bias_linear": false,
4
+ "add_qkv_bias": true,
5
+ "apply_query_key_layer_scaling": true,
6
+ "apply_residual_connection_post_layernorm": false,
7
+ "architectures": [
8
+ "ChatGLMModel"
9
+ ],
10
+ "attention_dropout": 0.0,
11
+ "attention_softmax_in_fp32": true,
12
+ "auto_map": {
13
+ "AutoConfig": "configuration_chatglm.ChatGLMConfig",
14
+ "AutoModel": "modeling_chatglm.ChatGLMForConditionalGeneration",
15
+ "AutoModelForCausalLM": "modeling_chatglm.ChatGLMForConditionalGeneration",
16
+ "AutoModelForSeq2SeqLM": "modeling_chatglm.ChatGLMForConditionalGeneration",
17
+ "AutoModelForSequenceClassification": "modeling_chatglm.ChatGLMForSequenceClassification"
18
+ },
19
+ "bias_dropout_fusion": true,
20
+ "classifier_dropout": null,
21
+ "eos_token_id": [
22
+ 151329,
23
+ 151336,
24
+ 151338
25
+ ],
26
+ "ffn_hidden_size": 13696,
27
+ "fp32_residual_connection": false,
28
+ "hidden_dropout": 0.0,
29
+ "hidden_size": 4096,
30
+ "kv_channels": 128,
31
+ "layernorm_epsilon": 1e-05,
32
+ "model_type": "chatglm",
33
+ "multi_query_attention": true,
34
+ "multi_query_group_num": 2,
35
+ "num_attention_heads": 32,
36
+ "num_hidden_layers": 40,
37
+ "num_layers": 40,
38
+ "original_rope": true,
39
+ "pad_token_id": 151329,
40
+ "padded_vocab_size": 151552,
41
+ "post_layer_norm": true,
42
+ "rmsnorm": true,
43
+ "rope_ratio": 500,
44
+ "seq_length": 131072,
45
+ "tie_word_embeddings": false,
46
+ "torch_dtype": "bfloat16",
47
+ "transformers_version": "4.40.2",
48
+ "use_cache": true,
49
+ "vocab_size": 151552
50
+ }
configuration.json ADDED
@@ -0,0 +1 @@
 
 
1
+ {"framework":"Pytorch","task":"text-generation"}
configuration_chatglm.py ADDED
@@ -0,0 +1,58 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from transformers import PretrainedConfig
2
+
3
+
4
+ class ChatGLMConfig(PretrainedConfig):
5
+ model_type = "chatglm"
6
+
7
+ def __init__(
8
+ self,
9
+ num_layers=28,
10
+ padded_vocab_size=65024,
11
+ hidden_size=4096,
12
+ ffn_hidden_size=13696,
13
+ kv_channels=128,
14
+ num_attention_heads=32,
15
+ seq_length=2048,
16
+ hidden_dropout=0.0,
17
+ classifier_dropout=None,
18
+ attention_dropout=0.0,
19
+ layernorm_epsilon=1e-5,
20
+ rmsnorm=True,
21
+ apply_residual_connection_post_layernorm=False,
22
+ post_layer_norm=True,
23
+ add_bias_linear=False,
24
+ add_qkv_bias=False,
25
+ bias_dropout_fusion=True,
26
+ multi_query_attention=False,
27
+ multi_query_group_num=1,
28
+ rope_ratio=1,
29
+ apply_query_key_layer_scaling=True,
30
+ attention_softmax_in_fp32=True,
31
+ fp32_residual_connection=False,
32
+ **kwargs
33
+ ):
34
+ self.num_layers = num_layers
35
+ self.vocab_size = padded_vocab_size
36
+ self.padded_vocab_size = padded_vocab_size
37
+ self.hidden_size = hidden_size
38
+ self.ffn_hidden_size = ffn_hidden_size
39
+ self.kv_channels = kv_channels
40
+ self.num_attention_heads = num_attention_heads
41
+ self.seq_length = seq_length
42
+ self.hidden_dropout = hidden_dropout
43
+ self.classifier_dropout = classifier_dropout
44
+ self.attention_dropout = attention_dropout
45
+ self.layernorm_epsilon = layernorm_epsilon
46
+ self.rmsnorm = rmsnorm
47
+ self.apply_residual_connection_post_layernorm = apply_residual_connection_post_layernorm
48
+ self.post_layer_norm = post_layer_norm
49
+ self.add_bias_linear = add_bias_linear
50
+ self.add_qkv_bias = add_qkv_bias
51
+ self.bias_dropout_fusion = bias_dropout_fusion
52
+ self.multi_query_attention = multi_query_attention
53
+ self.multi_query_group_num = multi_query_group_num
54
+ self.rope_ratio = rope_ratio
55
+ self.apply_query_key_layer_scaling = apply_query_key_layer_scaling
56
+ self.attention_softmax_in_fp32 = attention_softmax_in_fp32
57
+ self.fp32_residual_connection = fp32_residual_connection
58
+ super().__init__(**kwargs)
generation_config.json ADDED
@@ -0,0 +1,10 @@
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "_from_model_config": true,
3
+ "eos_token_id": [
4
+ 151329,
5
+ 151336,
6
+ 151338
7
+ ],
8
+ "pad_token_id": 151329,
9
+ "transformers_version": "4.40.2"
10
+ }
model-00001-of-00004.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:b5978eae170af0f148a02b359b8f0a530c3343263797cde7f2eaeeb01cee7150
3
+ size 4984147224
model-00002-of-00004.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:96cc587d630d8c83a16a013c7df1a7ee93393c6245ea67e670433026a00e7153
3
+ size 4895071360
model-00003-of-00004.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:6f7790d856338c832f48be28fa7d8a8d25a8c3c6d54c00be60e66ca25dcaf7c0
3
+ size 4895071384
model-00004-of-00004.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:a11a9b0d6f4ed05a97b74940a6d15a523c6ff8f7f697eb97999d8f03e51eb37c
3
+ size 4025651256
model.safetensors.index.json ADDED
@@ -0,0 +1,291 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "metadata": {
3
+ "total_size": 18799902784
4
+ },
5
+ "weight_map": {
6
+ "transformer.embedding.word_embeddings.weight": "model-00001-of-00004.safetensors",
7
+ "transformer.encoder.final_layernorm.weight": "model-00004-of-00004.safetensors",
8
+ "transformer.encoder.layers.0.input_layernorm.weight": "model-00001-of-00004.safetensors",
9
+ "transformer.encoder.layers.0.mlp.dense_4h_to_h.weight": "model-00001-of-00004.safetensors",
10
+ "transformer.encoder.layers.0.mlp.dense_h_to_4h.weight": "model-00001-of-00004.safetensors",
11
+ "transformer.encoder.layers.0.post_attention_layernorm.weight": "model-00001-of-00004.safetensors",
12
+ "transformer.encoder.layers.0.self_attention.dense.weight": "model-00001-of-00004.safetensors",
13
+ "transformer.encoder.layers.0.self_attention.query_key_value.bias": "model-00001-of-00004.safetensors",
14
+ "transformer.encoder.layers.0.self_attention.query_key_value.weight": "model-00001-of-00004.safetensors",
15
+ "transformer.encoder.layers.1.input_layernorm.weight": "model-00001-of-00004.safetensors",
16
+ "transformer.encoder.layers.1.mlp.dense_4h_to_h.weight": "model-00001-of-00004.safetensors",
17
+ "transformer.encoder.layers.1.mlp.dense_h_to_4h.weight": "model-00001-of-00004.safetensors",
18
+ "transformer.encoder.layers.1.post_attention_layernorm.weight": "model-00001-of-00004.safetensors",
19
+ "transformer.encoder.layers.1.self_attention.dense.weight": "model-00001-of-00004.safetensors",
20
+ "transformer.encoder.layers.1.self_attention.query_key_value.bias": "model-00001-of-00004.safetensors",
21
+ "transformer.encoder.layers.1.self_attention.query_key_value.weight": "model-00001-of-00004.safetensors",
22
+ "transformer.encoder.layers.10.input_layernorm.weight": "model-00002-of-00004.safetensors",
23
+ "transformer.encoder.layers.10.mlp.dense_4h_to_h.weight": "model-00002-of-00004.safetensors",
24
+ "transformer.encoder.layers.10.mlp.dense_h_to_4h.weight": "model-00002-of-00004.safetensors",
25
+ "transformer.encoder.layers.10.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
26
+ "transformer.encoder.layers.10.self_attention.dense.weight": "model-00002-of-00004.safetensors",
27
+ "transformer.encoder.layers.10.self_attention.query_key_value.bias": "model-00002-of-00004.safetensors",
28
+ "transformer.encoder.layers.10.self_attention.query_key_value.weight": "model-00002-of-00004.safetensors",
29
+ "transformer.encoder.layers.11.input_layernorm.weight": "model-00002-of-00004.safetensors",
30
+ "transformer.encoder.layers.11.mlp.dense_4h_to_h.weight": "model-00002-of-00004.safetensors",
31
+ "transformer.encoder.layers.11.mlp.dense_h_to_4h.weight": "model-00002-of-00004.safetensors",
32
+ "transformer.encoder.layers.11.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
33
+ "transformer.encoder.layers.11.self_attention.dense.weight": "model-00002-of-00004.safetensors",
34
+ "transformer.encoder.layers.11.self_attention.query_key_value.bias": "model-00002-of-00004.safetensors",
35
+ "transformer.encoder.layers.11.self_attention.query_key_value.weight": "model-00002-of-00004.safetensors",
36
+ "transformer.encoder.layers.12.input_layernorm.weight": "model-00002-of-00004.safetensors",
37
+ "transformer.encoder.layers.12.mlp.dense_4h_to_h.weight": "model-00002-of-00004.safetensors",
38
+ "transformer.encoder.layers.12.mlp.dense_h_to_4h.weight": "model-00002-of-00004.safetensors",
39
+ "transformer.encoder.layers.12.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
40
+ "transformer.encoder.layers.12.self_attention.dense.weight": "model-00002-of-00004.safetensors",
41
+ "transformer.encoder.layers.12.self_attention.query_key_value.bias": "model-00002-of-00004.safetensors",
42
+ "transformer.encoder.layers.12.self_attention.query_key_value.weight": "model-00002-of-00004.safetensors",
43
+ "transformer.encoder.layers.13.input_layernorm.weight": "model-00002-of-00004.safetensors",
44
+ "transformer.encoder.layers.13.mlp.dense_4h_to_h.weight": "model-00002-of-00004.safetensors",
45
+ "transformer.encoder.layers.13.mlp.dense_h_to_4h.weight": "model-00002-of-00004.safetensors",
46
+ "transformer.encoder.layers.13.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
47
+ "transformer.encoder.layers.13.self_attention.dense.weight": "model-00002-of-00004.safetensors",
48
+ "transformer.encoder.layers.13.self_attention.query_key_value.bias": "model-00002-of-00004.safetensors",
49
+ "transformer.encoder.layers.13.self_attention.query_key_value.weight": "model-00002-of-00004.safetensors",
50
+ "transformer.encoder.layers.14.input_layernorm.weight": "model-00002-of-00004.safetensors",
51
+ "transformer.encoder.layers.14.mlp.dense_4h_to_h.weight": "model-00002-of-00004.safetensors",
52
+ "transformer.encoder.layers.14.mlp.dense_h_to_4h.weight": "model-00002-of-00004.safetensors",
53
+ "transformer.encoder.layers.14.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
54
+ "transformer.encoder.layers.14.self_attention.dense.weight": "model-00002-of-00004.safetensors",
55
+ "transformer.encoder.layers.14.self_attention.query_key_value.bias": "model-00002-of-00004.safetensors",
56
+ "transformer.encoder.layers.14.self_attention.query_key_value.weight": "model-00002-of-00004.safetensors",
57
+ "transformer.encoder.layers.15.input_layernorm.weight": "model-00002-of-00004.safetensors",
58
+ "transformer.encoder.layers.15.mlp.dense_4h_to_h.weight": "model-00002-of-00004.safetensors",
59
+ "transformer.encoder.layers.15.mlp.dense_h_to_4h.weight": "model-00002-of-00004.safetensors",
60
+ "transformer.encoder.layers.15.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
61
+ "transformer.encoder.layers.15.self_attention.dense.weight": "model-00002-of-00004.safetensors",
62
+ "transformer.encoder.layers.15.self_attention.query_key_value.bias": "model-00002-of-00004.safetensors",
63
+ "transformer.encoder.layers.15.self_attention.query_key_value.weight": "model-00002-of-00004.safetensors",
64
+ "transformer.encoder.layers.16.input_layernorm.weight": "model-00002-of-00004.safetensors",
65
+ "transformer.encoder.layers.16.mlp.dense_4h_to_h.weight": "model-00002-of-00004.safetensors",
66
+ "transformer.encoder.layers.16.mlp.dense_h_to_4h.weight": "model-00002-of-00004.safetensors",
67
+ "transformer.encoder.layers.16.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
68
+ "transformer.encoder.layers.16.self_attention.dense.weight": "model-00002-of-00004.safetensors",
69
+ "transformer.encoder.layers.16.self_attention.query_key_value.bias": "model-00002-of-00004.safetensors",
70
+ "transformer.encoder.layers.16.self_attention.query_key_value.weight": "model-00002-of-00004.safetensors",
71
+ "transformer.encoder.layers.17.input_layernorm.weight": "model-00002-of-00004.safetensors",
72
+ "transformer.encoder.layers.17.mlp.dense_4h_to_h.weight": "model-00002-of-00004.safetensors",
73
+ "transformer.encoder.layers.17.mlp.dense_h_to_4h.weight": "model-00002-of-00004.safetensors",
74
+ "transformer.encoder.layers.17.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
75
+ "transformer.encoder.layers.17.self_attention.dense.weight": "model-00002-of-00004.safetensors",
76
+ "transformer.encoder.layers.17.self_attention.query_key_value.bias": "model-00002-of-00004.safetensors",
77
+ "transformer.encoder.layers.17.self_attention.query_key_value.weight": "model-00002-of-00004.safetensors",
78
+ "transformer.encoder.layers.18.input_layernorm.weight": "model-00002-of-00004.safetensors",
79
+ "transformer.encoder.layers.18.mlp.dense_4h_to_h.weight": "model-00002-of-00004.safetensors",
80
+ "transformer.encoder.layers.18.mlp.dense_h_to_4h.weight": "model-00002-of-00004.safetensors",
81
+ "transformer.encoder.layers.18.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
82
+ "transformer.encoder.layers.18.self_attention.dense.weight": "model-00002-of-00004.safetensors",
83
+ "transformer.encoder.layers.18.self_attention.query_key_value.bias": "model-00002-of-00004.safetensors",
84
+ "transformer.encoder.layers.18.self_attention.query_key_value.weight": "model-00002-of-00004.safetensors",
85
+ "transformer.encoder.layers.19.input_layernorm.weight": "model-00002-of-00004.safetensors",
86
+ "transformer.encoder.layers.19.mlp.dense_4h_to_h.weight": "model-00002-of-00004.safetensors",
87
+ "transformer.encoder.layers.19.mlp.dense_h_to_4h.weight": "model-00002-of-00004.safetensors",
88
+ "transformer.encoder.layers.19.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
89
+ "transformer.encoder.layers.19.self_attention.dense.weight": "model-00002-of-00004.safetensors",
90
+ "transformer.encoder.layers.19.self_attention.query_key_value.bias": "model-00002-of-00004.safetensors",
91
+ "transformer.encoder.layers.19.self_attention.query_key_value.weight": "model-00002-of-00004.safetensors",
92
+ "transformer.encoder.layers.2.input_layernorm.weight": "model-00001-of-00004.safetensors",
93
+ "transformer.encoder.layers.2.mlp.dense_4h_to_h.weight": "model-00001-of-00004.safetensors",
94
+ "transformer.encoder.layers.2.mlp.dense_h_to_4h.weight": "model-00001-of-00004.safetensors",
95
+ "transformer.encoder.layers.2.post_attention_layernorm.weight": "model-00001-of-00004.safetensors",
96
+ "transformer.encoder.layers.2.self_attention.dense.weight": "model-00001-of-00004.safetensors",
97
+ "transformer.encoder.layers.2.self_attention.query_key_value.bias": "model-00001-of-00004.safetensors",
98
+ "transformer.encoder.layers.2.self_attention.query_key_value.weight": "model-00001-of-00004.safetensors",
99
+ "transformer.encoder.layers.20.input_layernorm.weight": "model-00002-of-00004.safetensors",
100
+ "transformer.encoder.layers.20.mlp.dense_4h_to_h.weight": "model-00002-of-00004.safetensors",
101
+ "transformer.encoder.layers.20.mlp.dense_h_to_4h.weight": "model-00002-of-00004.safetensors",
102
+ "transformer.encoder.layers.20.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
103
+ "transformer.encoder.layers.20.self_attention.dense.weight": "model-00002-of-00004.safetensors",
104
+ "transformer.encoder.layers.20.self_attention.query_key_value.bias": "model-00002-of-00004.safetensors",
105
+ "transformer.encoder.layers.20.self_attention.query_key_value.weight": "model-00002-of-00004.safetensors",
106
+ "transformer.encoder.layers.21.input_layernorm.weight": "model-00002-of-00004.safetensors",
107
+ "transformer.encoder.layers.21.mlp.dense_4h_to_h.weight": "model-00003-of-00004.safetensors",
108
+ "transformer.encoder.layers.21.mlp.dense_h_to_4h.weight": "model-00003-of-00004.safetensors",
109
+ "transformer.encoder.layers.21.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
110
+ "transformer.encoder.layers.21.self_attention.dense.weight": "model-00002-of-00004.safetensors",
111
+ "transformer.encoder.layers.21.self_attention.query_key_value.bias": "model-00002-of-00004.safetensors",
112
+ "transformer.encoder.layers.21.self_attention.query_key_value.weight": "model-00002-of-00004.safetensors",
113
+ "transformer.encoder.layers.22.input_layernorm.weight": "model-00003-of-00004.safetensors",
114
+ "transformer.encoder.layers.22.mlp.dense_4h_to_h.weight": "model-00003-of-00004.safetensors",
115
+ "transformer.encoder.layers.22.mlp.dense_h_to_4h.weight": "model-00003-of-00004.safetensors",
116
+ "transformer.encoder.layers.22.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
117
+ "transformer.encoder.layers.22.self_attention.dense.weight": "model-00003-of-00004.safetensors",
118
+ "transformer.encoder.layers.22.self_attention.query_key_value.bias": "model-00003-of-00004.safetensors",
119
+ "transformer.encoder.layers.22.self_attention.query_key_value.weight": "model-00003-of-00004.safetensors",
120
+ "transformer.encoder.layers.23.input_layernorm.weight": "model-00003-of-00004.safetensors",
121
+ "transformer.encoder.layers.23.mlp.dense_4h_to_h.weight": "model-00003-of-00004.safetensors",
122
+ "transformer.encoder.layers.23.mlp.dense_h_to_4h.weight": "model-00003-of-00004.safetensors",
123
+ "transformer.encoder.layers.23.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
124
+ "transformer.encoder.layers.23.self_attention.dense.weight": "model-00003-of-00004.safetensors",
125
+ "transformer.encoder.layers.23.self_attention.query_key_value.bias": "model-00003-of-00004.safetensors",
126
+ "transformer.encoder.layers.23.self_attention.query_key_value.weight": "model-00003-of-00004.safetensors",
127
+ "transformer.encoder.layers.24.input_layernorm.weight": "model-00003-of-00004.safetensors",
128
+ "transformer.encoder.layers.24.mlp.dense_4h_to_h.weight": "model-00003-of-00004.safetensors",
129
+ "transformer.encoder.layers.24.mlp.dense_h_to_4h.weight": "model-00003-of-00004.safetensors",
130
+ "transformer.encoder.layers.24.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
131
+ "transformer.encoder.layers.24.self_attention.dense.weight": "model-00003-of-00004.safetensors",
132
+ "transformer.encoder.layers.24.self_attention.query_key_value.bias": "model-00003-of-00004.safetensors",
133
+ "transformer.encoder.layers.24.self_attention.query_key_value.weight": "model-00003-of-00004.safetensors",
134
+ "transformer.encoder.layers.25.input_layernorm.weight": "model-00003-of-00004.safetensors",
135
+ "transformer.encoder.layers.25.mlp.dense_4h_to_h.weight": "model-00003-of-00004.safetensors",
136
+ "transformer.encoder.layers.25.mlp.dense_h_to_4h.weight": "model-00003-of-00004.safetensors",
137
+ "transformer.encoder.layers.25.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
138
+ "transformer.encoder.layers.25.self_attention.dense.weight": "model-00003-of-00004.safetensors",
139
+ "transformer.encoder.layers.25.self_attention.query_key_value.bias": "model-00003-of-00004.safetensors",
140
+ "transformer.encoder.layers.25.self_attention.query_key_value.weight": "model-00003-of-00004.safetensors",
141
+ "transformer.encoder.layers.26.input_layernorm.weight": "model-00003-of-00004.safetensors",
142
+ "transformer.encoder.layers.26.mlp.dense_4h_to_h.weight": "model-00003-of-00004.safetensors",
143
+ "transformer.encoder.layers.26.mlp.dense_h_to_4h.weight": "model-00003-of-00004.safetensors",
144
+ "transformer.encoder.layers.26.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
145
+ "transformer.encoder.layers.26.self_attention.dense.weight": "model-00003-of-00004.safetensors",
146
+ "transformer.encoder.layers.26.self_attention.query_key_value.bias": "model-00003-of-00004.safetensors",
147
+ "transformer.encoder.layers.26.self_attention.query_key_value.weight": "model-00003-of-00004.safetensors",
148
+ "transformer.encoder.layers.27.input_layernorm.weight": "model-00003-of-00004.safetensors",
149
+ "transformer.encoder.layers.27.mlp.dense_4h_to_h.weight": "model-00003-of-00004.safetensors",
150
+ "transformer.encoder.layers.27.mlp.dense_h_to_4h.weight": "model-00003-of-00004.safetensors",
151
+ "transformer.encoder.layers.27.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
152
+ "transformer.encoder.layers.27.self_attention.dense.weight": "model-00003-of-00004.safetensors",
153
+ "transformer.encoder.layers.27.self_attention.query_key_value.bias": "model-00003-of-00004.safetensors",
154
+ "transformer.encoder.layers.27.self_attention.query_key_value.weight": "model-00003-of-00004.safetensors",
155
+ "transformer.encoder.layers.28.input_layernorm.weight": "model-00003-of-00004.safetensors",
156
+ "transformer.encoder.layers.28.mlp.dense_4h_to_h.weight": "model-00003-of-00004.safetensors",
157
+ "transformer.encoder.layers.28.mlp.dense_h_to_4h.weight": "model-00003-of-00004.safetensors",
158
+ "transformer.encoder.layers.28.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
159
+ "transformer.encoder.layers.28.self_attention.dense.weight": "model-00003-of-00004.safetensors",
160
+ "transformer.encoder.layers.28.self_attention.query_key_value.bias": "model-00003-of-00004.safetensors",
161
+ "transformer.encoder.layers.28.self_attention.query_key_value.weight": "model-00003-of-00004.safetensors",
162
+ "transformer.encoder.layers.29.input_layernorm.weight": "model-00003-of-00004.safetensors",
163
+ "transformer.encoder.layers.29.mlp.dense_4h_to_h.weight": "model-00003-of-00004.safetensors",
164
+ "transformer.encoder.layers.29.mlp.dense_h_to_4h.weight": "model-00003-of-00004.safetensors",
165
+ "transformer.encoder.layers.29.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
166
+ "transformer.encoder.layers.29.self_attention.dense.weight": "model-00003-of-00004.safetensors",
167
+ "transformer.encoder.layers.29.self_attention.query_key_value.bias": "model-00003-of-00004.safetensors",
168
+ "transformer.encoder.layers.29.self_attention.query_key_value.weight": "model-00003-of-00004.safetensors",
169
+ "transformer.encoder.layers.3.input_layernorm.weight": "model-00001-of-00004.safetensors",
170
+ "transformer.encoder.layers.3.mlp.dense_4h_to_h.weight": "model-00001-of-00004.safetensors",
171
+ "transformer.encoder.layers.3.mlp.dense_h_to_4h.weight": "model-00001-of-00004.safetensors",
172
+ "transformer.encoder.layers.3.post_attention_layernorm.weight": "model-00001-of-00004.safetensors",
173
+ "transformer.encoder.layers.3.self_attention.dense.weight": "model-00001-of-00004.safetensors",
174
+ "transformer.encoder.layers.3.self_attention.query_key_value.bias": "model-00001-of-00004.safetensors",
175
+ "transformer.encoder.layers.3.self_attention.query_key_value.weight": "model-00001-of-00004.safetensors",
176
+ "transformer.encoder.layers.30.input_layernorm.weight": "model-00003-of-00004.safetensors",
177
+ "transformer.encoder.layers.30.mlp.dense_4h_to_h.weight": "model-00003-of-00004.safetensors",
178
+ "transformer.encoder.layers.30.mlp.dense_h_to_4h.weight": "model-00003-of-00004.safetensors",
179
+ "transformer.encoder.layers.30.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
180
+ "transformer.encoder.layers.30.self_attention.dense.weight": "model-00003-of-00004.safetensors",
181
+ "transformer.encoder.layers.30.self_attention.query_key_value.bias": "model-00003-of-00004.safetensors",
182
+ "transformer.encoder.layers.30.self_attention.query_key_value.weight": "model-00003-of-00004.safetensors",
183
+ "transformer.encoder.layers.31.input_layernorm.weight": "model-00003-of-00004.safetensors",
184
+ "transformer.encoder.layers.31.mlp.dense_4h_to_h.weight": "model-00003-of-00004.safetensors",
185
+ "transformer.encoder.layers.31.mlp.dense_h_to_4h.weight": "model-00003-of-00004.safetensors",
186
+ "transformer.encoder.layers.31.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
187
+ "transformer.encoder.layers.31.self_attention.dense.weight": "model-00003-of-00004.safetensors",
188
+ "transformer.encoder.layers.31.self_attention.query_key_value.bias": "model-00003-of-00004.safetensors",
189
+ "transformer.encoder.layers.31.self_attention.query_key_value.weight": "model-00003-of-00004.safetensors",
190
+ "transformer.encoder.layers.32.input_layernorm.weight": "model-00003-of-00004.safetensors",
191
+ "transformer.encoder.layers.32.mlp.dense_4h_to_h.weight": "model-00003-of-00004.safetensors",
192
+ "transformer.encoder.layers.32.mlp.dense_h_to_4h.weight": "model-00003-of-00004.safetensors",
193
+ "transformer.encoder.layers.32.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
194
+ "transformer.encoder.layers.32.self_attention.dense.weight": "model-00003-of-00004.safetensors",
195
+ "transformer.encoder.layers.32.self_attention.query_key_value.bias": "model-00003-of-00004.safetensors",
196
+ "transformer.encoder.layers.32.self_attention.query_key_value.weight": "model-00003-of-00004.safetensors",
197
+ "transformer.encoder.layers.33.input_layernorm.weight": "model-00003-of-00004.safetensors",
198
+ "transformer.encoder.layers.33.mlp.dense_4h_to_h.weight": "model-00004-of-00004.safetensors",
199
+ "transformer.encoder.layers.33.mlp.dense_h_to_4h.weight": "model-00004-of-00004.safetensors",
200
+ "transformer.encoder.layers.33.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
201
+ "transformer.encoder.layers.33.self_attention.dense.weight": "model-00003-of-00004.safetensors",
202
+ "transformer.encoder.layers.33.self_attention.query_key_value.bias": "model-00003-of-00004.safetensors",
203
+ "transformer.encoder.layers.33.self_attention.query_key_value.weight": "model-00003-of-00004.safetensors",
204
+ "transformer.encoder.layers.34.input_layernorm.weight": "model-00004-of-00004.safetensors",
205
+ "transformer.encoder.layers.34.mlp.dense_4h_to_h.weight": "model-00004-of-00004.safetensors",
206
+ "transformer.encoder.layers.34.mlp.dense_h_to_4h.weight": "model-00004-of-00004.safetensors",
207
+ "transformer.encoder.layers.34.post_attention_layernorm.weight": "model-00004-of-00004.safetensors",
208
+ "transformer.encoder.layers.34.self_attention.dense.weight": "model-00004-of-00004.safetensors",
209
+ "transformer.encoder.layers.34.self_attention.query_key_value.bias": "model-00004-of-00004.safetensors",
210
+ "transformer.encoder.layers.34.self_attention.query_key_value.weight": "model-00004-of-00004.safetensors",
211
+ "transformer.encoder.layers.35.input_layernorm.weight": "model-00004-of-00004.safetensors",
212
+ "transformer.encoder.layers.35.mlp.dense_4h_to_h.weight": "model-00004-of-00004.safetensors",
213
+ "transformer.encoder.layers.35.mlp.dense_h_to_4h.weight": "model-00004-of-00004.safetensors",
214
+ "transformer.encoder.layers.35.post_attention_layernorm.weight": "model-00004-of-00004.safetensors",
215
+ "transformer.encoder.layers.35.self_attention.dense.weight": "model-00004-of-00004.safetensors",
216
+ "transformer.encoder.layers.35.self_attention.query_key_value.bias": "model-00004-of-00004.safetensors",
217
+ "transformer.encoder.layers.35.self_attention.query_key_value.weight": "model-00004-of-00004.safetensors",
218
+ "transformer.encoder.layers.36.input_layernorm.weight": "model-00004-of-00004.safetensors",
219
+ "transformer.encoder.layers.36.mlp.dense_4h_to_h.weight": "model-00004-of-00004.safetensors",
220
+ "transformer.encoder.layers.36.mlp.dense_h_to_4h.weight": "model-00004-of-00004.safetensors",
221
+ "transformer.encoder.layers.36.post_attention_layernorm.weight": "model-00004-of-00004.safetensors",
222
+ "transformer.encoder.layers.36.self_attention.dense.weight": "model-00004-of-00004.safetensors",
223
+ "transformer.encoder.layers.36.self_attention.query_key_value.bias": "model-00004-of-00004.safetensors",
224
+ "transformer.encoder.layers.36.self_attention.query_key_value.weight": "model-00004-of-00004.safetensors",
225
+ "transformer.encoder.layers.37.input_layernorm.weight": "model-00004-of-00004.safetensors",
226
+ "transformer.encoder.layers.37.mlp.dense_4h_to_h.weight": "model-00004-of-00004.safetensors",
227
+ "transformer.encoder.layers.37.mlp.dense_h_to_4h.weight": "model-00004-of-00004.safetensors",
228
+ "transformer.encoder.layers.37.post_attention_layernorm.weight": "model-00004-of-00004.safetensors",
229
+ "transformer.encoder.layers.37.self_attention.dense.weight": "model-00004-of-00004.safetensors",
230
+ "transformer.encoder.layers.37.self_attention.query_key_value.bias": "model-00004-of-00004.safetensors",
231
+ "transformer.encoder.layers.37.self_attention.query_key_value.weight": "model-00004-of-00004.safetensors",
232
+ "transformer.encoder.layers.38.input_layernorm.weight": "model-00004-of-00004.safetensors",
233
+ "transformer.encoder.layers.38.mlp.dense_4h_to_h.weight": "model-00004-of-00004.safetensors",
234
+ "transformer.encoder.layers.38.mlp.dense_h_to_4h.weight": "model-00004-of-00004.safetensors",
235
+ "transformer.encoder.layers.38.post_attention_layernorm.weight": "model-00004-of-00004.safetensors",
236
+ "transformer.encoder.layers.38.self_attention.dense.weight": "model-00004-of-00004.safetensors",
237
+ "transformer.encoder.layers.38.self_attention.query_key_value.bias": "model-00004-of-00004.safetensors",
238
+ "transformer.encoder.layers.38.self_attention.query_key_value.weight": "model-00004-of-00004.safetensors",
239
+ "transformer.encoder.layers.39.input_layernorm.weight": "model-00004-of-00004.safetensors",
240
+ "transformer.encoder.layers.39.mlp.dense_4h_to_h.weight": "model-00004-of-00004.safetensors",
241
+ "transformer.encoder.layers.39.mlp.dense_h_to_4h.weight": "model-00004-of-00004.safetensors",
242
+ "transformer.encoder.layers.39.post_attention_layernorm.weight": "model-00004-of-00004.safetensors",
243
+ "transformer.encoder.layers.39.self_attention.dense.weight": "model-00004-of-00004.safetensors",
244
+ "transformer.encoder.layers.39.self_attention.query_key_value.bias": "model-00004-of-00004.safetensors",
245
+ "transformer.encoder.layers.39.self_attention.query_key_value.weight": "model-00004-of-00004.safetensors",
246
+ "transformer.encoder.layers.4.input_layernorm.weight": "model-00001-of-00004.safetensors",
247
+ "transformer.encoder.layers.4.mlp.dense_4h_to_h.weight": "model-00001-of-00004.safetensors",
248
+ "transformer.encoder.layers.4.mlp.dense_h_to_4h.weight": "model-00001-of-00004.safetensors",
249
+ "transformer.encoder.layers.4.post_attention_layernorm.weight": "model-00001-of-00004.safetensors",
250
+ "transformer.encoder.layers.4.self_attention.dense.weight": "model-00001-of-00004.safetensors",
251
+ "transformer.encoder.layers.4.self_attention.query_key_value.bias": "model-00001-of-00004.safetensors",
252
+ "transformer.encoder.layers.4.self_attention.query_key_value.weight": "model-00001-of-00004.safetensors",
253
+ "transformer.encoder.layers.5.input_layernorm.weight": "model-00001-of-00004.safetensors",
254
+ "transformer.encoder.layers.5.mlp.dense_4h_to_h.weight": "model-00001-of-00004.safetensors",
255
+ "transformer.encoder.layers.5.mlp.dense_h_to_4h.weight": "model-00001-of-00004.safetensors",
256
+ "transformer.encoder.layers.5.post_attention_layernorm.weight": "model-00001-of-00004.safetensors",
257
+ "transformer.encoder.layers.5.self_attention.dense.weight": "model-00001-of-00004.safetensors",
258
+ "transformer.encoder.layers.5.self_attention.query_key_value.bias": "model-00001-of-00004.safetensors",
259
+ "transformer.encoder.layers.5.self_attention.query_key_value.weight": "model-00001-of-00004.safetensors",
260
+ "transformer.encoder.layers.6.input_layernorm.weight": "model-00001-of-00004.safetensors",
261
+ "transformer.encoder.layers.6.mlp.dense_4h_to_h.weight": "model-00001-of-00004.safetensors",
262
+ "transformer.encoder.layers.6.mlp.dense_h_to_4h.weight": "model-00001-of-00004.safetensors",
263
+ "transformer.encoder.layers.6.post_attention_layernorm.weight": "model-00001-of-00004.safetensors",
264
+ "transformer.encoder.layers.6.self_attention.dense.weight": "model-00001-of-00004.safetensors",
265
+ "transformer.encoder.layers.6.self_attention.query_key_value.bias": "model-00001-of-00004.safetensors",
266
+ "transformer.encoder.layers.6.self_attention.query_key_value.weight": "model-00001-of-00004.safetensors",
267
+ "transformer.encoder.layers.7.input_layernorm.weight": "model-00001-of-00004.safetensors",
268
+ "transformer.encoder.layers.7.mlp.dense_4h_to_h.weight": "model-00001-of-00004.safetensors",
269
+ "transformer.encoder.layers.7.mlp.dense_h_to_4h.weight": "model-00001-of-00004.safetensors",
270
+ "transformer.encoder.layers.7.post_attention_layernorm.weight": "model-00001-of-00004.safetensors",
271
+ "transformer.encoder.layers.7.self_attention.dense.weight": "model-00001-of-00004.safetensors",
272
+ "transformer.encoder.layers.7.self_attention.query_key_value.bias": "model-00001-of-00004.safetensors",
273
+ "transformer.encoder.layers.7.self_attention.query_key_value.weight": "model-00001-of-00004.safetensors",
274
+ "transformer.encoder.layers.8.input_layernorm.weight": "model-00001-of-00004.safetensors",
275
+ "transformer.encoder.layers.8.mlp.dense_4h_to_h.weight": "model-00001-of-00004.safetensors",
276
+ "transformer.encoder.layers.8.mlp.dense_h_to_4h.weight": "model-00001-of-00004.safetensors",
277
+ "transformer.encoder.layers.8.post_attention_layernorm.weight": "model-00001-of-00004.safetensors",
278
+ "transformer.encoder.layers.8.self_attention.dense.weight": "model-00001-of-00004.safetensors",
279
+ "transformer.encoder.layers.8.self_attention.query_key_value.bias": "model-00001-of-00004.safetensors",
280
+ "transformer.encoder.layers.8.self_attention.query_key_value.weight": "model-00001-of-00004.safetensors",
281
+ "transformer.encoder.layers.9.input_layernorm.weight": "model-00001-of-00004.safetensors",
282
+ "transformer.encoder.layers.9.mlp.dense_4h_to_h.weight": "model-00002-of-00004.safetensors",
283
+ "transformer.encoder.layers.9.mlp.dense_h_to_4h.weight": "model-00002-of-00004.safetensors",
284
+ "transformer.encoder.layers.9.post_attention_layernorm.weight": "model-00001-of-00004.safetensors",
285
+ "transformer.encoder.layers.9.self_attention.dense.weight": "model-00001-of-00004.safetensors",
286
+ "transformer.encoder.layers.9.self_attention.query_key_value.bias": "model-00001-of-00004.safetensors",
287
+ "transformer.encoder.layers.9.self_attention.query_key_value.weight": "model-00001-of-00004.safetensors",
288
+ "transformer.output_layer.weight": "model-00004-of-00004.safetensors",
289
+ "transformer.rotary_pos_emb.inv_freq": "model-00001-of-00004.safetensors"
290
+ }
291
+ }
modeling_chatglm.py ADDED
@@ -0,0 +1,1344 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """ PyTorch ChatGLM model. """
2
+ import json
3
+ import math
4
+ import copy
5
+ import warnings
6
+ import re
7
+ import sys
8
+
9
+ import torch
10
+ import torch.utils.checkpoint
11
+ import torch.nn.functional as F
12
+ from torch import nn
13
+ from torch.nn import CrossEntropyLoss, LayerNorm, MSELoss, BCEWithLogitsLoss
14
+ from torch.nn.utils import skip_init
15
+ from typing import Optional, Tuple, Union, List, Callable, Dict, Any
16
+ from copy import deepcopy
17
+
18
+ from transformers.modeling_outputs import (
19
+ BaseModelOutputWithPast,
20
+ CausalLMOutputWithPast,
21
+ SequenceClassifierOutputWithPast,
22
+ )
23
+ from transformers.modeling_utils import PreTrainedModel
24
+ from transformers.utils import logging, is_torch_npu_available
25
+ from transformers.generation.logits_process import LogitsProcessor
26
+ from transformers.generation.utils import LogitsProcessorList, StoppingCriteriaList, GenerationConfig, ModelOutput
27
+
28
+ from .configuration_chatglm import ChatGLMConfig
29
+
30
+ try:
31
+ from transformers.utils import is_flash_attn_greater_or_equal_2_10, is_flash_attn_2_available
32
+ if is_flash_attn_2_available():
33
+ from flash_attn import flash_attn_func, flash_attn_varlen_func
34
+ from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input # noqa
35
+ except:
36
+ pass
37
+
38
+
39
+ # flags required to enable jit fusion kernels
40
+
41
+ if sys.platform != 'darwin' and not is_torch_npu_available():
42
+ torch._C._jit_set_profiling_mode(False)
43
+ torch._C._jit_set_profiling_executor(False)
44
+ torch._C._jit_override_can_fuse_on_cpu(True)
45
+ torch._C._jit_override_can_fuse_on_gpu(True)
46
+
47
+ logger = logging.get_logger(__name__)
48
+
49
+ _CHECKPOINT_FOR_DOC = "THUDM/ChatGLM"
50
+ _CONFIG_FOR_DOC = "ChatGLMConfig"
51
+
52
+
53
+ def default_init(cls, *args, **kwargs):
54
+ return cls(*args, **kwargs)
55
+
56
+
57
+ class InvalidScoreLogitsProcessor(LogitsProcessor):
58
+ def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:
59
+ if torch.isnan(scores).any() or torch.isinf(scores).any():
60
+ scores.zero_()
61
+ scores[..., 198] = 5e4
62
+ return scores
63
+
64
+
65
+ def split_tensor_along_last_dim(
66
+ tensor: torch.Tensor,
67
+ num_partitions: int,
68
+ contiguous_split_chunks: bool = False,
69
+ ) -> List[torch.Tensor]:
70
+ """Split a tensor along its last dimension.
71
+
72
+ Arguments:
73
+ tensor: input tensor.
74
+ num_partitions: number of partitions to split the tensor
75
+ contiguous_split_chunks: If True, make each chunk contiguous
76
+ in memory.
77
+
78
+ Returns:
79
+ A list of Tensors
80
+ """
81
+ # Get the size and dimension.
82
+ last_dim = tensor.dim() - 1
83
+ last_dim_size = tensor.size()[last_dim] // num_partitions
84
+ # Split.
85
+ tensor_list = torch.split(tensor, last_dim_size, dim=last_dim)
86
+ # Note: torch.split does not create contiguous tensors by default.
87
+ if contiguous_split_chunks:
88
+ return tuple(chunk.contiguous() for chunk in tensor_list)
89
+
90
+ return tensor_list
91
+
92
+
93
+ class RotaryEmbedding(nn.Module):
94
+ def __init__(self, dim, rope_ratio=1, original_impl=False, device=None, dtype=None):
95
+ super().__init__()
96
+ inv_freq = 1.0 / (10000 ** (torch.arange(0, dim, 2, device=device).to(dtype=dtype) / dim))
97
+ self.register_buffer("inv_freq", inv_freq)
98
+ self.dim = dim
99
+ self.original_impl = original_impl
100
+ self.rope_ratio = rope_ratio
101
+
102
+ def forward_impl(
103
+ self, seq_len: int, n_elem: int, dtype: torch.dtype, device: torch.device, base: int = 10000
104
+ ):
105
+ """Enhanced Transformer with Rotary Position Embedding.
106
+
107
+ Derived from: https://github.com/labmlai/annotated_deep_learning_paper_implementations/blob/master/labml_nn/
108
+ transformers/rope/__init__.py. MIT License:
109
+ https://github.com/labmlai/annotated_deep_learning_paper_implementations/blob/master/license.
110
+ """
111
+ # $\Theta = {\theta_i = 10000^{\frac{2(i-1)}{d}}, i \in [1, 2, ..., \frac{d}{2}]}$
112
+ base = base * self.rope_ratio
113
+ theta = 1.0 / (base ** (torch.arange(0, n_elem, 2, dtype=torch.float, device=device) / n_elem))
114
+
115
+ # Create position indexes `[0, 1, ..., seq_len - 1]`
116
+ seq_idx = torch.arange(seq_len, dtype=torch.float, device=device)
117
+
118
+ # Calculate the product of position index and $\theta_i$
119
+ idx_theta = torch.outer(seq_idx, theta).float()
120
+
121
+ cache = torch.stack([torch.cos(idx_theta), torch.sin(idx_theta)], dim=-1)
122
+
123
+ # this is to mimic the behaviour of complex32, else we will get different results
124
+ if dtype in (torch.float16, torch.bfloat16, torch.int8):
125
+ cache = cache.bfloat16() if dtype == torch.bfloat16 else cache.half()
126
+ return cache
127
+
128
+ def forward(self, max_seq_len, offset=0):
129
+ return self.forward_impl(
130
+ max_seq_len, self.dim, dtype=self.inv_freq.dtype, device=self.inv_freq.device
131
+ )
132
+
133
+
134
+ @torch.jit.script
135
+ def apply_rotary_pos_emb(x: torch.Tensor, rope_cache: torch.Tensor) -> torch.Tensor:
136
+ # x: [b, np, sq, hn]
137
+ b, np, sq, hn = x.size(0), x.size(1), x.size(2), x.size(3)
138
+ rot_dim = rope_cache.shape[-2] * 2
139
+ x, x_pass = x[..., :rot_dim], x[..., rot_dim:]
140
+ # truncate to support variable sizes
141
+ rope_cache = rope_cache[:, :sq]
142
+ xshaped = x.reshape(b, np, sq, rot_dim // 2, 2)
143
+ rope_cache = rope_cache.view(-1, 1, sq, xshaped.size(3), 2)
144
+ x_out2 = torch.stack(
145
+ [
146
+ xshaped[..., 0] * rope_cache[..., 0] - xshaped[..., 1] * rope_cache[..., 1],
147
+ xshaped[..., 1] * rope_cache[..., 0] + xshaped[..., 0] * rope_cache[..., 1],
148
+ ],
149
+ -1,
150
+ )
151
+ x_out2 = x_out2.flatten(3)
152
+ return torch.cat((x_out2, x_pass), dim=-1)
153
+
154
+
155
+ class RMSNorm(torch.nn.Module):
156
+ def __init__(self, normalized_shape, eps=1e-5, device=None, dtype=None, **kwargs):
157
+ super().__init__()
158
+ self.weight = torch.nn.Parameter(torch.empty(normalized_shape, device=device, dtype=dtype))
159
+ self.eps = eps
160
+
161
+ def forward(self, hidden_states: torch.Tensor):
162
+ input_dtype = hidden_states.dtype
163
+ variance = hidden_states.to(torch.float32).pow(2).mean(-1, keepdim=True)
164
+ hidden_states = hidden_states * torch.rsqrt(variance + self.eps)
165
+
166
+ return (self.weight * hidden_states).to(input_dtype)
167
+
168
+
169
+ class CoreAttention(torch.nn.Module):
170
+ def __init__(self, config: ChatGLMConfig, layer_number):
171
+ super(CoreAttention, self).__init__()
172
+ self.config = config
173
+ self.apply_query_key_layer_scaling = config.apply_query_key_layer_scaling
174
+ self.attention_softmax_in_fp32 = config.attention_softmax_in_fp32
175
+ if self.apply_query_key_layer_scaling:
176
+ self.attention_softmax_in_fp32 = True
177
+ self.layer_number = max(1, layer_number)
178
+ self.is_causal = True
179
+
180
+ projection_size = config.kv_channels * config.num_attention_heads
181
+
182
+ # Per attention head and per partition values.
183
+ self.hidden_size_per_partition = projection_size
184
+ self.hidden_size_per_attention_head = projection_size // config.num_attention_heads
185
+ self.num_attention_heads_per_partition = config.num_attention_heads
186
+
187
+ coeff = None
188
+ self.norm_factor = math.sqrt(self.hidden_size_per_attention_head)
189
+ if self.apply_query_key_layer_scaling:
190
+ coeff = self.layer_number
191
+ self.norm_factor *= coeff
192
+ self.coeff = coeff
193
+
194
+ self.attention_dropout = torch.nn.Dropout(config.attention_dropout)
195
+
196
+ def forward(self, query_layer, key_layer, value_layer, attention_mask):
197
+ # [b, np, sq, sk]
198
+ output_size = (query_layer.size(0), query_layer.size(1), query_layer.size(2), key_layer.size(2))
199
+
200
+ # [b, np, sq, hn] -> [b * np, sq, hn]
201
+ query_layer = query_layer.view(output_size[0] * output_size[1], output_size[2], -1)
202
+ # [b, np, sk, hn] -> [b * np, sk, hn]
203
+ key_layer = key_layer.view(output_size[0] * output_size[1], output_size[3], -1)
204
+
205
+ # preallocting input tensor: [b * np, sq, sk]
206
+ matmul_input_buffer = torch.empty(
207
+ output_size[0] * output_size[1], output_size[2], output_size[3], dtype=query_layer.dtype,
208
+ device=query_layer.device
209
+ )
210
+
211
+ # Raw attention scores. [b * np, sq, sk]
212
+ matmul_result = torch.baddbmm(
213
+ matmul_input_buffer,
214
+ query_layer, # [b * np, sq, hn]
215
+ key_layer.transpose(1, 2), # [b * np, hn, sk]
216
+ beta=0.0,
217
+ alpha=(1.0 / self.norm_factor),
218
+ )
219
+
220
+ # change view to [b, np, sq, sk]
221
+ attention_scores = matmul_result.view(*output_size)
222
+
223
+ # ===========================
224
+ # Attention probs and dropout
225
+ # ===========================
226
+
227
+ # attention scores and attention mask [b, np, sq, sk]
228
+ if self.attention_softmax_in_fp32:
229
+ attention_scores = attention_scores.float()
230
+ if self.coeff is not None:
231
+ attention_scores = attention_scores * self.coeff
232
+ if attention_mask is None and attention_scores.shape[2] == attention_scores.shape[3]:
233
+ attention_mask = torch.ones(output_size[0], 1, output_size[2], output_size[3],
234
+ device=attention_scores.device, dtype=torch.bool)
235
+ attention_mask.tril_()
236
+ attention_mask = ~attention_mask
237
+ if attention_mask is not None:
238
+ attention_scores = attention_scores.masked_fill(attention_mask, float("-inf"))
239
+ attention_probs = F.softmax(attention_scores, dim=-1)
240
+ attention_probs = attention_probs.type_as(value_layer)
241
+
242
+ # This is actually dropping out entire tokens to attend to, which might
243
+ # seem a bit unusual, but is taken from the original Transformer paper.
244
+ attention_probs = self.attention_dropout(attention_probs)
245
+
246
+ # query layer shape: [b * np, sq, hn]
247
+ # value layer shape: [b, np, sk, hn]
248
+ # attention shape: [b, np, sq, sk]
249
+ # context layer shape: [b, np, sq, hn]
250
+ output_size = (value_layer.size(0), value_layer.size(1), query_layer.size(1), value_layer.size(3))
251
+ # change view [b * np, sk, hn]
252
+ value_layer = value_layer.view(output_size[0] * output_size[1], value_layer.size(2), -1)
253
+ # change view [b * np, sq, sk]
254
+ attention_probs = attention_probs.view(output_size[0] * output_size[1], output_size[2], -1)
255
+ # matmul: [b * np, sq, hn]
256
+ context_layer = torch.bmm(attention_probs, value_layer)
257
+ # change view [b, np, sq, hn]
258
+ context_layer = context_layer.view(*output_size)
259
+ # [b, np, sq, hn] --> [b, sq, np, hn]
260
+ context_layer = context_layer.transpose(1, 2).contiguous()
261
+ # [b, sq, np, hn] --> [b, sq, hp]
262
+ new_context_layer_shape = context_layer.size()[:-2] + (self.hidden_size_per_partition,)
263
+ context_layer = context_layer.reshape(*new_context_layer_shape)
264
+
265
+ return context_layer
266
+
267
+
268
+ class SdpaAttention(CoreAttention):
269
+ def forward(self, query_layer, key_layer, value_layer, attention_mask):
270
+ if attention_mask is None and query_layer.shape[2] == key_layer.shape[2]:
271
+ context_layer = torch.nn.functional.scaled_dot_product_attention(query_layer, key_layer, value_layer,
272
+ is_causal=True,
273
+ dropout_p=self.config.attention_dropout if self.training else 0.0)
274
+ else:
275
+ if attention_mask is not None:
276
+ attention_mask = ~attention_mask
277
+ context_layer = torch.nn.functional.scaled_dot_product_attention(query_layer, key_layer, value_layer,
278
+ attention_mask,
279
+ dropout_p=self.config.attention_dropout if self.training else 0.0)
280
+ context_layer = context_layer.transpose(1, 2).contiguous()
281
+ new_context_layer_shape = context_layer.size()[:-2] + (self.hidden_size_per_partition,)
282
+ context_layer = context_layer.reshape(*new_context_layer_shape)
283
+ return context_layer
284
+
285
+
286
+ def _get_unpad_data(attention_mask):
287
+ seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32)
288
+ indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()
289
+ max_seqlen_in_batch = seqlens_in_batch.max().item()
290
+ cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0))
291
+ return (
292
+ indices,
293
+ cu_seqlens,
294
+ max_seqlen_in_batch,
295
+ )
296
+
297
+
298
+ # Copied from transformers.models.llama.modeling_llama.LlamaFlashAttention2
299
+ class FlashAttention2(CoreAttention):
300
+ def __init__(self, *args, **kwargs):
301
+ super().__init__(*args, **kwargs)
302
+ self._flash_attn_uses_top_left_mask = not is_flash_attn_greater_or_equal_2_10()
303
+
304
+ def forward(self, query_states, key_states, value_states, attention_mask):
305
+ query_states = query_states.transpose(1, 2)
306
+ key_states = key_states.transpose(1, 2)
307
+ value_states = value_states.transpose(1, 2)
308
+ batch_size, query_length = query_states.shape[:2]
309
+ if not self._flash_attn_uses_top_left_mask:
310
+ causal = self.is_causal
311
+ else:
312
+ # TODO: Remove the `query_length != 1` check once Flash Attention for RoCm is bumped to 2.1. For details, please see the comment in LlamaFlashAttention2 __init__.
313
+ causal = self.is_causal and query_length != 1
314
+ dropout = self.config.attention_dropout if self.training else 0.0
315
+ # Contains at least one padding token in the sequence
316
+ if attention_mask is not None:
317
+ query_states, key_states, value_states, indices_q, cu_seq_lens, max_seq_lens = self._upad_input(
318
+ query_states, key_states, value_states, attention_mask, query_length
319
+ )
320
+
321
+ cu_seqlens_q, cu_seqlens_k = cu_seq_lens
322
+ max_seqlen_in_batch_q, max_seqlen_in_batch_k = max_seq_lens
323
+
324
+ attn_output_unpad = flash_attn_varlen_func(
325
+ query_states,
326
+ key_states,
327
+ value_states,
328
+ cu_seqlens_q=cu_seqlens_q,
329
+ cu_seqlens_k=cu_seqlens_k,
330
+ max_seqlen_q=max_seqlen_in_batch_q,
331
+ max_seqlen_k=max_seqlen_in_batch_k,
332
+ dropout_p=dropout,
333
+ softmax_scale=None,
334
+ causal=causal,
335
+ )
336
+
337
+ attn_output = pad_input(attn_output_unpad, indices_q, batch_size, query_length)
338
+ else:
339
+ attn_output = flash_attn_func(
340
+ query_states, key_states, value_states, dropout, softmax_scale=None, causal=causal
341
+ )
342
+ attn_output = attn_output.reshape(batch_size, query_length, self.hidden_size_per_partition).contiguous()
343
+ return attn_output
344
+
345
+ def _upad_input(self, query_layer, key_layer, value_layer, attention_mask, query_length):
346
+ indices_k, cu_seqlens_k, max_seqlen_in_batch_k = _get_unpad_data(attention_mask)
347
+ batch_size, kv_seq_len, num_key_value_heads, head_dim = key_layer.shape
348
+
349
+ key_layer = index_first_axis(
350
+ key_layer.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim), indices_k
351
+ )
352
+ value_layer = index_first_axis(
353
+ value_layer.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim), indices_k
354
+ )
355
+ if query_length == kv_seq_len:
356
+ query_layer = index_first_axis(
357
+ query_layer.reshape(batch_size * kv_seq_len, self.num_attention_heads_per_partition, head_dim), indices_k
358
+ )
359
+ cu_seqlens_q = cu_seqlens_k
360
+ max_seqlen_in_batch_q = max_seqlen_in_batch_k
361
+ indices_q = indices_k
362
+ elif query_length == 1:
363
+ max_seqlen_in_batch_q = 1
364
+ cu_seqlens_q = torch.arange(
365
+ batch_size + 1, dtype=torch.int32, device=query_layer.device
366
+ ) # There is a memcpy here, that is very bad.
367
+ indices_q = cu_seqlens_q[:-1]
368
+ query_layer = query_layer.squeeze(1)
369
+ else:
370
+ # The -q_len: slice assumes left padding.
371
+ attention_mask = attention_mask[:, -query_length:]
372
+ query_layer, indices_q, cu_seqlens_q, max_seqlen_in_batch_q = unpad_input(query_layer, attention_mask)
373
+
374
+ return (
375
+ query_layer,
376
+ key_layer,
377
+ value_layer,
378
+ indices_q,
379
+ (cu_seqlens_q, cu_seqlens_k),
380
+ (max_seqlen_in_batch_q, max_seqlen_in_batch_k),
381
+ )
382
+
383
+
384
+ CORE_ATTENTION_CLASSES = {
385
+ "eager": CoreAttention,
386
+ "sdpa": SdpaAttention,
387
+ "flash_attention_2": FlashAttention2
388
+ }
389
+
390
+
391
+ class SelfAttention(torch.nn.Module):
392
+ """Parallel self-attention layer abstract class.
393
+
394
+ Self-attention layer takes input with size [s, b, h]
395
+ and returns output of the same size.
396
+ """
397
+
398
+ def __init__(self, config: ChatGLMConfig, layer_number, device=None):
399
+ super(SelfAttention, self).__init__()
400
+ self.layer_number = max(1, layer_number)
401
+
402
+ self.projection_size = config.kv_channels * config.num_attention_heads
403
+
404
+ # Per attention head and per partition values.
405
+ self.hidden_size_per_attention_head = self.projection_size // config.num_attention_heads
406
+ self.num_attention_heads_per_partition = config.num_attention_heads
407
+
408
+ self.multi_query_attention = config.multi_query_attention
409
+ self.qkv_hidden_size = 3 * self.projection_size
410
+ if self.multi_query_attention:
411
+ self.num_multi_query_groups_per_partition = config.multi_query_group_num
412
+ self.qkv_hidden_size = (
413
+ self.projection_size + 2 * self.hidden_size_per_attention_head * config.multi_query_group_num
414
+ )
415
+ self.query_key_value = nn.Linear(config.hidden_size, self.qkv_hidden_size,
416
+ bias=config.add_bias_linear or config.add_qkv_bias,
417
+ device=device, **_config_to_kwargs(config)
418
+ )
419
+
420
+ self.core_attention = CORE_ATTENTION_CLASSES[config._attn_implementation](config, self.layer_number)
421
+
422
+ # Output.
423
+ self.dense = nn.Linear(self.projection_size, config.hidden_size, bias=config.add_bias_linear,
424
+ device=device, **_config_to_kwargs(config)
425
+ )
426
+
427
+ def _allocate_memory(self, inference_max_sequence_len, batch_size, device=None, dtype=None):
428
+ if self.multi_query_attention:
429
+ num_attention_heads = self.num_multi_query_groups_per_partition
430
+ else:
431
+ num_attention_heads = self.num_attention_heads_per_partition
432
+ return torch.empty(
433
+ inference_max_sequence_len,
434
+ batch_size,
435
+ num_attention_heads,
436
+ self.hidden_size_per_attention_head,
437
+ dtype=dtype,
438
+ device=device,
439
+ )
440
+
441
+ def forward(
442
+ self, hidden_states, attention_mask, rotary_pos_emb, kv_cache=None, use_cache=True
443
+ ):
444
+ # hidden_states: [b, sq, h]
445
+
446
+ # =================================================
447
+ # Pre-allocate memory for key-values for inference.
448
+ # =================================================
449
+ # =====================
450
+ # Query, Key, and Value
451
+ # =====================
452
+
453
+ # Attention heads [b, sq, h] --> [b, sq, (np * 3 * hn)]
454
+ mixed_x_layer = self.query_key_value(hidden_states)
455
+
456
+ if self.multi_query_attention:
457
+ (query_layer, key_layer, value_layer) = mixed_x_layer.split(
458
+ [
459
+ self.num_attention_heads_per_partition * self.hidden_size_per_attention_head,
460
+ self.num_multi_query_groups_per_partition * self.hidden_size_per_attention_head,
461
+ self.num_multi_query_groups_per_partition * self.hidden_size_per_attention_head,
462
+ ],
463
+ dim=-1,
464
+ )
465
+ query_layer = query_layer.view(
466
+ query_layer.size()[:-1] + (self.num_attention_heads_per_partition, self.hidden_size_per_attention_head)
467
+ )
468
+ key_layer = key_layer.view(
469
+ key_layer.size()[:-1] + (self.num_multi_query_groups_per_partition, self.hidden_size_per_attention_head)
470
+ )
471
+ value_layer = value_layer.view(
472
+ value_layer.size()[:-1]
473
+ + (self.num_multi_query_groups_per_partition, self.hidden_size_per_attention_head)
474
+ )
475
+ else:
476
+ new_tensor_shape = mixed_x_layer.size()[:-1] + \
477
+ (self.num_attention_heads_per_partition,
478
+ 3 * self.hidden_size_per_attention_head)
479
+ mixed_x_layer = mixed_x_layer.view(*new_tensor_shape)
480
+
481
+ # [b, sq, np, 3 * hn] --> 3 [b, sq, np, hn]
482
+ (query_layer, key_layer, value_layer) = split_tensor_along_last_dim(mixed_x_layer, 3)
483
+
484
+ # [b, sq, np, hn] -> [b, np, sq, hn]
485
+ query_layer, key_layer, value_layer = [k.transpose(1, 2) for k in [query_layer, key_layer, value_layer]]
486
+
487
+ # apply relative positional encoding (rotary embedding)
488
+ if rotary_pos_emb is not None:
489
+ query_layer = apply_rotary_pos_emb(query_layer, rotary_pos_emb)
490
+ key_layer = apply_rotary_pos_emb(key_layer, rotary_pos_emb)
491
+
492
+ # adjust key and value for inference
493
+ if kv_cache is not None:
494
+ cache_k, cache_v = kv_cache
495
+ key_layer = torch.cat((cache_k, key_layer), dim=2)
496
+ value_layer = torch.cat((cache_v, value_layer), dim=2)
497
+ if use_cache:
498
+ if kv_cache is None:
499
+ kv_cache = torch.cat((key_layer.unsqueeze(0).unsqueeze(0), value_layer.unsqueeze(0).unsqueeze(0)),
500
+ dim=1)
501
+ else:
502
+ kv_cache = (key_layer, value_layer)
503
+ else:
504
+ kv_cache = None
505
+
506
+ if self.multi_query_attention:
507
+ key_layer = key_layer.unsqueeze(2)
508
+ key_layer = key_layer.expand(
509
+ -1, -1, self.num_attention_heads_per_partition // self.num_multi_query_groups_per_partition, -1, -1
510
+ )
511
+ key_layer = key_layer.contiguous().view(
512
+ key_layer.size()[:1] + (self.num_attention_heads_per_partition,) + key_layer.size()[3:]
513
+ )
514
+ value_layer = value_layer.unsqueeze(2)
515
+ value_layer = value_layer.expand(
516
+ -1, -1, self.num_attention_heads_per_partition // self.num_multi_query_groups_per_partition, -1, -1
517
+ )
518
+ value_layer = value_layer.contiguous().view(
519
+ value_layer.size()[:1] + (self.num_attention_heads_per_partition,) + value_layer.size()[3:]
520
+ )
521
+
522
+ # ==================================
523
+ # core attention computation
524
+ # ==================================
525
+
526
+ context_layer = self.core_attention(query_layer, key_layer, value_layer, attention_mask)
527
+
528
+ # =================
529
+ # Output. [sq, b, h]
530
+ # =================
531
+
532
+ output = self.dense(context_layer)
533
+
534
+ return output, kv_cache
535
+
536
+
537
+ def _config_to_kwargs(args):
538
+ common_kwargs = {
539
+ "dtype": args.torch_dtype,
540
+ }
541
+ return common_kwargs
542
+
543
+
544
+ class MLP(torch.nn.Module):
545
+ """MLP.
546
+
547
+ MLP will take the input with h hidden state, project it to 4*h
548
+ hidden dimension, perform nonlinear transformation, and project the
549
+ state back into h hidden dimension.
550
+ """
551
+
552
+ def __init__(self, config: ChatGLMConfig, device=None):
553
+ super(MLP, self).__init__()
554
+
555
+ self.add_bias = config.add_bias_linear
556
+
557
+ # Project to 4h. If using swiglu double the output width, see https://arxiv.org/pdf/2002.05202.pdf
558
+ self.dense_h_to_4h = nn.Linear(
559
+ config.hidden_size,
560
+ config.ffn_hidden_size * 2,
561
+ bias=self.add_bias,
562
+ device=device,
563
+ **_config_to_kwargs(config)
564
+ )
565
+
566
+ def swiglu(x):
567
+ x = torch.chunk(x, 2, dim=-1)
568
+ return F.silu(x[0]) * x[1]
569
+
570
+ self.activation_func = swiglu
571
+
572
+ # Project back to h.
573
+ self.dense_4h_to_h = nn.Linear(
574
+ config.ffn_hidden_size,
575
+ config.hidden_size,
576
+ bias=self.add_bias,
577
+ device=device,
578
+ **_config_to_kwargs(config)
579
+ )
580
+
581
+ def forward(self, hidden_states):
582
+ # [s, b, 4hp]
583
+ intermediate_parallel = self.dense_h_to_4h(hidden_states)
584
+ intermediate_parallel = self.activation_func(intermediate_parallel)
585
+ # [s, b, h]
586
+ output = self.dense_4h_to_h(intermediate_parallel)
587
+ return output
588
+
589
+
590
+ class GLMBlock(torch.nn.Module):
591
+ """A single transformer layer.
592
+
593
+ Transformer layer takes input with size [s, b, h] and returns an
594
+ output of the same size.
595
+ """
596
+
597
+ def __init__(self, config: ChatGLMConfig, layer_number, device=None):
598
+ super(GLMBlock, self).__init__()
599
+ self.layer_number = layer_number
600
+
601
+ self.apply_residual_connection_post_layernorm = config.apply_residual_connection_post_layernorm
602
+
603
+ self.fp32_residual_connection = config.fp32_residual_connection
604
+
605
+ LayerNormFunc = RMSNorm if config.rmsnorm else LayerNorm
606
+ # Layernorm on the input data.
607
+ self.input_layernorm = LayerNormFunc(config.hidden_size, eps=config.layernorm_epsilon, device=device,
608
+ dtype=config.torch_dtype)
609
+
610
+ # Self attention.
611
+ self.self_attention = SelfAttention(config, layer_number, device=device)
612
+ self.hidden_dropout = config.hidden_dropout
613
+
614
+ # Layernorm on the attention output
615
+ self.post_attention_layernorm = LayerNormFunc(config.hidden_size, eps=config.layernorm_epsilon, device=device,
616
+ dtype=config.torch_dtype)
617
+
618
+ # MLP
619
+ self.mlp = MLP(config, device=device)
620
+
621
+ def forward(
622
+ self, hidden_states, attention_mask, rotary_pos_emb, kv_cache=None, use_cache=True,
623
+ ):
624
+ # hidden_states: [s, b, h]
625
+
626
+ # Layer norm at the beginning of the transformer layer.
627
+ layernorm_output = self.input_layernorm(hidden_states)
628
+ # Self attention.
629
+ attention_output, kv_cache = self.self_attention(
630
+ layernorm_output,
631
+ attention_mask,
632
+ rotary_pos_emb,
633
+ kv_cache=kv_cache,
634
+ use_cache=use_cache
635
+ )
636
+
637
+ # Residual connection.
638
+ if self.apply_residual_connection_post_layernorm:
639
+ residual = layernorm_output
640
+ else:
641
+ residual = hidden_states
642
+
643
+ layernorm_input = torch.nn.functional.dropout(attention_output, p=self.hidden_dropout, training=self.training)
644
+ layernorm_input = residual + layernorm_input
645
+
646
+ # Layer norm post the self attention.
647
+ layernorm_output = self.post_attention_layernorm(layernorm_input)
648
+
649
+ # MLP.
650
+ mlp_output = self.mlp(layernorm_output)
651
+
652
+ # Second residual connection.
653
+ if self.apply_residual_connection_post_layernorm:
654
+ residual = layernorm_output
655
+ else:
656
+ residual = layernorm_input
657
+
658
+ output = torch.nn.functional.dropout(mlp_output, p=self.hidden_dropout, training=self.training)
659
+ output = residual + output
660
+
661
+ return output, kv_cache
662
+
663
+
664
+ class GLMTransformer(torch.nn.Module):
665
+ """Transformer class."""
666
+
667
+ def __init__(self, config: ChatGLMConfig, device=None):
668
+ super(GLMTransformer, self).__init__()
669
+
670
+ self.fp32_residual_connection = config.fp32_residual_connection
671
+ self.post_layer_norm = config.post_layer_norm
672
+
673
+ # Number of layers.
674
+ self.num_layers = config.num_layers
675
+
676
+ # Transformer layers.
677
+ def build_layer(layer_number):
678
+ return GLMBlock(config, layer_number, device=device)
679
+
680
+ self.layers = torch.nn.ModuleList([build_layer(i + 1) for i in range(self.num_layers)])
681
+
682
+ if self.post_layer_norm:
683
+ LayerNormFunc = RMSNorm if config.rmsnorm else LayerNorm
684
+ # Final layer norm before output.
685
+ self.final_layernorm = LayerNormFunc(config.hidden_size, eps=config.layernorm_epsilon, device=device,
686
+ dtype=config.torch_dtype)
687
+
688
+ self.gradient_checkpointing = False
689
+
690
+ def _get_layer(self, layer_number):
691
+ return self.layers[layer_number]
692
+
693
+ def forward(
694
+ self, hidden_states, attention_mask, rotary_pos_emb, kv_caches=None,
695
+ use_cache: Optional[bool] = True,
696
+ output_hidden_states: Optional[bool] = False,
697
+ ):
698
+ if not kv_caches:
699
+ kv_caches = [None for _ in range(self.num_layers)]
700
+ presents = () if use_cache else None
701
+ if self.gradient_checkpointing and self.training:
702
+ if use_cache:
703
+ logger.warning_once(
704
+ "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
705
+ )
706
+ use_cache = False
707
+
708
+ all_self_attentions = None
709
+ all_hidden_states = () if output_hidden_states else None
710
+ for index in range(self.num_layers):
711
+ if output_hidden_states:
712
+ all_hidden_states = all_hidden_states + (hidden_states,)
713
+
714
+ layer = self._get_layer(index)
715
+ if self.gradient_checkpointing and self.training:
716
+ layer_ret = torch.utils.checkpoint.checkpoint(
717
+ layer,
718
+ hidden_states,
719
+ attention_mask,
720
+ rotary_pos_emb,
721
+ kv_caches[index],
722
+ use_cache,
723
+ use_reentrant=False
724
+ )
725
+ else:
726
+ layer_ret = layer(
727
+ hidden_states,
728
+ attention_mask,
729
+ rotary_pos_emb,
730
+ kv_cache=kv_caches[index],
731
+ use_cache=use_cache
732
+ )
733
+ hidden_states, kv_cache = layer_ret
734
+ if use_cache:
735
+ # token by token decoding, use tuple format
736
+ if kv_caches[0] is not None:
737
+ presents = presents + (kv_cache,)
738
+ # prefilling in decoding, use tensor format to save cuda memory
739
+ else:
740
+ if len(presents) == 0:
741
+ presents = kv_cache
742
+ else:
743
+ presents = torch.cat((presents, kv_cache.to(presents.device)), dim=0)
744
+
745
+ if output_hidden_states:
746
+ all_hidden_states = all_hidden_states + (hidden_states,)
747
+
748
+ # Final layer norm.
749
+ if self.post_layer_norm:
750
+ hidden_states = self.final_layernorm(hidden_states)
751
+
752
+ return hidden_states, presents, all_hidden_states, all_self_attentions
753
+
754
+
755
+ class ChatGLMPreTrainedModel(PreTrainedModel):
756
+ """
757
+ An abstract class to handle weights initialization and
758
+ a simple interface for downloading and loading pretrained models.
759
+ """
760
+
761
+ is_parallelizable = False
762
+ supports_gradient_checkpointing = True
763
+ config_class = ChatGLMConfig
764
+ base_model_prefix = "transformer"
765
+ _no_split_modules = ["GLMBlock"]
766
+ _supports_flash_attn_2 = True
767
+ _supports_sdpa = True
768
+
769
+ def _init_weights(self, module: nn.Module):
770
+ """Initialize the weights."""
771
+ return
772
+
773
+ def get_masks(self, input_ids, past_key_values, padding_mask=None):
774
+ if self.config._attn_implementation == "flash_attention_2":
775
+ if padding_mask is not None and not padding_mask.all():
776
+ return padding_mask
777
+ return None
778
+ batch_size, seq_length = input_ids.shape
779
+ full_attention_mask = torch.ones(batch_size, seq_length, seq_length, device=input_ids.device)
780
+ full_attention_mask.tril_()
781
+ past_length = 0
782
+ if past_key_values:
783
+ past_length = past_key_values[0][0].shape[2]
784
+ if past_length:
785
+ full_attention_mask = torch.cat((torch.ones(batch_size, seq_length, past_length,
786
+ device=input_ids.device), full_attention_mask), dim=-1)
787
+ if padding_mask is not None:
788
+ full_attention_mask = full_attention_mask * padding_mask.unsqueeze(1)
789
+ if not past_length and padding_mask is not None:
790
+ full_attention_mask -= padding_mask.unsqueeze(-1) - 1
791
+ full_attention_mask = (full_attention_mask < 0.5).bool()
792
+ full_attention_mask.unsqueeze_(1)
793
+ return full_attention_mask
794
+
795
+ def get_position_ids(self, input_ids, device):
796
+ batch_size, seq_length = input_ids.shape
797
+ position_ids = torch.arange(seq_length, dtype=torch.long, device=device).unsqueeze(0).repeat(batch_size, 1)
798
+ return position_ids
799
+
800
+ def gradient_checkpointing_enable(self, gradient_checkpointing_kwargs=None):
801
+ if not self.supports_gradient_checkpointing:
802
+ raise ValueError(f"{self.__class__.__name__} does not support gradient checkpointing.")
803
+
804
+
805
+ class Embedding(torch.nn.Module):
806
+ """Language model embeddings."""
807
+
808
+ def __init__(self, config: ChatGLMConfig, device=None):
809
+ super(Embedding, self).__init__()
810
+
811
+ self.hidden_size = config.hidden_size
812
+ # Word embeddings (parallel).
813
+ self.word_embeddings = nn.Embedding(
814
+ config.padded_vocab_size,
815
+ self.hidden_size,
816
+ dtype=config.torch_dtype,
817
+ device=device
818
+ )
819
+ self.fp32_residual_connection = config.fp32_residual_connection
820
+
821
+ def forward(self, input_ids):
822
+ # Embeddings.
823
+ words_embeddings = self.word_embeddings(input_ids)
824
+ embeddings = words_embeddings
825
+ # If the input flag for fp32 residual connection is set, convert for float.
826
+ if self.fp32_residual_connection:
827
+ embeddings = embeddings.float()
828
+ return embeddings
829
+
830
+
831
+ class ChatGLMModel(ChatGLMPreTrainedModel):
832
+ def __init__(self, config: ChatGLMConfig, device=None, empty_init=True):
833
+ super().__init__(config)
834
+ if empty_init:
835
+ init_method = skip_init
836
+ else:
837
+ init_method = default_init
838
+ init_kwargs = {}
839
+ if device is not None:
840
+ init_kwargs["device"] = device
841
+ self.embedding = init_method(Embedding, config, **init_kwargs)
842
+ self.num_layers = config.num_layers
843
+ self.multi_query_group_num = config.multi_query_group_num
844
+ self.kv_channels = config.kv_channels
845
+
846
+ # Rotary positional embeddings
847
+ self.seq_length = config.seq_length
848
+ rotary_dim = (
849
+ config.hidden_size // config.num_attention_heads if config.kv_channels is None else config.kv_channels
850
+ )
851
+
852
+ self.rotary_pos_emb = RotaryEmbedding(rotary_dim // 2, rope_ratio=config.rope_ratio,
853
+ original_impl=config.original_rope,
854
+ device=device, dtype=config.torch_dtype)
855
+ self.encoder = init_method(GLMTransformer, config, **init_kwargs)
856
+ self.output_layer = init_method(nn.Linear, config.hidden_size, config.padded_vocab_size, bias=False,
857
+ dtype=config.torch_dtype, **init_kwargs)
858
+
859
+ def get_input_embeddings(self):
860
+ return self.embedding.word_embeddings
861
+
862
+ def set_input_embeddings(self, value):
863
+ self.embedding.word_embeddings = value
864
+
865
+ def forward(
866
+ self,
867
+ input_ids,
868
+ position_ids: Optional[torch.Tensor] = None,
869
+ attention_mask: Optional[torch.BoolTensor] = None,
870
+ full_attention_mask: Optional[torch.BoolTensor] = None,
871
+ past_key_values: Optional[Tuple[Tuple[torch.Tensor, torch.Tensor], ...]] = None,
872
+ inputs_embeds: Optional[torch.Tensor] = None,
873
+ use_cache: Optional[bool] = None,
874
+ output_attentions: Optional[bool] = None,
875
+ output_hidden_states: Optional[bool] = None,
876
+ return_dict: Optional[bool] = None,
877
+ ):
878
+ output_hidden_states = (
879
+ output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
880
+ )
881
+ use_cache = use_cache if use_cache is not None else self.config.use_cache
882
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
883
+
884
+ batch_size, seq_length = input_ids.shape
885
+
886
+ if inputs_embeds is None:
887
+ inputs_embeds = self.embedding(input_ids)
888
+
889
+ if full_attention_mask is None:
890
+ if (attention_mask is not None and not attention_mask.all()) or (past_key_values and seq_length != 1):
891
+ full_attention_mask = self.get_masks(input_ids, past_key_values, padding_mask=attention_mask)
892
+
893
+ # Rotary positional embeddings
894
+ rotary_pos_emb = self.rotary_pos_emb(self.seq_length)
895
+ if position_ids is not None:
896
+ rotary_pos_emb = rotary_pos_emb[position_ids]
897
+ else:
898
+ rotary_pos_emb = rotary_pos_emb[None, :seq_length]
899
+
900
+ # Run encoder.
901
+ hidden_states, presents, all_hidden_states, all_self_attentions = self.encoder(
902
+ inputs_embeds, full_attention_mask, rotary_pos_emb=rotary_pos_emb,
903
+ kv_caches=past_key_values, use_cache=use_cache, output_hidden_states=output_hidden_states
904
+ )
905
+ if presents is not None and type(presents) is torch.Tensor:
906
+ presents = presents.split(1, dim=0)
907
+ presents = list(presents)
908
+ presents = [list(x.squeeze(0).split(1, dim=0)) for x in presents]
909
+ presents = [tuple([x.squeeze(0) for x in y]) for y in presents]
910
+ presents = tuple(presents)
911
+
912
+ if not return_dict:
913
+ return tuple(v for v in [hidden_states, presents, all_hidden_states, all_self_attentions] if v is not None)
914
+
915
+ return BaseModelOutputWithPast(
916
+ last_hidden_state=hidden_states,
917
+ past_key_values=presents,
918
+ hidden_states=all_hidden_states,
919
+ attentions=all_self_attentions,
920
+ )
921
+
922
+
923
+ class ChatGLMForConditionalGeneration(ChatGLMPreTrainedModel):
924
+ def __init__(self, config: ChatGLMConfig, empty_init=True, device=None):
925
+ super().__init__(config)
926
+
927
+ self.max_sequence_length = config.max_length
928
+ self.transformer = ChatGLMModel(config, empty_init=empty_init, device=device)
929
+ self.config = config
930
+
931
+ def _update_model_kwargs_for_generation(
932
+ self,
933
+ outputs: ModelOutput,
934
+ model_kwargs: Dict[str, Any],
935
+ is_encoder_decoder: bool = False,
936
+ standardize_cache_format: bool = False,
937
+ ) -> Dict[str, Any]:
938
+ # update past_key_values
939
+ model_kwargs["past_key_values"] = self._extract_past_from_model_output(
940
+ outputs, standardize_cache_format=standardize_cache_format
941
+ )
942
+
943
+ # update attention mask
944
+ if "attention_mask" in model_kwargs:
945
+ attention_mask = model_kwargs["attention_mask"]
946
+ model_kwargs["attention_mask"] = torch.cat(
947
+ [attention_mask, attention_mask.new_ones((attention_mask.shape[0], 1))], dim=-1
948
+ )
949
+
950
+ # update position ids
951
+ if "position_ids" in model_kwargs:
952
+ position_ids = model_kwargs["position_ids"]
953
+ new_position_id = position_ids[..., -1:].clone()
954
+ new_position_id += 1
955
+ model_kwargs["position_ids"] = torch.cat(
956
+ [position_ids, new_position_id], dim=-1
957
+ )
958
+
959
+ model_kwargs["is_first_forward"] = False
960
+ return model_kwargs
961
+
962
+ def prepare_inputs_for_generation(
963
+ self,
964
+ input_ids: torch.LongTensor,
965
+ past_key_values: Optional[torch.Tensor] = None,
966
+ attention_mask: Optional[torch.Tensor] = None,
967
+ position_ids: Optional[torch.Tensor] = None,
968
+ use_cache: Optional[bool] = None,
969
+ is_first_forward: bool = True,
970
+ **kwargs
971
+ ) -> dict:
972
+ # only last token for input_ids if past is not None
973
+ if position_ids is None:
974
+ position_ids = self.get_position_ids(input_ids, device=input_ids.device)
975
+ if not is_first_forward:
976
+ if past_key_values is not None:
977
+ position_ids = position_ids[..., -1:]
978
+ input_ids = input_ids[:, -1:]
979
+ return {
980
+ "input_ids": input_ids,
981
+ "past_key_values": past_key_values,
982
+ "position_ids": position_ids,
983
+ "attention_mask": attention_mask,
984
+ "return_last_logit": True,
985
+ "use_cache": use_cache
986
+ }
987
+
988
+ def forward(
989
+ self,
990
+ input_ids: Optional[torch.Tensor] = None,
991
+ position_ids: Optional[torch.Tensor] = None,
992
+ attention_mask: Optional[torch.Tensor] = None,
993
+ past_key_values: Optional[Tuple[torch.FloatTensor]] = None,
994
+ inputs_embeds: Optional[torch.Tensor] = None,
995
+ labels: Optional[torch.Tensor] = None,
996
+ use_cache: Optional[bool] = None,
997
+ output_attentions: Optional[bool] = None,
998
+ output_hidden_states: Optional[bool] = None,
999
+ return_dict: Optional[bool] = None,
1000
+ return_last_logit: Optional[bool] = False,
1001
+ ):
1002
+ use_cache = use_cache if use_cache is not None else self.config.use_cache
1003
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
1004
+
1005
+ transformer_outputs = self.transformer(
1006
+ input_ids=input_ids,
1007
+ position_ids=position_ids,
1008
+ attention_mask=attention_mask,
1009
+ past_key_values=past_key_values,
1010
+ inputs_embeds=inputs_embeds,
1011
+ use_cache=use_cache,
1012
+ output_hidden_states=output_hidden_states,
1013
+ return_dict=return_dict,
1014
+ )
1015
+
1016
+ hidden_states = transformer_outputs[0]
1017
+ if return_last_logit:
1018
+ hidden_states = hidden_states[:, -1:]
1019
+ lm_logits = self.transformer.output_layer(hidden_states)
1020
+
1021
+ loss = None
1022
+ if labels is not None:
1023
+ lm_logits = lm_logits.to(torch.float32)
1024
+
1025
+ # Shift so that tokens < n predict n
1026
+ shift_logits = lm_logits[..., :-1, :].contiguous()
1027
+ shift_labels = labels[..., 1:].contiguous()
1028
+ # Flatten the tokens
1029
+ loss_fct = CrossEntropyLoss(ignore_index=-100)
1030
+ loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))
1031
+
1032
+ lm_logits = lm_logits.to(hidden_states.dtype)
1033
+ loss = loss.to(hidden_states.dtype)
1034
+
1035
+ if not return_dict:
1036
+ output = (lm_logits,) + transformer_outputs[1:]
1037
+ return ((loss,) + output) if loss is not None else output
1038
+
1039
+ return CausalLMOutputWithPast(
1040
+ loss=loss,
1041
+ logits=lm_logits,
1042
+ past_key_values=transformer_outputs.past_key_values,
1043
+ hidden_states=transformer_outputs.hidden_states,
1044
+ attentions=transformer_outputs.attentions,
1045
+ )
1046
+
1047
+ @staticmethod
1048
+ def _reorder_cache(
1049
+ past: Tuple[Tuple[torch.Tensor, torch.Tensor], ...], beam_idx: torch.LongTensor
1050
+ ) -> Tuple[Tuple[torch.Tensor, torch.Tensor], ...]:
1051
+ """
1052
+ This function is used to re-order the `past_key_values` cache if [`~PreTrainedModel.beam_search`] or
1053
+ [`~PreTrainedModel.beam_sample`] is called. This is required to match `past_key_values` with the correct
1054
+ beam_idx at every generation step.
1055
+
1056
+ Output shares the same memory storage as `past`.
1057
+ """
1058
+ return tuple(
1059
+ (
1060
+ layer_past[0].index_select(0, beam_idx.to(layer_past[0].device)),
1061
+ layer_past[1].index_select(0, beam_idx.to(layer_past[1].device)),
1062
+ )
1063
+ for layer_past in past
1064
+ )
1065
+
1066
+ def process_response(self, output, history):
1067
+ content = ""
1068
+ history = deepcopy(history)
1069
+ for response in output.split("<|assistant|>"):
1070
+ if "\n" in response:
1071
+ metadata, content = response.split("\n", maxsplit=1)
1072
+ else:
1073
+ metadata, content = "", response
1074
+ if not metadata.strip():
1075
+ content = content.strip()
1076
+ history.append({"role": "assistant", "metadata": metadata, "content": content})
1077
+ content = content.replace("[[训练时间]]", "2023年")
1078
+ else:
1079
+ history.append({"role": "assistant", "metadata": metadata, "content": content})
1080
+ if history[0]["role"] == "system" and "tools" in history[0]:
1081
+ parameters = json.loads(content)
1082
+ content = {"name": metadata.strip(), "parameters": parameters}
1083
+ else:
1084
+ content = {"name": metadata.strip(), "content": content}
1085
+ return content, history
1086
+
1087
+ @torch.inference_mode()
1088
+ def chat(self, tokenizer, query: str, history: List[Dict] = None, role: str = "user",
1089
+ max_length: int = 8192, num_beams=1, do_sample=True, top_p=0.8, temperature=0.8, logits_processor=None,
1090
+ **kwargs):
1091
+ if history is None:
1092
+ history = []
1093
+ if logits_processor is None:
1094
+ logits_processor = LogitsProcessorList()
1095
+ logits_processor.append(InvalidScoreLogitsProcessor())
1096
+ gen_kwargs = {"max_length": max_length, "num_beams": num_beams, "do_sample": do_sample, "top_p": top_p,
1097
+ "temperature": temperature, "logits_processor": logits_processor, **kwargs}
1098
+ history.append({"role": role, "content": query})
1099
+ inputs = tokenizer.apply_chat_template(history, add_generation_prompt=True, tokenize=True,
1100
+ return_tensors="pt", return_dict=True)
1101
+ inputs = inputs.to(self.device)
1102
+ eos_token_id = [tokenizer.eos_token_id, tokenizer.convert_tokens_to_ids("<|user|>"),
1103
+ tokenizer.convert_tokens_to_ids("<|observation|>")]
1104
+ outputs = self.generate(**inputs, **gen_kwargs, eos_token_id=eos_token_id)
1105
+ outputs = outputs.tolist()[0][len(inputs["input_ids"][0]):-1]
1106
+ response = tokenizer.decode(outputs)
1107
+ response, history = self.process_response(response, history)
1108
+ return response, history
1109
+
1110
+ @torch.inference_mode()
1111
+ def stream_chat(self, tokenizer, query: str, history: List[Dict] = None, role: str = "user",
1112
+ past_key_values=None, max_length: int = 8192, do_sample=True, top_p=0.8, temperature=0.8,
1113
+ logits_processor=None, return_past_key_values=False, **kwargs):
1114
+ if history is None:
1115
+ history = []
1116
+ if logits_processor is None:
1117
+ logits_processor = LogitsProcessorList()
1118
+ logits_processor.append(InvalidScoreLogitsProcessor())
1119
+ eos_token_id = [tokenizer.eos_token_id, tokenizer.convert_tokens_to_ids("<|user|>"),
1120
+ tokenizer.convert_tokens_to_ids("<|observation|>")]
1121
+ gen_kwargs = {"max_length": max_length, "do_sample": do_sample, "top_p": top_p,
1122
+ "temperature": temperature, "logits_processor": logits_processor, **kwargs}
1123
+ if past_key_values is None:
1124
+ inputs = tokenizer.apply_chat_template(history + [{"role": role, "content": query}],
1125
+ add_generation_prompt=True, tokenize=True, return_tensors="pt",
1126
+ return_dict=True)
1127
+ else:
1128
+ inputs = tokenizer.apply_chat_template([{"role": role, "content": query}], add_special_tokens=False,
1129
+ add_generation_prompt=True, tokenize=True, return_tensors="pt",
1130
+ return_dict=True)
1131
+ inputs = inputs.to(self.device)
1132
+ if past_key_values is not None:
1133
+ past_length = past_key_values[0][0].shape[2]
1134
+ inputs.position_ids += past_length
1135
+ attention_mask = inputs.attention_mask
1136
+ attention_mask = torch.cat((attention_mask.new_ones(1, past_length), attention_mask), dim=1)
1137
+ inputs['attention_mask'] = attention_mask
1138
+ history.append({"role": role, "content": query})
1139
+ for outputs in self.stream_generate(**inputs, past_key_values=past_key_values,
1140
+ eos_token_id=eos_token_id, return_past_key_values=return_past_key_values,
1141
+ **gen_kwargs):
1142
+ if return_past_key_values:
1143
+ outputs, past_key_values = outputs
1144
+ outputs = outputs.tolist()[0][len(inputs["input_ids"][0]):-1]
1145
+ response = tokenizer.decode(outputs)
1146
+ if response and response[-1] != "�":
1147
+ response, new_history = self.process_response(response, history)
1148
+ if return_past_key_values:
1149
+ yield response, new_history, past_key_values
1150
+ else:
1151
+ yield response, new_history
1152
+
1153
+ @torch.inference_mode()
1154
+ def stream_generate(
1155
+ self,
1156
+ input_ids,
1157
+ generation_config: Optional[GenerationConfig] = None,
1158
+ logits_processor: Optional[LogitsProcessorList] = None,
1159
+ stopping_criteria: Optional[StoppingCriteriaList] = None,
1160
+ prefix_allowed_tokens_fn: Optional[Callable[[int, torch.Tensor], List[int]]] = None,
1161
+ return_past_key_values=False,
1162
+ **kwargs,
1163
+ ):
1164
+ batch_size, input_ids_seq_length = input_ids.shape[0], input_ids.shape[-1]
1165
+
1166
+ if generation_config is None:
1167
+ generation_config = self.generation_config
1168
+ generation_config = copy.deepcopy(generation_config)
1169
+ model_kwargs = generation_config.update(**kwargs)
1170
+ model_kwargs["use_cache"] = generation_config.use_cache
1171
+ bos_token_id, eos_token_id = generation_config.bos_token_id, generation_config.eos_token_id
1172
+
1173
+ if isinstance(eos_token_id, int):
1174
+ eos_token_id = [eos_token_id]
1175
+ eos_token_id_tensor = torch.tensor(eos_token_id).to(input_ids.device) if eos_token_id is not None else None
1176
+
1177
+ has_default_max_length = kwargs.get("max_length") is None and generation_config.max_length is not None
1178
+ if has_default_max_length and generation_config.max_new_tokens is None:
1179
+ warnings.warn(
1180
+ f"Using `max_length`'s default ({generation_config.max_length}) to control the generation length. "
1181
+ "This behaviour is deprecated and will be removed from the config in v5 of Transformers -- we"
1182
+ " recommend using `max_new_tokens` to control the maximum length of the generation.",
1183
+ UserWarning,
1184
+ )
1185
+ elif generation_config.max_new_tokens is not None:
1186
+ generation_config.max_length = generation_config.max_new_tokens + input_ids_seq_length
1187
+ if not has_default_max_length:
1188
+ logger.warn(
1189
+ f"Both `max_new_tokens` (={generation_config.max_new_tokens}) and `max_length`(="
1190
+ f"{generation_config.max_length}) seem to have been set. `max_new_tokens` will take precedence. "
1191
+ "Please refer to the documentation for more information. "
1192
+ "(https://huggingface.co/docs/transformers/main/en/main_classes/text_generation)",
1193
+ UserWarning,
1194
+ )
1195
+
1196
+ if input_ids_seq_length >= generation_config.max_length:
1197
+ input_ids_string = "decoder_input_ids" if self.config.is_encoder_decoder else "input_ids"
1198
+ logger.warning(
1199
+ f"Input length of {input_ids_string} is {input_ids_seq_length}, but `max_length` is set to"
1200
+ f" {generation_config.max_length}. This can lead to unexpected behavior. You should consider"
1201
+ " increasing `max_new_tokens`."
1202
+ )
1203
+
1204
+ # 2. Set generation parameters if not already defined
1205
+ logits_processor = logits_processor if logits_processor is not None else LogitsProcessorList()
1206
+ stopping_criteria = stopping_criteria if stopping_criteria is not None else StoppingCriteriaList()
1207
+
1208
+ logits_processor = self._get_logits_processor(
1209
+ generation_config=generation_config,
1210
+ input_ids_seq_length=input_ids_seq_length,
1211
+ encoder_input_ids=input_ids,
1212
+ prefix_allowed_tokens_fn=prefix_allowed_tokens_fn,
1213
+ logits_processor=logits_processor,
1214
+ )
1215
+
1216
+ stopping_criteria = self._get_stopping_criteria(
1217
+ generation_config=generation_config, stopping_criteria=stopping_criteria
1218
+ )
1219
+ logits_warper = self._get_logits_warper(generation_config)
1220
+
1221
+ unfinished_sequences = input_ids.new(input_ids.shape[0]).fill_(1)
1222
+ scores = None
1223
+ while True:
1224
+ model_inputs = self.prepare_inputs_for_generation(input_ids, **model_kwargs)
1225
+ # forward pass to get next token
1226
+ outputs = self(
1227
+ **model_inputs,
1228
+ return_dict=True,
1229
+ output_attentions=False,
1230
+ output_hidden_states=False,
1231
+ )
1232
+
1233
+ next_token_logits = outputs.logits[:, -1, :]
1234
+
1235
+ # pre-process distribution
1236
+ next_token_scores = logits_processor(input_ids, next_token_logits)
1237
+ next_token_scores = logits_warper(input_ids, next_token_scores)
1238
+
1239
+ # sample
1240
+ probs = nn.functional.softmax(next_token_scores, dim=-1)
1241
+ if generation_config.do_sample:
1242
+ next_tokens = torch.multinomial(probs, num_samples=1).squeeze(1)
1243
+ else:
1244
+ next_tokens = torch.argmax(probs, dim=-1)
1245
+ # update generated ids, model inputs, and length for next step
1246
+ input_ids = torch.cat([input_ids, next_tokens[:, None]], dim=-1)
1247
+ model_kwargs = self._update_model_kwargs_for_generation(
1248
+ outputs, model_kwargs, is_encoder_decoder=self.config.is_encoder_decoder
1249
+ )
1250
+ unfinished_sequences = unfinished_sequences.mul(
1251
+ next_tokens.tile(eos_token_id_tensor.shape[0], 1).ne(eos_token_id_tensor.unsqueeze(1)).prod(dim=0)
1252
+ )
1253
+ if return_past_key_values:
1254
+ yield input_ids, outputs.past_key_values
1255
+ else:
1256
+ yield input_ids
1257
+ # stop when each sentence is finished, or if we exceed the maximum length
1258
+ if unfinished_sequences.max() == 0 or stopping_criteria(input_ids, scores):
1259
+ break
1260
+
1261
+
1262
+ class ChatGLMForSequenceClassification(ChatGLMPreTrainedModel):
1263
+ def __init__(self, config: ChatGLMConfig, empty_init=True, device=None):
1264
+ super().__init__(config)
1265
+
1266
+ self.num_labels = config.num_labels
1267
+ self.transformer = ChatGLMModel(config, empty_init=empty_init, device=device)
1268
+
1269
+ self.classifier_head = nn.Linear(config.hidden_size, config.num_labels, bias=True, dtype=config.torch_dtype)
1270
+ if config.classifier_dropout is not None:
1271
+ self.dropout = nn.Dropout(config.classifier_dropout)
1272
+ else:
1273
+ self.dropout = None
1274
+ self.config = config
1275
+
1276
+ def forward(
1277
+ self,
1278
+ input_ids: Optional[torch.LongTensor] = None,
1279
+ position_ids: Optional[torch.LongTensor] = None,
1280
+ attention_mask: Optional[torch.Tensor] = None,
1281
+ full_attention_mask: Optional[torch.Tensor] = None,
1282
+ past_key_values: Optional[Tuple[Tuple[torch.Tensor, torch.Tensor], ...]] = None,
1283
+ inputs_embeds: Optional[torch.LongTensor] = None,
1284
+ labels: Optional[torch.LongTensor] = None,
1285
+ use_cache: Optional[bool] = None,
1286
+ output_attentions: Optional[bool] = None,
1287
+ output_hidden_states: Optional[bool] = None,
1288
+ return_dict: Optional[bool] = None,
1289
+ ) -> Union[Tuple[torch.Tensor, ...], SequenceClassifierOutputWithPast]:
1290
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
1291
+
1292
+ transformer_outputs = self.transformer(
1293
+ input_ids=input_ids,
1294
+ position_ids=position_ids,
1295
+ attention_mask=attention_mask,
1296
+ full_attention_mask=full_attention_mask,
1297
+ past_key_values=past_key_values,
1298
+ inputs_embeds=inputs_embeds,
1299
+ use_cache=use_cache,
1300
+ output_attentions=output_attentions,
1301
+ output_hidden_states=output_hidden_states,
1302
+ return_dict=return_dict,
1303
+ )
1304
+
1305
+ hidden_states = transformer_outputs[0]
1306
+ pooled_hidden_states = hidden_states[:, -1]
1307
+ if self.dropout is not None:
1308
+ pooled_hidden_states = self.dropout(pooled_hidden_states)
1309
+ logits = self.classifier_head(pooled_hidden_states)
1310
+
1311
+ loss = None
1312
+ if labels is not None:
1313
+ if self.config.problem_type is None:
1314
+ if self.num_labels == 1:
1315
+ self.config.problem_type = "regression"
1316
+ elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
1317
+ self.config.problem_type = "single_label_classification"
1318
+ else:
1319
+ self.config.problem_type = "multi_label_classification"
1320
+
1321
+ if self.config.problem_type == "regression":
1322
+ loss_fct = MSELoss()
1323
+ if self.num_labels == 1:
1324
+ loss = loss_fct(logits.squeeze().float(), labels.squeeze())
1325
+ else:
1326
+ loss = loss_fct(logits.float(), labels)
1327
+ elif self.config.problem_type == "single_label_classification":
1328
+ loss_fct = CrossEntropyLoss()
1329
+ loss = loss_fct(logits.view(-1, self.num_labels).float(), labels.view(-1))
1330
+ elif self.config.problem_type == "multi_label_classification":
1331
+ loss_fct = BCEWithLogitsLoss()
1332
+ loss = loss_fct(logits.float(), labels.view(-1, self.num_labels))
1333
+
1334
+ if not return_dict:
1335
+ output = (logits,) + transformer_outputs[1:]
1336
+ return ((loss,) + output) if loss is not None else output
1337
+
1338
+ return SequenceClassifierOutputWithPast(
1339
+ loss=loss,
1340
+ logits=logits,
1341
+ past_key_values=transformer_outputs.past_key_values,
1342
+ hidden_states=transformer_outputs.hidden_states,
1343
+ attentions=transformer_outputs.attentions,
1344
+ )
tokenization_chatglm.py ADDED
@@ -0,0 +1,395 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import base64
2
+ import json
3
+ import os
4
+ from typing import List, Optional, Union, Dict, Any
5
+
6
+ import regex as re
7
+ import tiktoken
8
+ from torch import TensorType
9
+ from transformers import PreTrainedTokenizer
10
+ from transformers.tokenization_utils_base import EncodedInput, BatchEncoding
11
+ from transformers.utils import PaddingStrategy
12
+
13
+
14
+ class ChatGLM4Tokenizer(PreTrainedTokenizer):
15
+ vocab_files_names = {"vocab_file": "tokenizer.model"}
16
+ model_input_names = ["input_ids", "attention_mask", "position_ids"]
17
+
18
+ def __init__(
19
+ self,
20
+ vocab_file,
21
+ padding_side="left",
22
+ clean_up_tokenization_spaces=False,
23
+ encode_special_tokens=False,
24
+ **kwargs
25
+ ):
26
+ self.name = "GLM4Tokenizer"
27
+ self.vocab_file = vocab_file
28
+ pat_str = "(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\\r\\n\\p{L}\\p{N}]?\\p{L}+|\\p{N}{1,3}| ?[^\\s\\p{L}\\p{N}]+[\\r\\n]*|\\s*[\\r\\n]+|\\s+(?!\\S)|\\s+"
29
+ self.pat_str = re.compile(pat_str)
30
+ self.encode_special_tokens = encode_special_tokens
31
+
32
+ mergeable_ranks = {}
33
+ with open(vocab_file) as f:
34
+ for line in f:
35
+ token, rank = line.strip().split()
36
+ rank = int(rank)
37
+ token = base64.b64decode(token)
38
+ mergeable_ranks[token] = rank
39
+
40
+ self.mergeable_ranks = mergeable_ranks
41
+
42
+ self.tokenizer = tiktoken.Encoding(
43
+ name="my_tokenizer",
44
+ pat_str=pat_str,
45
+ mergeable_ranks=mergeable_ranks,
46
+ special_tokens={v.content: int(k) for k, v in kwargs['added_tokens_decoder'].items()}
47
+ # special_tokens={}
48
+ )
49
+ self.decoder = {rank: token for token, rank in mergeable_ranks.items()}
50
+ self.n_words = len(self.decoder)
51
+
52
+ super().__init__(
53
+ padding_side=padding_side,
54
+ clean_up_tokenization_spaces=clean_up_tokenization_spaces,
55
+ **kwargs
56
+ )
57
+
58
+ @property
59
+ def vocab_size(self):
60
+ return self.n_words
61
+
62
+ def get_vocab(self):
63
+ """ Returns vocab as a dict """
64
+ vocab = {self._convert_id_to_token(i): i for i in range(self.vocab_size)}
65
+ vocab.update(self.added_tokens_encoder)
66
+ return vocab
67
+
68
+ @staticmethod
69
+ def convert_tokens_to_string(tokens: List[Union[bytes, str]]) -> str:
70
+ """
71
+ Converts a sequence of tokens in a single string.
72
+ """
73
+ text = ""
74
+ temp = b""
75
+ for t in tokens:
76
+ if isinstance(t, str):
77
+ if temp:
78
+ text += temp.decode("utf-8", errors="replace")
79
+ temp = b""
80
+ text += t
81
+ elif isinstance(t, bytes):
82
+ temp += t
83
+ else:
84
+ raise TypeError("token should only be of type types or str")
85
+ if temp:
86
+ text += temp.decode("utf-8", errors="replace")
87
+ return text
88
+
89
+ def _tokenize(self, text, **kwargs):
90
+ tokens = []
91
+ ids = self.tokenizer.encode(text)
92
+ for t in ids:
93
+ tokens.append(self.decoder[t])
94
+ return tokens
95
+
96
+ def _convert_token_to_id(self, token):
97
+ """ Converts a token (str) in an id using the vocab. """
98
+ return self.mergeable_ranks[token]
99
+
100
+ def _convert_id_to_token(self, index):
101
+ """Converts an index (integer) in a token (str) using the vocab."""
102
+ return self.decoder.get(index, "")
103
+
104
+ def save_vocabulary(self, save_directory, filename_prefix=None):
105
+ """
106
+ Save the vocabulary and special tokens file to a directory.
107
+
108
+ Args:
109
+ save_directory (`str`):
110
+ The directory in which to save the vocabulary.
111
+ filename_prefix (`str`, *optional*):
112
+ An optional prefix to add to the named of the saved files.
113
+
114
+ Returns:
115
+ `Tuple(str)`: Paths to the files saved.
116
+ """
117
+ if os.path.isdir(save_directory):
118
+ vocab_file = os.path.join(
119
+ save_directory, self.vocab_files_names["vocab_file"]
120
+ )
121
+ else:
122
+ vocab_file = save_directory
123
+
124
+ with open(self.vocab_file, 'rb') as fin:
125
+ proto_str = fin.read()
126
+
127
+ with open(vocab_file, "wb") as writer:
128
+ writer.write(proto_str)
129
+
130
+ return (vocab_file,)
131
+
132
+ def get_prefix_tokens(self):
133
+ prefix_tokens = [self.convert_tokens_to_ids("[gMASK]"), self.convert_tokens_to_ids("<sop>")]
134
+ return prefix_tokens
135
+
136
+ def apply_chat_template(
137
+ self,
138
+ conversation: Union[List[Dict[str, str]], List[List[Dict[str, str]]]],
139
+ add_generation_prompt: bool = False,
140
+ tokenize: bool = True,
141
+ padding: bool = False,
142
+ truncation: bool = False,
143
+ max_length: Optional[int] = None,
144
+ return_tensors: Optional[Union[str, TensorType]] = None,
145
+ return_dict: bool = False,
146
+ tokenizer_kwargs: Optional[Dict[str, Any]] = None,
147
+ add_special_tokens: bool = True,
148
+ **kwargs,
149
+ ) -> Union[str, List[int], List[str], List[List[int]], BatchEncoding]:
150
+
151
+ if return_dict and not tokenize:
152
+ raise ValueError(
153
+ "`return_dict=True` is incompatible with `tokenize=False`, because there is no dict "
154
+ "of tokenizer outputs to return."
155
+ )
156
+
157
+ def handle_single_conversation(messages):
158
+ content = "你是一位智能编程助手,你叫CodeGeeX。你会为用户回答关于编程、代码、计算机方面的任何问题,并提供格式规范、可以执行、准确安全的代码,并在必要时提供详细的解释。"
159
+ input_message = self.build_single_message("system", "", content)
160
+ for item in messages:
161
+ role = item.get("role", "")
162
+ if not role:
163
+ raise ValueError("Invalid conversation format, 'role' must be given")
164
+ # function call
165
+ elif role == "tool":
166
+ content = self.build_function_sys_prompt(item["content"])
167
+ input_message = self.build_single_message("system", "", content)
168
+ # chat
169
+ elif role == "system":
170
+ input_message = self.build_single_message("system", item.get("metadata", ""), item["content"])
171
+ else:
172
+ input_message += self.build_single_message(item["role"], item.get("metadata", ""), item["content"])
173
+
174
+ if add_generation_prompt:
175
+ input_message += "<|assistant|>\n"
176
+ if tokenize:
177
+ input_ids = self.get_prefix_tokens() if add_special_tokens else []
178
+ input_ids += self.tokenizer.encode(input_message, allowed_special='all', disallowed_special=set())
179
+ return input_ids
180
+ else:
181
+ return input_message
182
+
183
+ # Main logic to handle different conversation formats
184
+ if isinstance(conversation, list) and all(isinstance(i, dict) for i in conversation):
185
+ result = handle_single_conversation(conversation)
186
+ elif isinstance(conversation, list) and all(isinstance(i, list) for i in conversation):
187
+ result = [handle_single_conversation(c) for c in conversation]
188
+ elif hasattr(conversation, "messages"):
189
+ result = handle_single_conversation(conversation.messages)
190
+ else:
191
+ raise ValueError("Invalid conversation format")
192
+
193
+ if tokenize:
194
+ output = self.batch_encode_plus(
195
+ [result] if isinstance(result[0], int) else result,
196
+ padding=padding,
197
+ truncation=truncation,
198
+ max_length=max_length,
199
+ return_tensors=return_tensors,
200
+ is_split_into_words=True,
201
+ add_special_tokens=False
202
+ )
203
+ if return_dict:
204
+ return output
205
+ else:
206
+ return output["input_ids"]
207
+ else:
208
+ return result
209
+
210
+ def build_inputs_with_special_tokens(
211
+ self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None
212
+ ) -> List[int]:
213
+ """
214
+ Build model inputs from a sequence or a pair of sequence for sequence classification tasks by concatenating and
215
+ adding special tokens. A BERT sequence has the following format:
216
+
217
+ - single sequence: `[CLS] X [SEP]`
218
+ - pair of sequences: `[CLS] A [SEP] B [SEP]`
219
+
220
+ Args:
221
+ token_ids_0 (`List[int]`):
222
+ List of IDs to which the special tokens will be added.
223
+ token_ids_1 (`List[int]`, *optional*):
224
+ Optional second list of IDs for sequence pairs.
225
+
226
+ Returns:
227
+ `List[int]`: List of [input IDs](../glossary#input-ids) with the appropriate special tokens.
228
+ """
229
+ prefix_tokens = self.get_prefix_tokens()
230
+ token_ids_0 = prefix_tokens + token_ids_0
231
+ if token_ids_1 is not None:
232
+ token_ids_0 = token_ids_0 + token_ids_1 + [self.convert_tokens_to_ids("<eos>")]
233
+ return token_ids_0
234
+
235
+ def _pad(
236
+ self,
237
+ encoded_inputs: Union[Dict[str, EncodedInput], BatchEncoding],
238
+ max_length: Optional[int] = None,
239
+ padding_strategy: PaddingStrategy = PaddingStrategy.DO_NOT_PAD,
240
+ pad_to_multiple_of: Optional[int] = None,
241
+ return_attention_mask: Optional[bool] = None,
242
+ ) -> dict:
243
+ """
244
+ Pad encoded inputs (on left/right and up to predefined length or max length in the batch)
245
+
246
+ Args:
247
+ encoded_inputs:
248
+ Dictionary of tokenized inputs (`List[int]`) or batch of tokenized inputs (`List[List[int]]`).
249
+ max_length: maximum length of the returned list and optionally padding length (see below).
250
+ Will truncate by taking into account the special tokens.
251
+ padding_strategy: PaddingStrategy to use for padding.
252
+
253
+ - PaddingStrategy.LONGEST Pad to the longest sequence in the batch
254
+ - PaddingStrategy.MAX_LENGTH: Pad to the max length (default)
255
+ - PaddingStrategy.DO_NOT_PAD: Do not pad
256
+ The tokenizer padding sides are defined in self.padding_side:
257
+
258
+ - 'left': pads on the left of the sequences
259
+ - 'right': pads on the right of the sequences
260
+ pad_to_multiple_of: (optional) Integer if set will pad the sequence to a multiple of the provided value.
261
+ This is especially useful to enable the use of Tensor Core on NVIDIA hardware with compute capability
262
+ `>= 7.5` (Volta).
263
+ return_attention_mask:
264
+ (optional) Set to 'False' to avoid returning attention mask (default: set to model specifics)
265
+ """
266
+ # Load from model defaults
267
+ assert self.padding_side == "left"
268
+
269
+ required_input = encoded_inputs[self.model_input_names[0]]
270
+ seq_length = len(required_input)
271
+
272
+ if padding_strategy == PaddingStrategy.LONGEST:
273
+ max_length = len(required_input)
274
+
275
+ if max_length is not None and pad_to_multiple_of is not None and (max_length % pad_to_multiple_of != 0):
276
+ max_length = ((max_length // pad_to_multiple_of) + 1) * pad_to_multiple_of
277
+
278
+ needs_to_be_padded = padding_strategy != PaddingStrategy.DO_NOT_PAD and len(required_input) != max_length
279
+
280
+ # Initialize attention mask if not present.
281
+ if "attention_mask" not in encoded_inputs:
282
+ encoded_inputs["attention_mask"] = [1] * seq_length
283
+
284
+ if "position_ids" not in encoded_inputs:
285
+ encoded_inputs["position_ids"] = list(range(seq_length))
286
+
287
+ if needs_to_be_padded:
288
+ difference = max_length - len(required_input)
289
+
290
+ if "attention_mask" in encoded_inputs:
291
+ encoded_inputs["attention_mask"] = [0] * difference + encoded_inputs["attention_mask"]
292
+ if "position_ids" in encoded_inputs:
293
+ encoded_inputs["position_ids"] = [0] * difference + encoded_inputs["position_ids"]
294
+ encoded_inputs[self.model_input_names[0]] = [self.pad_token_id] * difference + required_input
295
+
296
+ return encoded_inputs
297
+
298
+ @staticmethod
299
+ def build_single_message(role, metadata, message):
300
+ assert role in ["system", "user", "assistant", "observation"], role
301
+ return f"<|{role}|>{metadata}\n{message}"
302
+
303
+ @staticmethod
304
+ def build_function_sys_prompt(item: dict) -> str:
305
+ prompt = """
306
+ 你将接收到一个用户提出的问题,并请撰写清晰、简洁且准确的答案。
307
+
308
+ # Note
309
+ - 我将给你提供一些函数工具的接口信息,包括函数的定义、用途、名字、参数名和参数类型。
310
+ - 请根据这些信息,为用户的指令,从中选择最合适的函数,并给出调用时需要使用的参数。
311
+ - **返回类型为一个json格式的字符串,包含函数名和参数字典。**
312
+ - name: 函数名
313
+ - arguments: 参数字典,其中key为参数名,value为参数类型。
314
+ - **只需要生成答案即可,无需在你的回答之前或之后做出解释,也不要直接回答用户的问题。**
315
+ - 只用当提供的函数工具不足以完成任务时,请你用正常的语气告知用户并解释原因。
316
+
317
+ # Functions
318
+ 以下是可使用的函数工具的接口信息。
319
+ """.lstrip()
320
+
321
+ if isinstance(item['function'], dict):
322
+ func = item['function']
323
+ prompt += f"\n## Function 1\n"
324
+ prompt += f"\n### Name\n{func['name']}\n"
325
+ prompt += f"\n### Description\n{func['description']}\n"
326
+ prompt += f"\n### Parameters\n```json\n{json.dumps(func['parameters'], ensure_ascii=False)}\n```\n"
327
+ return prompt
328
+ elif isinstance(item['function'], list):
329
+ for idx, func in enumerate(item['function']):
330
+ prompt += f"\n## Function {idx + 1}\n"
331
+ prompt += f"\n### Name\n{func['name']}\n"
332
+ prompt += f"\n### Description\n{func['description']}\n"
333
+ prompt += f"\n### Parameters\n```json\n{json.dumps(func['parameters'], ensure_ascii=False)}\n```\n"
334
+ return prompt
335
+
336
+ def apply_infilling_template(
337
+ self,
338
+ message: dict,
339
+ add_generation_prompt: bool = False,
340
+ tokenize: bool = True,
341
+ padding: bool = False,
342
+ truncation: bool = False,
343
+ max_length: Optional[int] = None,
344
+ return_tensors: Optional[Union[str, TensorType]] = None,
345
+ return_dict: bool = False,
346
+ add_special_tokens: bool = True,
347
+ ) -> Union[str, List[int], List[str], List[List[int]], BatchEncoding]:
348
+ if return_dict and not tokenize:
349
+ raise ValueError(
350
+ "`return_dict=True` is incompatible with `tokenize=False`, because there is no dict "
351
+ "of tokenizer outputs to return."
352
+ )
353
+
354
+ if not isinstance(message, dict):
355
+ raise ValueError("Invalid conversation format")
356
+ content = self.build_infilling_prompt(message)
357
+ input_message = self.build_single_message("user", "", content)
358
+ if add_generation_prompt:
359
+ input_message += "<|assistant|>\n"
360
+ if not tokenize:
361
+ return input_message
362
+
363
+ input_ids = self.get_prefix_tokens() if add_special_tokens else []
364
+ input_ids += self.tokenizer.encode(input_message, allowed_special='all', disallowed_special=set())
365
+ output = self.batch_encode_plus(
366
+ [input_ids] if isinstance(input_ids[0], int) else input_ids,
367
+ padding=padding,
368
+ truncation=truncation,
369
+ max_length=max_length,
370
+ return_tensors=return_tensors,
371
+ is_split_into_words=True,
372
+ add_special_tokens=False
373
+ )
374
+ if return_dict:
375
+ return output
376
+ else:
377
+ return output["input_ids"]
378
+
379
+ @staticmethod
380
+ def build_infilling_prompt(item: dict) -> str:
381
+ prompt = ""
382
+ if "path" in item:
383
+ prompt += f"###PATH:{item['path']}\n"
384
+ if "language" in item:
385
+ prompt += f"###LANGUAGE:{item['language']}\n"
386
+ elif "lang" in item:
387
+ prompt += f"###LANGUAGE:{item['lang']}\n"
388
+ if "mode" in item and item['mode'].lower() == "line":
389
+ prompt += "###MODE:LINE\n"
390
+ else:
391
+ prompt += "###MODE:BLOCK\n"
392
+ prompt += f"<|code_suffix|>{item['suffix']}"
393
+ prompt += f"<|code_prefix|>{item['prefix']}"
394
+ prompt += "<|code_middle|>"
395
+ return prompt
tokenizer.json ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:b3bbfa7d61e43c6ce23a59c2fd2c8a4365684a48d69d33dafa67e75e9a884c93
3
+ size 10743038
tokenizer.model ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:5a493598071550244b2ee7f26118f3edec2150b9dfa967929a99052ac83fe716
3
+ size 2623634
tokenizer_config.json ADDED
@@ -0,0 +1,165 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "auto_map": {
3
+ "AutoTokenizer": [
4
+ "tokenization_chatglm.ChatGLM4Tokenizer",
5
+ null
6
+ ]
7
+ },
8
+ "added_tokens_decoder": {
9
+ "151329": {
10
+ "content": "<|endoftext|>",
11
+ "lstrip": false,
12
+ "normalized": false,
13
+ "rstrip": false,
14
+ "single_word": false,
15
+ "special": true
16
+ },
17
+ "151330": {
18
+ "content": "[MASK]",
19
+ "lstrip": false,
20
+ "normalized": false,
21
+ "rstrip": false,
22
+ "single_word": false,
23
+ "special": true
24
+ },
25
+ "151331": {
26
+ "content": "[gMASK]",
27
+ "lstrip": false,
28
+ "normalized": false,
29
+ "rstrip": false,
30
+ "single_word": false,
31
+ "special": true
32
+ },
33
+ "151332": {
34
+ "content": "[sMASK]",
35
+ "lstrip": false,
36
+ "normalized": false,
37
+ "rstrip": false,
38
+ "single_word": false,
39
+ "special": true
40
+ },
41
+ "151333": {
42
+ "content": "<sop>",
43
+ "lstrip": false,
44
+ "normalized": false,
45
+ "rstrip": false,
46
+ "single_word": false,
47
+ "special": true
48
+ },
49
+ "151334": {
50
+ "content": "<eop>",
51
+ "lstrip": false,
52
+ "normalized": false,
53
+ "rstrip": false,
54
+ "single_word": false,
55
+ "special": true
56
+ },
57
+ "151335": {
58
+ "content": "<|system|>",
59
+ "lstrip": false,
60
+ "normalized": false,
61
+ "rstrip": false,
62
+ "single_word": false,
63
+ "special": true
64
+ },
65
+ "151336": {
66
+ "content": "<|user|>",
67
+ "lstrip": false,
68
+ "normalized": false,
69
+ "rstrip": false,
70
+ "single_word": false,
71
+ "special": true
72
+ },
73
+ "151337": {
74
+ "content": "<|assistant|>",
75
+ "lstrip": false,
76
+ "normalized": false,
77
+ "rstrip": false,
78
+ "single_word": false,
79
+ "special": true
80
+ },
81
+ "151338": {
82
+ "content": "<|observation|>",
83
+ "lstrip": false,
84
+ "normalized": false,
85
+ "rstrip": false,
86
+ "single_word": false,
87
+ "special": true
88
+ },
89
+ "151339": {
90
+ "content": "<|begin_of_image|>",
91
+ "lstrip": false,
92
+ "normalized": false,
93
+ "rstrip": false,
94
+ "single_word": false,
95
+ "special": true
96
+ },
97
+ "151340": {
98
+ "content": "<|end_of_image|>",
99
+ "lstrip": false,
100
+ "normalized": false,
101
+ "rstrip": false,
102
+ "single_word": false,
103
+ "special": true
104
+ },
105
+ "151341": {
106
+ "content": "<|begin_of_video|>",
107
+ "lstrip": false,
108
+ "normalized": false,
109
+ "rstrip": false,
110
+ "single_word": false,
111
+ "special": true
112
+ },
113
+ "151342": {
114
+ "content": "<|end_of_video|>",
115
+ "lstrip": false,
116
+ "normalized": false,
117
+ "rstrip": false,
118
+ "single_word": false,
119
+ "special": true
120
+ },
121
+ "151343": {
122
+ "content": "<|code_prefix|>",
123
+ "lstrip": false,
124
+ "normalized": false,
125
+ "rstrip": false,
126
+ "single_word": false,
127
+ "special": false
128
+ },
129
+ "151344": {
130
+ "content": "<|code_middle|>",
131
+ "lstrip": false,
132
+ "normalized": false,
133
+ "rstrip": false,
134
+ "single_word": false,
135
+ "special": false
136
+ },
137
+ "151345": {
138
+ "content": "<|code_suffix|>",
139
+ "lstrip": false,
140
+ "normalized": false,
141
+ "rstrip": false,
142
+ "single_word": false,
143
+ "special": false
144
+ },
145
+ "151346": {
146
+ "content": "<|cursor|>",
147
+ "lstrip": false,
148
+ "normalized": false,
149
+ "rstrip": false,
150
+ "single_word": false,
151
+ "special": false
152
+ }
153
+ },
154
+ "additional_special_tokens": ["<|endoftext|>", "[MASK]", "[gMASK]", "[sMASK]", "<sop>", "<eop>", "<|system|>",
155
+ "<|user|>", "<|assistant|>", "<|observation|>", "<|begin_of_image|>", "<|end_of_image|>",
156
+ "<|begin_of_video|>", "<|end_of_video|>", "<|code_prefix|>", "<|code_middle|>", "<|code_suffix|>", "<|cursor|>"],
157
+ "clean_up_tokenization_spaces": false,
158
+ "do_lower_case": false,
159
+ "eos_token": "<|endoftext|>",
160
+ "pad_token": "<|endoftext|>",
161
+ "model_max_length": 1000000000000000019884624838656,
162
+ "padding_side": "left",
163
+ "remove_space": false,
164
+ "tokenizer_class": "ChatGLM4Tokenizer"
165
+ }