add chat_template, fix metadata

.gitignore

@@ -1,3 +0,0 @@
-*venv
-*.DS_Store
-*.idea/

README.md

@@ -3,31 +3,27 @@ license: other
 license_name: glm-4
 license_link: https://huggingface.co/THUDM/glm-4-9b-chat/blob/main/LICENSE
 language:
-  - zh
-  - en
+- zh
+- en
 tags:
-  - glm
-  - chatglm
-  - thudm
+- glm
+- chatglm
+- thudm
 inference: false
+pipeline_tag: text-generation
 ---
 
 # GLM-4-9B-Chat
 
-Read this in [English](README_en.md).
+Read this in [English](README_en.md)
 
-**2024/11/25**, 我们建议使用从 `transformers>=4.46.0` 开始，使用 [glm-4-9b-chat-hf](https://huggingface.co/THUDM/glm-4-9b-chat-hf) 以减少后续 transformers 升级导致的兼容性问题。
-
-**2024/08/12, 本仓库代码已更新并使用 `transformers>=4.44.0`, 请及时更新依赖。**
-
-**2024/07/24，我们发布了与长文本相关的最新技术解读，关注 [这里](https://medium.com/@ChatGLM/glm-long-scaling-pre-trained-model-contexts-to-millions-caa3c48dea85) 查看我们在训练 GLM-4-9B 开源模型中关于长文本技术的技术报告**
-
-## 模型介绍
-GLM-4-9B 是智谱 AI 推出的最新一代预训练模型 GLM-4 系列中的开源版本。
-在语义、数学、推理、代码和知识等多方面的数据集测评中，GLM-4-9B 及其人类偏好对齐的版本 GLM-4-9B-Chat 均表现出较高的性能。
-除了能进行多轮对话，GLM-4-9B-Chat 还具备网页浏览、代码执行、自定义工具调用（Function Call）和长文本推理（支持最大 128K
-上下文）等高级功能。
-本代模型增加了多语言支持，支持包括日语，韩语，德语在内的 26 种语言。我们还推出了支持 1M 上下文长度（约 200 万中文字符）的模型。
+GLM-4-9B 是智谱 AI 推出的最新一代预训练模型 GLM-4 系列中的开源版本。 在语义、数学、推理、代码和知识等多方面的数据集测评中，
+**GLM-4-9B** 及其人类偏好对齐的版本 **GLM-4-9B-Chat** 均表现出超越 Llama-3-8B 的卓越性能。除了能进行多轮对话，GLM-4-9B-Chat
+还具备网页浏览、代码执行、自定义工具调用（Function Call）和长文本推理（支持最大 128K 上下文）等高级功能。本代模型增加了多语言支持，支持包括日语，韩语，德语在内的
+26 种语言。我们还推出了支持 1M 上下文长度（约 200 万中文字符）的 **GLM-4-9B-Chat-1M** 模型和基于 GLM-4-9B 的多模态模型
+GLM-4V-9B。**GLM-4V-9B** 具备 1120 * 1120 高分辨率下的中英双语多轮对话能力，在中英文综合能力、感知推理、文字识别、图表理解等多方面多模态评测中，GLM-4V-9B
+表现出超越 GPT-4-turbo-2024-04-09、Gemini
+1.0 Pro、Qwen-VL-Max 和 Claude 3 Opus 的卓越性能。
 
 ## 评测结果
 
@@ -53,14 +49,14 @@ GLM-4-9B 是智谱 AI 推出的最新一代预训练模型 GLM-4 系列中的开
 
 在六个多语言数据集上对 GLM-4-9B-Chat 和 Llama-3-8B-Instruct 进行了测试，测试结果及数据集对应选取语言如下表
 
-| Dataset     | Llama-3-8B-Instruct | GLM-4-9B-Chat |                                           Languages                                            
+| Dataset     | Llama-3-8B-Instruct | GLM-4-9B-Chat |                                           Languages                                            |
 |:------------|:-------------------:|:-------------:|:----------------------------------------------------------------------------------------------:|
-| M-MMLU      |        49.6         |     56.6      |                                              all                                               
-| FLORES      |        25.0         |     28.8      | ru, es, de, fr, it, pt, pl, ja, nl, ar, tr, cs, vi, fa, hu, el, ro, sv, uk, fi, ko, da, bg, no 
-| MGSM        |        54.0         |     65.3      |                           zh, en, bn, de, es, fr, ja, ru, sw, te, th                           
-| XWinograd   |        61.7         |     73.1      |                                     zh, en, fr, jp, ru, pt                                     
-| XStoryCloze |        84.7         |     90.7      |                           zh, en, ar, es, eu, hi, id, my, ru, sw, te                           
-| XCOPA       |        73.3         |     80.1      |                           zh, et, ht, id, it, qu, sw, ta, th, tr, vi                           
+| M-MMLU      |        49.6         |     56.6      |                                              all                                               |
+| FLORES      |        25.0         |     28.8      | ru, es, de, fr, it, pt, pl, ja, nl, ar, tr, cs, vi, fa, hu, el, ro, sv, uk, fi, ko, da, bg, no |
+| MGSM        |        54.0         |     65.3      |                           zh, en, bn, de, es, fr, ja, ru, sw, te, th                           |
+| XWinograd   |        61.7         |     73.1      |                                     zh, en, fr, jp, ru, pt                                     |
+| XStoryCloze |        84.7         |     90.7      |                           zh, en, ar, es, eu, hi, id, my, ru, sw, te                           |
+| XCOPA       |        73.3         |     80.1      |                           zh, et, ht, id, it, qu, sw, ta, th, tr, vi                           |
 
 ### 工具调用能力
 
@@ -78,12 +74,7 @@ GLM-4-9B 是智谱 AI 推出的最新一代预训练模型 GLM-4 系列中的开
 
 ## 运行模型
 
-**更多推理代码和依赖信息，请访问我们的 [github](https://github.com/THUDM/GLM-4)。**
-
-**请严格按照[依赖](https://github.com/THUDM/GLM-4/blob/main/basic_demo/requirements.txt)安装，否则无法正常运行。**
-
-### 使用 transformers 后端进行推理:
-
+使用 transformers 后端进行推理:
 
 ```python
 import torch
@@ -161,12 +152,20 @@ GLM-4 模型的权重的使用则需要遵循 [LICENSE](LICENSE)。
 如果你觉得我们的工作有帮助的话，请考虑引用下列论文。
 
 ```
-@misc{glm2024chatglm,
-      title={ChatGLM: A Family of Large Language Models from GLM-130B to GLM-4 All Tools}, 
-      author={Team GLM and Aohan Zeng and Bin Xu and Bowen Wang and Chenhui Zhang and Da Yin and Diego Rojas and Guanyu Feng and Hanlin Zhao and Hanyu Lai and Hao Yu and Hongning Wang and Jiadai Sun and Jiajie Zhang and Jiale Cheng and Jiayi Gui and Jie Tang and Jing Zhang and Juanzi Li and Lei Zhao and Lindong Wu and Lucen Zhong and Mingdao Liu and Minlie Huang and Peng Zhang and Qinkai Zheng and Rui Lu and Shuaiqi Duan and Shudan Zhang and Shulin Cao and Shuxun Yang and Weng Lam Tam and Wenyi Zhao and Xiao Liu and Xiao Xia and Xiaohan Zhang and Xiaotao Gu and Xin Lv and Xinghan Liu and Xinyi Liu and Xinyue Yang and Xixuan Song and Xunkai Zhang and Yifan An and Yifan Xu and Yilin Niu and Yuantao Yang and Yueyan Li and Yushi Bai and Yuxiao Dong and Zehan Qi and Zhaoyu Wang and Zhen Yang and Zhengxiao Du and Zhenyu Hou and Zihan Wang},
-      year={2024},
-      eprint={2406.12793},
-      archivePrefix={arXiv},
-      primaryClass={id='cs.CL' full_name='Computation and Language' is_active=True alt_name='cmp-lg' in_archive='cs' is_general=False description='Covers natural language processing. Roughly includes material in ACM Subject Class I.2.7. Note that work on artificial languages (programming languages, logics, formal systems) that does not explicitly address natural-language issues broadly construed (natural-language processing, computational linguistics, speech, text retrieval, etc.) is not appropriate for this area.'}
+@article{zeng2022glm,
+  title={Glm-130b: An open bilingual pre-trained model},
+  author={Zeng, Aohan and Liu, Xiao and Du, Zhengxiao and Wang, Zihan and Lai, Hanyu and Ding, Ming and Yang, Zhuoyi and Xu, Yifan and Zheng, Wendi and Xia, Xiao and others},
+  journal={arXiv preprint arXiv:2210.02414},
+  year={2022}
 }
 ```
+
+```
+@inproceedings{du2022glm,
+  title={GLM: General Language Model Pretraining with Autoregressive Blank Infilling},
+  author={Du, Zhengxiao and Qian, Yujie and Liu, Xiao and Ding, Ming and Qiu, Jiezhong and Yang, Zhilin and Tang, Jie},
+  booktitle={Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers)},
+  pages={320--335},
+  year={2022}
+}
+```

README_en.md

@@ -1,18 +1,12 @@
 # GLM-4-9B-Chat
 
-**2024/08/12, The repository code has been updated and now requires `transformers>=4.44.0`. Please update your dependencies accordingly.**
-
-**On July 24, 2024, we released the latest technical interpretation related to long texts. Check
-out [here](https://medium.com/@ChatGLM/glm-long-scaling-pre-trained-model-contexts-to-millions-caa3c48dea85) to view our
-technical report on long context technology in the training of the open-source GLM-4-9B model.**
-
 ## Model Introduction
 
 GLM-4-9B is the open-source version of the latest generation of pre-trained models in the GLM-4 series launched by Zhipu
 AI. In the evaluation of data sets in semantics, mathematics, reasoning, code, and knowledge, **GLM-4-9B**
 and its human preference-aligned version **GLM-4-9B-Chat** have shown superior performance beyond Llama-3-8B. In
 addition to multi-round conversations, GLM-4-9B-Chat also has advanced features such as web browsing, code execution,
-custom tool calls (Function Call), and long context
+custom tool calls (Function Call), and long text
 reasoning (supporting up to 128K context). This generation of models has added multi-language support, supporting 26
 languages including Japanese, Korean, and German. We have also launched the **GLM-4-9B-Chat-1M** model that supports 1M
 context length (about 2 million Chinese characters) and the multimodal model GLM-4V-9B based on GLM-4-9B.
@@ -70,12 +64,9 @@ on [Berkeley Function Calling Leaderboard](https://github.com/ShishirPatil/goril
 
 **This repository is the model repository of GLM-4-9B-Chat, supporting `128K` context length.**
 
-## Quick Start
+## Quick call
 
-**For more inference code and requirements, please visit our [github page](https://github.com/THUDM/GLM-4).**
-
-**Please strictly follow the [dependencies](https://github.com/THUDM/GLM-4/blob/main/basic_demo/requirements.txt) to
-install, otherwise it will not run properly**
+**For hardware configuration and system requirements, please check [here](basic_demo/README_en.md).**
 
 ### Use the following method to quickly call the GLM-4-9B-Chat language model
 
@@ -144,6 +135,7 @@ sampling_params = SamplingParams(temperature=0.95, max_tokens=1024, stop_token_i
 
 inputs = tokenizer.apply_chat_template(prompt, tokenize=False, add_generation_prompt=True)
 outputs = llm.generate(prompts=inputs, sampling_params=sampling_params)
+
 print(outputs[0].outputs[0].text)
 ```
 
@@ -156,12 +148,20 @@ The weights of the GLM-4 model are available under the terms of [LICENSE](LICENS
 If you find our work useful, please consider citing the following paper.
 
 ```
-@misc{glm2024chatglm,
-      title={ChatGLM: A Family of Large Language Models from GLM-130B to GLM-4 All Tools}, 
-      author={Team GLM and Aohan Zeng and Bin Xu and Bowen Wang and Chenhui Zhang and Da Yin and Diego Rojas and Guanyu Feng and Hanlin Zhao and Hanyu Lai and Hao Yu and Hongning Wang and Jiadai Sun and Jiajie Zhang and Jiale Cheng and Jiayi Gui and Jie Tang and Jing Zhang and Juanzi Li and Lei Zhao and Lindong Wu and Lucen Zhong and Mingdao Liu and Minlie Huang and Peng Zhang and Qinkai Zheng and Rui Lu and Shuaiqi Duan and Shudan Zhang and Shulin Cao and Shuxun Yang and Weng Lam Tam and Wenyi Zhao and Xiao Liu and Xiao Xia and Xiaohan Zhang and Xiaotao Gu and Xin Lv and Xinghan Liu and Xinyi Liu and Xinyue Yang and Xixuan Song and Xunkai Zhang and Yifan An and Yifan Xu and Yilin Niu and Yuantao Yang and Yueyan Li and Yushi Bai and Yuxiao Dong and Zehan Qi and Zhaoyu Wang and Zhen Yang and Zhengxiao Du and Zhenyu Hou and Zihan Wang},
-      year={2024},
-      eprint={2406.12793},
-      archivePrefix={arXiv},
-      primaryClass={id='cs.CL' full_name='Computation and Language' is_active=True alt_name='cmp-lg' in_archive='cs' is_general=False description='Covers natural language processing. Roughly includes material in ACM Subject Class I.2.7. Note that work on artificial languages (programming languages, logics, formal systems) that does not explicitly address natural-language issues broadly construed (natural-language processing, computational linguistics, speech, text retrieval, etc.) is not appropriate for this area.'}
+@article{zeng2022glm,
+  title={Glm-130b: An open bilingual pre-trained model},
+  author={Zeng, Aohan and Liu, Xiao and Du, Zhengxiao and Wang, Zihan and Lai, Hanyu and Ding, Ming and Yang, Zhuoyi and Xu, Yifan and Zheng, Wendi and Xia, Xiao and others},
+  journal={arXiv preprint arXiv:2210.02414},
+  year={2022}
+}
+```
+
+```
+@inproceedings{du2022glm,
+  title={GLM: General Language Model Pretraining with Autoregressive Blank Infilling},
+  author={Du, Zhengxiao and Qian, Yujie and Liu, Xiao and Ding, Ming and Qiu, Jiezhong and Yang, Zhilin and Tang, Jie},
+  booktitle={Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers)},
+  pages={320--335},
+  year={2022}
 }
 ```

config.json

@@ -1,5 +1,5 @@
 {
-  "_name_or_path": "THUDM/glm-4-9b-chat",
+  "_name_or_path": "THUDM/glm4-9b-chat",
   "model_type": "chatglm",
   "architectures": [
     "ChatGLMModel"
@@ -17,7 +17,6 @@
   "apply_residual_connection_post_layernorm": false,
   "attention_dropout": 0.0,
   "attention_softmax_in_fp32": true,
-  "attn_implementation": "sdpa",
   "bias_dropout_fusion": true,
   "ffn_hidden_size": 13696,
   "fp32_residual_connection": false,
@@ -38,8 +37,9 @@
   "seq_length": 131072,
   "use_cache": true,
   "torch_dtype": "bfloat16",
-  "transformers_version": "4.44.0",
+  "transformers_version": "4.30.2",
   "tie_word_embeddings": false,
   "eos_token_id": [151329, 151336, 151338],
   "pad_token_id": 151329
-}
+}
+

generation_config.json

@@ -9,5 +9,5 @@
   "temperature": 0.8,
   "max_length": 128000,
   "top_p": 0.8,
-  "transformers_version": "4.44.0"
+  "transformers_version": "4.38.2"
 }

modeling_chatglm.py

@@ -1,14 +1,19 @@
 """ PyTorch ChatGLM model. """
-
+import json
 import math
+import copy
+import warnings
+import re
 import sys
+
 import torch
 import torch.utils.checkpoint
 import torch.nn.functional as F
 from torch import nn
 from torch.nn import CrossEntropyLoss, LayerNorm, MSELoss, BCEWithLogitsLoss
 from torch.nn.utils import skip_init
-from typing import Optional, Tuple, Union, List, Dict, Any
+from typing import Optional, Tuple, Union, List, Callable, Dict, Any
+from copy import deepcopy
 
 from transformers.modeling_outputs import (
     BaseModelOutputWithPast,
@@ -16,24 +21,15 @@ from transformers.modeling_outputs import (
     SequenceClassifierOutputWithPast,
 )
 from transformers.modeling_utils import PreTrainedModel
-from transformers.utils import logging, is_torch_npu_available
+from transformers.utils import logging
 from transformers.generation.logits_process import LogitsProcessor
-from transformers.generation.utils import ModelOutput
+from transformers.generation.utils import LogitsProcessorList, StoppingCriteriaList, GenerationConfig, ModelOutput
 
 from .configuration_chatglm import ChatGLMConfig
 
-try:
-    from transformers.utils import is_flash_attn_greater_or_equal_2_10, is_flash_attn_2_available
-
-    if is_flash_attn_2_available():
-        from flash_attn import flash_attn_func, flash_attn_varlen_func
-        from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input  # noqa
-except:
-    pass
-
 # flags required to enable jit fusion kernels
 
-if sys.platform != 'darwin' and not is_torch_npu_available():
+if sys.platform != 'darwin':
     torch._C._jit_set_profiling_mode(False)
     torch._C._jit_set_profiling_executor(False)
     torch._C._jit_override_can_fuse_on_cpu(True)
@@ -44,7 +40,6 @@ logger = logging.get_logger(__name__)
 _CHECKPOINT_FOR_DOC = "THUDM/ChatGLM"
 _CONFIG_FOR_DOC = "ChatGLMConfig"
 
-
 def default_init(cls, *args, **kwargs):
     return cls(*args, **kwargs)
 
@@ -164,13 +159,12 @@ class RMSNorm(torch.nn.Module):
 class CoreAttention(torch.nn.Module):
     def __init__(self, config: ChatGLMConfig, layer_number):
         super(CoreAttention, self).__init__()
-        self.config = config
+
         self.apply_query_key_layer_scaling = config.apply_query_key_layer_scaling
         self.attention_softmax_in_fp32 = config.attention_softmax_in_fp32
         if self.apply_query_key_layer_scaling:
             self.attention_softmax_in_fp32 = True
         self.layer_number = max(1, layer_number)
-        self.is_causal = True
 
         projection_size = config.kv_channels * config.num_attention_heads
 
@@ -189,199 +183,91 @@ class CoreAttention(torch.nn.Module):
         self.attention_dropout = torch.nn.Dropout(config.attention_dropout)
 
     def forward(self, query_layer, key_layer, value_layer, attention_mask):
-        # [b, np, sq, sk]
-        output_size = (query_layer.size(0), query_layer.size(1), query_layer.size(2), key_layer.size(2))
-
-        # [b, np, sq, hn] -> [b * np, sq, hn]
-        query_layer = query_layer.view(output_size[0] * output_size[1], output_size[2], -1)
-        # [b, np, sk, hn] -> [b * np, sk, hn]
-        key_layer = key_layer.view(output_size[0] * output_size[1], output_size[3], -1)
-
-        # preallocting input tensor: [b * np, sq, sk]
-        matmul_input_buffer = torch.empty(
-            output_size[0] * output_size[1], output_size[2], output_size[3], dtype=query_layer.dtype,
-            device=query_layer.device
-        )
-
-        # Raw attention scores. [b * np, sq, sk]
-        matmul_result = torch.baddbmm(
-            matmul_input_buffer,
-            query_layer,  # [b * np, sq, hn]
-            key_layer.transpose(1, 2),  # [b * np, hn, sk]
-            beta=0.0,
-            alpha=(1.0 / self.norm_factor),
-        )
-
-        # change view to [b, np, sq, sk]
-        attention_scores = matmul_result.view(*output_size)
-
-        # ===========================
-        # Attention probs and dropout
-        # ===========================
-
-        # attention scores and attention mask [b, np, sq, sk]
-        if self.attention_softmax_in_fp32:
-            attention_scores = attention_scores.float()
-        if self.coeff is not None:
-            attention_scores = attention_scores * self.coeff
-        if attention_mask is None and attention_scores.shape[2] == attention_scores.shape[3]:
-            attention_mask = torch.ones(output_size[0], 1, output_size[2], output_size[3],
-                                        device=attention_scores.device, dtype=torch.bool)
-            attention_mask.tril_()
-            attention_mask = ~attention_mask
-        if attention_mask is not None:
-            attention_scores = attention_scores.masked_fill(attention_mask, float("-inf"))
-        attention_probs = F.softmax(attention_scores, dim=-1)
-        attention_probs = attention_probs.type_as(value_layer)
-
-        # This is actually dropping out entire tokens to attend to, which might
-        # seem a bit unusual, but is taken from the original Transformer paper.
-        attention_probs = self.attention_dropout(attention_probs)
-
-        # query layer shape: [b * np, sq, hn]
-        # value layer shape: [b, np, sk, hn]
-        # attention shape: [b, np, sq, sk]
-        # context layer shape: [b, np, sq, hn]
-        output_size = (value_layer.size(0), value_layer.size(1), query_layer.size(1), value_layer.size(3))
-        # change view [b * np, sk, hn]
-        value_layer = value_layer.view(output_size[0] * output_size[1], value_layer.size(2), -1)
-        # change view [b * np, sq, sk]
-        attention_probs = attention_probs.view(output_size[0] * output_size[1], output_size[2], -1)
-        # matmul: [b * np, sq, hn]
-        context_layer = torch.bmm(attention_probs, value_layer)
-        # change view [b, np, sq, hn]
-        context_layer = context_layer.view(*output_size)
-        # [b, np, sq, hn] --> [b, sq, np, hn]
-        context_layer = context_layer.transpose(1, 2).contiguous()
-        # [b, sq, np, hn] --> [b, sq, hp]
-        new_context_layer_shape = context_layer.size()[:-2] + (self.hidden_size_per_partition,)
-        context_layer = context_layer.reshape(*new_context_layer_shape)
-
-        return context_layer
-
-
-class SdpaAttention(CoreAttention):
-    def forward(self, query_layer, key_layer, value_layer, attention_mask):
-        if attention_mask is None and query_layer.shape[2] == key_layer.shape[2]:
-            context_layer = torch.nn.functional.scaled_dot_product_attention(query_layer, key_layer, value_layer,
-                                                                             is_causal=True,
-                                                                             dropout_p=self.config.attention_dropout if self.training else 0.0)
+        pytorch_major_version = int(torch.__version__.split('.')[0])
+        if pytorch_major_version >= 2:
+            if attention_mask is None and query_layer.shape[2] == key_layer.shape[2]:
+                context_layer = torch.nn.functional.scaled_dot_product_attention(query_layer, key_layer, value_layer,
+                                                                                 is_causal=True)
+            else:
+                if attention_mask is not None:
+                    attention_mask = ~attention_mask
+                context_layer = torch.nn.functional.scaled_dot_product_attention(query_layer, key_layer, value_layer,
+                                                                                 attention_mask)
+            context_layer = context_layer.transpose(1, 2).contiguous()
+            new_context_layer_shape = context_layer.size()[:-2] + (self.hidden_size_per_partition,)
+            context_layer = context_layer.reshape(*new_context_layer_shape)
         else:
-            if attention_mask is not None:
-                attention_mask = ~attention_mask
-            context_layer = torch.nn.functional.scaled_dot_product_attention(query_layer, key_layer, value_layer,
-                                                                             attention_mask,
-                                                                             dropout_p=self.config.attention_dropout if self.training else 0.0)
-        context_layer = context_layer.transpose(1, 2).contiguous()
-        new_context_layer_shape = context_layer.size()[:-2] + (self.hidden_size_per_partition,)
-        context_layer = context_layer.reshape(*new_context_layer_shape)
-        return context_layer
-
-
-def _get_unpad_data(attention_mask):
-    seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32)
-    indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()
-    max_seqlen_in_batch = seqlens_in_batch.max().item()
-    cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0))
-    return (
-        indices,
-        cu_seqlens,
-        max_seqlen_in_batch,
-    )
+            # Raw attention scores
 
+            # [b, np, sq, sk]
+            output_size = (query_layer.size(0), query_layer.size(1), query_layer.size(2), key_layer.size(2))
 
-# Copied from transformers.models.llama.modeling_llama.LlamaFlashAttention2
-class FlashAttention2(CoreAttention):
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-        self._flash_attn_uses_top_left_mask = not is_flash_attn_greater_or_equal_2_10()
+            # [b, np, sq, hn] -> [b * np, sq, hn]
+            query_layer = query_layer.view(output_size[0] * output_size[1], output_size[2], -1)
+            # [b, np, sk, hn] -> [b * np, sk, hn]
+            key_layer = key_layer.view(output_size[0] * output_size[1], output_size[3], -1)
 
-    def forward(self, query_states, key_states, value_states, attention_mask):
-        query_states = query_states.transpose(1, 2)
-        key_states = key_states.transpose(1, 2)
-        value_states = value_states.transpose(1, 2)
-        batch_size, query_length = query_states.shape[:2]
-        if not self._flash_attn_uses_top_left_mask:
-            causal = self.is_causal
-        else:
-            # 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__.
-            causal = self.is_causal and query_length != 1
-        dropout = self.config.attention_dropout if self.training else 0.0
-        # Contains at least one padding token in the sequence
-        if attention_mask is not None:
-            query_states, key_states, value_states, indices_q, cu_seq_lens, max_seq_lens = self._upad_input(
-                query_states, key_states, value_states, attention_mask, query_length
-            )
-
-            cu_seqlens_q, cu_seqlens_k = cu_seq_lens
-            max_seqlen_in_batch_q, max_seqlen_in_batch_k = max_seq_lens
-
-            attn_output_unpad = flash_attn_varlen_func(
-                query_states,
-                key_states,
-                value_states,
-                cu_seqlens_q=cu_seqlens_q,
-                cu_seqlens_k=cu_seqlens_k,
-                max_seqlen_q=max_seqlen_in_batch_q,
-                max_seqlen_k=max_seqlen_in_batch_k,
-                dropout_p=dropout,
-                softmax_scale=None,
-                causal=causal,
+            # preallocting input tensor: [b * np, sq, sk]
+            matmul_input_buffer = torch.empty(
+                output_size[0] * output_size[1], output_size[2], output_size[3], dtype=query_layer.dtype,
+                device=query_layer.device
             )
 
-            attn_output = pad_input(attn_output_unpad, indices_q, batch_size, query_length)
-        else:
-            attn_output = flash_attn_func(
-                query_states, key_states, value_states, dropout, softmax_scale=None, causal=causal
+            # Raw attention scores. [b * np, sq, sk]
+            matmul_result = torch.baddbmm(
+                matmul_input_buffer,
+                query_layer,  # [b * np, sq, hn]
+                key_layer.transpose(1, 2),  # [b * np, hn, sk]
+                beta=0.0,
+                alpha=(1.0 / self.norm_factor),
             )
-        attn_output = attn_output.reshape(batch_size, query_length, self.hidden_size_per_partition).contiguous()
-        return attn_output
-
-    def _upad_input(self, query_layer, key_layer, value_layer, attention_mask, query_length):
-        indices_k, cu_seqlens_k, max_seqlen_in_batch_k = _get_unpad_data(attention_mask)
-        batch_size, kv_seq_len, num_key_value_heads, head_dim = key_layer.shape
-
-        key_layer = index_first_axis(
-            key_layer.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim), indices_k
-        )
-        value_layer = index_first_axis(
-            value_layer.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim), indices_k
-        )
-        if query_length == kv_seq_len:
-            query_layer = index_first_axis(
-                query_layer.reshape(batch_size * kv_seq_len, self.num_attention_heads_per_partition, head_dim),
-                indices_k
-            )
-            cu_seqlens_q = cu_seqlens_k
-            max_seqlen_in_batch_q = max_seqlen_in_batch_k
-            indices_q = indices_k
-        elif query_length == 1:
-            max_seqlen_in_batch_q = 1
-            cu_seqlens_q = torch.arange(
-                batch_size + 1, dtype=torch.int32, device=query_layer.device
-            )  # There is a memcpy here, that is very bad.
-            indices_q = cu_seqlens_q[:-1]
-            query_layer = query_layer.squeeze(1)
-        else:
-            # The -q_len: slice assumes left padding.
-            attention_mask = attention_mask[:, -query_length:]
-            query_layer, indices_q, cu_seqlens_q, max_seqlen_in_batch_q = unpad_input(query_layer, attention_mask)
-
-        return (
-            query_layer,
-            key_layer,
-            value_layer,
-            indices_q,
-            (cu_seqlens_q, cu_seqlens_k),
-            (max_seqlen_in_batch_q, max_seqlen_in_batch_k),
-        )
 
+            # change view to [b, np, sq, sk]
+            attention_scores = matmul_result.view(*output_size)
+
+            # ===========================
+            # Attention probs and dropout
+            # ===========================
+
+            # attention scores and attention mask [b, np, sq, sk]
+            if self.attention_softmax_in_fp32:
+                attention_scores = attention_scores.float()
+            if self.coeff is not None:
+                attention_scores = attention_scores * self.coeff
+            if attention_mask is None and attention_scores.shape[2] == attention_scores.shape[3]:
+                attention_mask = torch.ones(output_size[0], 1, output_size[2], output_size[3],
+                                            device=attention_scores.device, dtype=torch.bool)
+                attention_mask.tril_()
+                attention_mask = ~attention_mask
+            if attention_mask is not None:
+                attention_scores = attention_scores.masked_fill(attention_mask, float("-inf"))
+            attention_probs = F.softmax(attention_scores, dim=-1)
+            attention_probs = attention_probs.type_as(value_layer)
+
+            # This is actually dropping out entire tokens to attend to, which might
+            # seem a bit unusual, but is taken from the original Transformer paper.
+            attention_probs = self.attention_dropout(attention_probs)
+
+            # query layer shape: [b * np, sq, hn]
+            # value layer shape: [b, np, sk, hn]
+            # attention shape: [b, np, sq, sk]
+            # context layer shape: [b, np, sq, hn]
+            output_size = (value_layer.size(0), value_layer.size(1), query_layer.size(1), value_layer.size(3))
+            # change view [b * np, sk, hn]
+            value_layer = value_layer.view(output_size[0] * output_size[1], value_layer.size(2), -1)
+            # change view [b * np, sq, sk]
+            attention_probs = attention_probs.view(output_size[0] * output_size[1], output_size[2], -1)
+            # matmul: [b * np, sq, hn]
+            context_layer = torch.bmm(attention_probs, value_layer)
+            # change view [b, np, sq, hn]
+            context_layer = context_layer.view(*output_size)
+            # [b, np, sq, hn] --> [b, sq, np, hn]
+            context_layer = context_layer.transpose(1, 2).contiguous()
+            # [b, sq, np, hn] --> [b, sq, hp]
+            new_context_layer_shape = context_layer.size()[:-2] + (self.hidden_size_per_partition,)
+            context_layer = context_layer.reshape(*new_context_layer_shape)
 
-CORE_ATTENTION_CLASSES = {
-    "eager": CoreAttention,
-    "sdpa": SdpaAttention,
-    "flash_attention_2": FlashAttention2
-}
+        return context_layer
 
 
 class SelfAttention(torch.nn.Module):
@@ -413,7 +299,7 @@ class SelfAttention(torch.nn.Module):
                                          device=device, **_config_to_kwargs(config)
                                          )
 
-        self.core_attention = CORE_ATTENTION_CLASSES[config._attn_implementation](config, self.layer_number)
+        self.core_attention = CoreAttention(config, self.layer_number)
 
         # Output.
         self.dense = nn.Linear(self.projection_size, config.hidden_size, bias=config.add_bias_linear,
@@ -492,8 +378,7 @@ class SelfAttention(torch.nn.Module):
             value_layer = torch.cat((cache_v, value_layer), dim=2)
         if use_cache:
             if kv_cache is None:
-                kv_cache = torch.cat((key_layer.unsqueeze(0).unsqueeze(0), value_layer.unsqueeze(0).unsqueeze(0)),
-                                     dim=1)
+                kv_cache = torch.cat((key_layer.unsqueeze(0).unsqueeze(0), value_layer.unsqueeze(0).unsqueeze(0)), dim=1)
             else:
                 kv_cache = (key_layer, value_layer)
         else:
@@ -759,18 +644,12 @@ class ChatGLMPreTrainedModel(PreTrainedModel):
     config_class = ChatGLMConfig
     base_model_prefix = "transformer"
     _no_split_modules = ["GLMBlock"]
-    _supports_flash_attn_2 = True
-    _supports_sdpa = True
 
     def _init_weights(self, module: nn.Module):
         """Initialize the weights."""
         return
 
     def get_masks(self, input_ids, past_key_values, padding_mask=None):
-        if self.config._attn_implementation == "flash_attention_2":
-            if padding_mask is not None and not padding_mask.all():
-                return padding_mask
-            return None
         batch_size, seq_length = input_ids.shape
         full_attention_mask = torch.ones(batch_size, seq_length, seq_length, device=input_ids.device)
         full_attention_mask.tril_()
@@ -793,6 +672,11 @@ class ChatGLMPreTrainedModel(PreTrainedModel):
         position_ids = torch.arange(seq_length, dtype=torch.long, device=device).unsqueeze(0).repeat(batch_size, 1)
         return position_ids
 
+    def gradient_checkpointing_enable(self, gradient_checkpointing_kwargs=None):
+        if not self.supports_gradient_checkpointing:
+            raise ValueError(f"{self.__class__.__name__} does not support gradient checkpointing.")
+
+
 class Embedding(torch.nn.Module):
     """Language model embeddings."""
 
@@ -840,8 +724,7 @@ class ChatGLMModel(ChatGLMPreTrainedModel):
             config.hidden_size // config.num_attention_heads if config.kv_channels is None else config.kv_channels
         )
 
-        self.rotary_pos_emb = RotaryEmbedding(rotary_dim // 2, rope_ratio=config.rope_ratio,
-                                              original_impl=config.original_rope,
+        self.rotary_pos_emb = RotaryEmbedding(rotary_dim // 2, rope_ratio=config.rope_ratio, original_impl=config.original_rope, 
                                               device=device, dtype=config.torch_dtype)
         self.encoder = init_method(GLMTransformer, config, **init_kwargs)
         self.output_layer = init_method(nn.Linear, config.hidden_size, config.padded_vocab_size, bias=False,
@@ -862,7 +745,6 @@ class ChatGLMModel(ChatGLMPreTrainedModel):
             past_key_values: Optional[Tuple[Tuple[torch.Tensor, torch.Tensor], ...]] = None,
             inputs_embeds: Optional[torch.Tensor] = None,
             use_cache: Optional[bool] = None,
-            output_attentions: Optional[bool] = None,
             output_hidden_states: Optional[bool] = None,
             return_dict: Optional[bool] = None,
     ):
@@ -924,10 +806,12 @@ class ChatGLMForConditionalGeneration(ChatGLMPreTrainedModel):
             outputs: ModelOutput,
             model_kwargs: Dict[str, Any],
             is_encoder_decoder: bool = False,
+            standardize_cache_format: bool = False,
     ) -> Dict[str, Any]:
         # update past_key_values
-        cache_name, cache = self._extract_past_from_model_output(outputs)
-        model_kwargs[cache_name] = cache
+        model_kwargs["past_key_values"] = self._extract_past_from_model_output(
+            outputs, standardize_cache_format=standardize_cache_format
+        )
 
         # update attention mask
         if "attention_mask" in model_kwargs:
@@ -1052,6 +936,201 @@ class ChatGLMForConditionalGeneration(ChatGLMPreTrainedModel):
             for layer_past in past
         )
 
+    def process_response(self, output, history):
+        content = ""
+        history = deepcopy(history)
+        for response in output.split("<|assistant|>"):
+            if "\n" in response:
+                metadata, content = response.split("\n", maxsplit=1)
+            else:
+                metadata, content = "", response
+            if not metadata.strip():
+                content = content.strip()
+                history.append({"role": "assistant", "metadata": metadata, "content": content})
+                content = content.replace("[[训练时间]]", "2023年")
+            else:
+                history.append({"role": "assistant", "metadata": metadata, "content": content})
+                if history[0]["role"] == "system" and "tools" in history[0]:
+                    parameters = json.loads(content)
+                    content = {"name": metadata.strip(), "parameters": parameters}
+                else:
+                    content = {"name": metadata.strip(), "content": content}
+        return content, history
+
+    @torch.inference_mode()
+    def chat(self, tokenizer, query: str, history: List[Dict] = None, role: str = "user",
+             max_length: int = 8192, num_beams=1, do_sample=True, top_p=0.8, temperature=0.8, logits_processor=None,
+             **kwargs):
+        if history is None:
+            history = []
+        if logits_processor is None:
+            logits_processor = LogitsProcessorList()
+        logits_processor.append(InvalidScoreLogitsProcessor())
+        gen_kwargs = {"max_length": max_length, "num_beams": num_beams, "do_sample": do_sample, "top_p": top_p,
+                      "temperature": temperature, "logits_processor": logits_processor, **kwargs}
+        history.append({"role": role, "content": query})
+        inputs = tokenizer.apply_chat_template(history, add_generation_prompt=True, tokenize=True,
+                                               return_tensors="pt", return_dict=True)
+        inputs = inputs.to(self.device)
+        eos_token_id = [tokenizer.eos_token_id, tokenizer.convert_tokens_to_ids("<|user|>"),
+                        tokenizer.convert_tokens_to_ids("<|observation|>")]
+        outputs = self.generate(**inputs, **gen_kwargs, eos_token_id=eos_token_id)
+        outputs = outputs.tolist()[0][len(inputs["input_ids"][0]):-1]
+        response = tokenizer.decode(outputs)
+        response, history = self.process_response(response, history)
+        return response, history
+
+    @torch.inference_mode()
+    def stream_chat(self, tokenizer, query: str, history: List[Dict] = None, role: str = "user",
+                    past_key_values=None, max_length: int = 8192, do_sample=True, top_p=0.8, temperature=0.8,
+                    logits_processor=None, return_past_key_values=False, **kwargs):
+        if history is None:
+            history = []
+        if logits_processor is None:
+            logits_processor = LogitsProcessorList()
+        logits_processor.append(InvalidScoreLogitsProcessor())
+        eos_token_id = [tokenizer.eos_token_id, tokenizer.convert_tokens_to_ids("<|user|>"),
+                        tokenizer.convert_tokens_to_ids("<|observation|>")]
+        gen_kwargs = {"max_length": max_length, "do_sample": do_sample, "top_p": top_p,
+                      "temperature": temperature, "logits_processor": logits_processor, **kwargs}
+        if past_key_values is None:
+            inputs = tokenizer.apply_chat_template(history + [{"role": role, "content": query}],
+                                                   add_generation_prompt=True, tokenize=True, return_tensors="pt",
+                                                   return_dict=True)
+        else:
+            inputs = tokenizer.apply_chat_template([{"role": role, "content": query}], add_special_tokens=False,
+                                                   add_generation_prompt=True, tokenize=True, return_tensors="pt",
+                                                   return_dict=True)
+        inputs = inputs.to(self.device)
+        if past_key_values is not None:
+            past_length = past_key_values[0][0].shape[2]
+            inputs.position_ids += past_length
+            attention_mask = inputs.attention_mask
+            attention_mask = torch.cat((attention_mask.new_ones(1, past_length), attention_mask), dim=1)
+            inputs['attention_mask'] = attention_mask
+        history.append({"role": role, "content": query})
+        for outputs in self.stream_generate(**inputs, past_key_values=past_key_values,
+                                            eos_token_id=eos_token_id, return_past_key_values=return_past_key_values,
+                                            **gen_kwargs):
+            if return_past_key_values:
+                outputs, past_key_values = outputs
+            outputs = outputs.tolist()[0][len(inputs["input_ids"][0]):-1]
+            response = tokenizer.decode(outputs)
+            if response and response[-1] != "�":
+                response, new_history = self.process_response(response, history)
+                if return_past_key_values:
+                    yield response, new_history, past_key_values
+                else:
+                    yield response, new_history
+
+    @torch.inference_mode()
+    def stream_generate(
+            self,
+            input_ids,
+            generation_config: Optional[GenerationConfig] = None,
+            logits_processor: Optional[LogitsProcessorList] = None,
+            stopping_criteria: Optional[StoppingCriteriaList] = None,
+            prefix_allowed_tokens_fn: Optional[Callable[[int, torch.Tensor], List[int]]] = None,
+            return_past_key_values=False,
+            **kwargs,
+    ):
+        batch_size, input_ids_seq_length = input_ids.shape[0], input_ids.shape[-1]
+
+        if generation_config is None:
+            generation_config = self.generation_config
+        generation_config = copy.deepcopy(generation_config)
+        model_kwargs = generation_config.update(**kwargs)
+        model_kwargs["use_cache"] = generation_config.use_cache
+        bos_token_id, eos_token_id = generation_config.bos_token_id, generation_config.eos_token_id
+
+        if isinstance(eos_token_id, int):
+            eos_token_id = [eos_token_id]
+        eos_token_id_tensor = torch.tensor(eos_token_id).to(input_ids.device) if eos_token_id is not None else None
+
+        has_default_max_length = kwargs.get("max_length") is None and generation_config.max_length is not None
+        if has_default_max_length and generation_config.max_new_tokens is None:
+            warnings.warn(
+                f"Using `max_length`'s default ({generation_config.max_length}) to control the generation length. "
+                "This behaviour is deprecated and will be removed from the config in v5 of Transformers -- we"
+                " recommend using `max_new_tokens` to control the maximum length of the generation.",
+                UserWarning,
+            )
+        elif generation_config.max_new_tokens is not None:
+            generation_config.max_length = generation_config.max_new_tokens + input_ids_seq_length
+            if not has_default_max_length:
+                logger.warn(
+                    f"Both `max_new_tokens` (={generation_config.max_new_tokens}) and `max_length`(="
+                    f"{generation_config.max_length}) seem to have been set. `max_new_tokens` will take precedence. "
+                    "Please refer to the documentation for more information. "
+                    "(https://huggingface.co/docs/transformers/main/en/main_classes/text_generation)",
+                    UserWarning,
+                )
+
+        if input_ids_seq_length >= generation_config.max_length:
+            input_ids_string = "decoder_input_ids" if self.config.is_encoder_decoder else "input_ids"
+            logger.warning(
+                f"Input length of {input_ids_string} is {input_ids_seq_length}, but `max_length` is set to"
+                f" {generation_config.max_length}. This can lead to unexpected behavior. You should consider"
+                " increasing `max_new_tokens`."
+            )
+
+        # 2. Set generation parameters if not already defined
+        logits_processor = logits_processor if logits_processor is not None else LogitsProcessorList()
+        stopping_criteria = stopping_criteria if stopping_criteria is not None else StoppingCriteriaList()
+
+        logits_processor = self._get_logits_processor(
+            generation_config=generation_config,
+            input_ids_seq_length=input_ids_seq_length,
+            encoder_input_ids=input_ids,
+            prefix_allowed_tokens_fn=prefix_allowed_tokens_fn,
+            logits_processor=logits_processor,
+        )
+
+        stopping_criteria = self._get_stopping_criteria(
+            generation_config=generation_config, stopping_criteria=stopping_criteria
+        )
+        logits_warper = self._get_logits_warper(generation_config)
+
+        unfinished_sequences = input_ids.new(input_ids.shape[0]).fill_(1)
+        scores = None
+        while True:
+            model_inputs = self.prepare_inputs_for_generation(input_ids, **model_kwargs)
+            # forward pass to get next token
+            outputs = self(
+                **model_inputs,
+                return_dict=True,
+                output_attentions=False,
+                output_hidden_states=False,
+            )
+
+            next_token_logits = outputs.logits[:, -1, :]
+
+            # pre-process distribution
+            next_token_scores = logits_processor(input_ids, next_token_logits)
+            next_token_scores = logits_warper(input_ids, next_token_scores)
+
+            # sample
+            probs = nn.functional.softmax(next_token_scores, dim=-1)
+            if generation_config.do_sample:
+                next_tokens = torch.multinomial(probs, num_samples=1).squeeze(1)
+            else:
+                next_tokens = torch.argmax(probs, dim=-1)
+            # update generated ids, model inputs, and length for next step
+            input_ids = torch.cat([input_ids, next_tokens[:, None]], dim=-1)
+            model_kwargs = self._update_model_kwargs_for_generation(
+                outputs, model_kwargs, is_encoder_decoder=self.config.is_encoder_decoder
+            )
+            unfinished_sequences = unfinished_sequences.mul(
+                next_tokens.tile(eos_token_id_tensor.shape[0], 1).ne(eos_token_id_tensor.unsqueeze(1)).prod(dim=0)
+            )
+            if return_past_key_values:
+                yield input_ids, outputs.past_key_values
+            else:
+                yield input_ids
+            # stop when each sentence is finished, or if we exceed the maximum length
+            if unfinished_sequences.max() == 0 or stopping_criteria(input_ids, scores):
+                break
+
 
 class ChatGLMForSequenceClassification(ChatGLMPreTrainedModel):
     def __init__(self, config: ChatGLMConfig, empty_init=True, device=None):
@@ -1060,7 +1139,7 @@ class ChatGLMForSequenceClassification(ChatGLMPreTrainedModel):
         self.num_labels = config.num_labels
         self.transformer = ChatGLMModel(config, empty_init=empty_init, device=device)
 
-        self.classifier_head = nn.Linear(config.hidden_size, config.num_labels, bias=True, dtype=config.torch_dtype)
+        self.classifier_head = nn.Linear(config.hidden_size, config.num_labels, bias=True, dtype=torch.half)
         if config.classifier_dropout is not None:
             self.dropout = nn.Dropout(config.classifier_dropout)
         else:
@@ -1077,7 +1156,6 @@ class ChatGLMForSequenceClassification(ChatGLMPreTrainedModel):
             inputs_embeds: Optional[torch.LongTensor] = None,
             labels: Optional[torch.LongTensor] = None,
             use_cache: Optional[bool] = None,
-            output_attentions: Optional[bool] = None,
             output_hidden_states: Optional[bool] = None,
             return_dict: Optional[bool] = None,
     ) -> Union[Tuple[torch.Tensor, ...], SequenceClassifierOutputWithPast]:
@@ -1091,7 +1169,6 @@ class ChatGLMForSequenceClassification(ChatGLMPreTrainedModel):
             past_key_values=past_key_values,
             inputs_embeds=inputs_embeds,
             use_cache=use_cache,
-            output_attentions=output_attentions,
             output_hidden_states=output_hidden_states,
             return_dict=return_dict,
         )

tokenization_chatglm.py

@@ -1,10 +1,12 @@
 import regex as re
 import base64
 import os
+import json
 import tiktoken
-from typing import List, Optional, Union, Dict
+from torch import TensorType
+from typing import List, Optional, Union, Dict, Any
 from transformers import PreTrainedTokenizer
-from transformers.utils import PaddingStrategy
+from transformers.utils import logging, PaddingStrategy
 from transformers.tokenization_utils_base import EncodedInput, BatchEncoding
 
 
@@ -15,13 +17,16 @@ class ChatGLM4Tokenizer(PreTrainedTokenizer):
     def __init__(
             self,
             vocab_file,
+            padding_side="left",
             clean_up_tokenization_spaces=False,
+            encode_special_tokens=False,
             **kwargs
     ):
         self.name = "GLM4Tokenizer"
         self.vocab_file = vocab_file
         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+"
         self.pat_str = re.compile(pat_str)
+        self.encode_special_tokens = encode_special_tokens
 
         mergeable_ranks = {}
         with open(vocab_file) as f:
@@ -43,6 +48,7 @@ class ChatGLM4Tokenizer(PreTrainedTokenizer):
         self.n_words = len(self.decoder)
 
         super().__init__(
+            padding_side=padding_side,
             clean_up_tokenization_spaces=clean_up_tokenization_spaces,
             **kwargs
         )
@@ -57,22 +63,22 @@ class ChatGLM4Tokenizer(PreTrainedTokenizer):
         vocab.update(self.added_tokens_encoder)
         return vocab
 
-    def convert_tokens_to_string(self, tokens: List[Union[bytes, str, int]]) -> str:
+    def convert_tokens_to_string(self, tokens: List[Union[bytes, str]]) -> str:
         """
         Converts a sequence of tokens in a single string.
         """
         text = ""
         temp = b""
         for t in tokens:
-            if isinstance(t, int):
-                t = chr(t)
             if isinstance(t, str):
                 if temp:
                     text += temp.decode("utf-8", errors="replace")
+                    temp = b""
+                text += t
             elif isinstance(t, bytes):
                 temp += t
             else:
-                raise TypeError("token should only be of type int, bytes or str")
+                raise TypeError("token should only be of type types or str")
         if temp:
             text += temp.decode("utf-8", errors="replace")
         return text
@@ -135,6 +141,99 @@ class ChatGLM4Tokenizer(PreTrainedTokenizer):
         else:
             return str(f"<|{role}|>{metadata}\n{message}")
 
+    def apply_chat_template(
+            self,
+            conversation: Union[List[Dict[str, str]], List[List[Dict[str, str]]], "Conversation"],
+            add_generation_prompt: bool = False,
+            tokenize: bool = True,
+            padding: bool = False,
+            truncation: bool = False,
+            max_length: Optional[int] = None,
+            return_tensors: Optional[Union[str, TensorType]] = None,
+            return_dict: bool = False,
+            tokenizer_kwargs: Optional[Dict[str, Any]] = None,
+            add_special_tokens: bool = True,
+            **kwargs,
+    ) -> Union[str, List[int], List[str], List[List[int]], BatchEncoding]:
+
+        if return_dict and not tokenize:
+            raise ValueError(
+                "`return_dict=True` is incompatible with `tokenize=False`, because there is no dict "
+                "of tokenizer outputs to return."
+            )
+
+        def handle_single_conversation(conversation):
+            input_ids = self.get_prefix_tokens() if add_special_tokens else []
+            input_message = "[gMASK]<sop>" if add_special_tokens else ""
+            for item in conversation:
+                if item.get("tools"):
+                    tools = item["tools"]
+                    content = "你是一个名为 GLM-4 的人工智能助手。你是基于智谱AI训练的语言模型 GLM-4 模型开发的，你的任务是针对用户的问题和要求提供适当的答复和支持。"
+                    for tool in tools:
+                        if tool["type"] == "function":
+                            function = tool["function"]
+                            content += f"\n\n## {function['name']}\n\n{json.dumps(function, ensure_ascii=False, indent=4)}"
+                            content += "\n在调用上述函数时，请使用 Json 格式表示调用的参数。"
+                        elif tool["type"] == "python":
+                            content += "\n\n## python\n\n当你向 `python` 发送包含 Python 代码的消息时，该代码将会在一个有状态的 Jupyter notebook 环境中执行。\n`python` 返回代码执行的输出，或在执行 60 秒后返回超时。\n`/mnt/data` 将会持久化存储你的文件。在此会话中，`python` 无法访问互联网。不要使用 `python` 进行任何网络请求或者在线 API 调用，这些在线内容的访问将不会成功。"
+                        elif tool["type"] == "simple_browser":
+                            content += "\n\n## simple_browser\n\n你可以使用 `simple_browser` 工具。该工具支持以下函数：\n`search(query: str, recency_days: int)`：使用搜索引擎进行查询并显示结果，可以使用 `recency_days` 参数控制搜索内容的时效性。\n`mclick(ids: list[int])`：获取一系列指定 id 的页面内容。每次调用时，须选择3-10个页面。选择多个角度的页面，同时尽可能选择可信任的信息来源。考虑到部分页面是无法加载的，你也可以多打开一些可能有用的页面而不用担心内容过多。\n`open_url(url: str)`：打开指定的 URL。\n\n使用 `【{引用 id}†{引用文本}】` 来引用内容。\n\n操作步骤：1. 使用 `search` 来获得信息列表; 2. 使用 `mclick` 来获取指定 ID 页面的内容; 3. 根据获得的内容进行回复。在回复中应当引用信息来源。\n 如果用户提供了 URL，也可以用 `open_url` 直接打开页面。\n如果初次搜索结果没有找到合适的信息，也可以再次使用 `search` 进行搜索。"
+                        elif tool["type"] == "cogview":
+                            content += "\n\n## cogview\n\n如果用户的请求中包含了对图像的描述，你可以使用 `cogview` 来生成图像并展示给用户。你需要向 `cogview` 发送图像描述，规则：\n- 发送给 `cogview` 的消息必须使用英语。用户的中文描述必须完全翻译为英语。\n- 应当尽可能详细地描述图像生成的需求，需求描述约 100 英文单词。\n- 保持用户原始描述的意图。不要虚构内容或者没见过的人物。\n- 如无特殊说明，所在地为中国，持有中国立场并遵循中国社会主义价值观。"
+                        else:
+                            raise NotImplementedError(f"Unknown tool type {tool['type']}")
+                    input = self.build_single_message("system", "", content, tokenize=tokenize)
+                    if tokenize:
+                        input_ids.extend(input)
+                    else:
+                        input_message += input
+                if item["content"]:
+                    input = self.build_single_message(
+                        item["role"],
+                        item.get("metadata", ""),
+                        item["content"],
+                        tokenize=tokenize
+                    )
+                    if tokenize:
+                        input_ids.extend(input)
+                    else:
+                        input_message += input
+            if add_generation_prompt:
+                if tokenize:
+                    input_ids.extend([self.convert_tokens_to_ids("<|assistant|>")])
+                else:
+                    input_message += "<|assistant|>"
+
+            return input_ids if tokenize else input_message
+
+        # Main logic to handle different conversation formats
+        if isinstance(conversation, list) and all(isinstance(i, dict) for i in conversation):
+            result = handle_single_conversation(conversation)
+        elif isinstance(conversation, list) and all(isinstance(i, list) for i in conversation):
+            result = [handle_single_conversation(c) for c in conversation]
+        elif hasattr(conversation, "messages"):
+            result = handle_single_conversation(conversation.messages)
+        else:
+            raise ValueError("Invalid conversation format")
+
+        if tokenize:
+            output = self.batch_encode_plus(
+                [result] if isinstance(result[0], int) else result,
+                padding=padding,
+                truncation=truncation,
+                max_length=max_length,
+                return_tensors=return_tensors,
+                is_split_into_words=True,
+                add_special_tokens=False
+            )
+            if return_dict:
+                return output
+            else:
+                return output["input_ids"]
+        else:
+            return result
+
+
     def build_inputs_with_special_tokens(
             self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None
     ) -> List[int]:
@@ -164,7 +263,6 @@ class ChatGLM4Tokenizer(PreTrainedTokenizer):
             self,
             encoded_inputs: Union[Dict[str, EncodedInput], BatchEncoding],
             max_length: Optional[int] = None,
-            padding_side: str = "left",
             padding_strategy: PaddingStrategy = PaddingStrategy.DO_NOT_PAD,
             pad_to_multiple_of: Optional[int] = None,
             return_attention_mask: Optional[bool] = None,
@@ -193,6 +291,7 @@ class ChatGLM4Tokenizer(PreTrainedTokenizer):
                 (optional) Set to False to avoid returning attention mask (default: set to model specifics)
         """
         # Load from model defaults
+        assert self.padding_side == "left"
 
         required_input = encoded_inputs[self.model_input_names[0]]
         seq_length = len(required_input)
@@ -221,4 +320,4 @@ class ChatGLM4Tokenizer(PreTrainedTokenizer):
                 encoded_inputs["position_ids"] = [0] * difference + encoded_inputs["position_ids"]
             encoded_inputs[self.model_input_names[0]] = [self.pad_token_id] * difference + required_input
 
-        return encoded_inputs
+        return encoded_inputs

tokenizer_config.json

@@ -123,7 +123,6 @@
                                "<|user|>", "<|assistant|>", "<|observation|>", "<|begin_of_image|>", "<|end_of_image|>",
                                "<|begin_of_video|>", "<|end_of_video|>"],
   "clean_up_tokenization_spaces": false,
-  "chat_template": "[gMASK]<sop>{% for item in messages %}{% if item['tools'] is defined %}<|system|>\n你是一个名为 ChatGLM 的人工智能助手。你是基于智谱AI训练的语言模型 GLM-4 模型开发的，你的任务是针对用户的问题和要求提供适当的答复和支持。\n\n# 可用工具{% set tools = item['tools'] %}{% for tool in tools %}{% if tool['type'] == 'function' %}\n\n## {{ tool['function']['name'] }}\n\n{{ tool['function'] | tojson(indent=4) }}\n在调用上述函数时，请使用 Json 格式表示调用的参数。{% elif tool['type'] == 'python' %}\n\n## python\n\n当你向 `python` 发送包含 Python 代码的消息时，该代码将会在一个有状态的 Jupyter notebook 环境中执行。\n`python` 返回代码执行的输出，或在执行 60 秒后返回超时。\n`/mnt/data` 将会持久化存储你的文件。在此会话中，`python` 无法访问互联网。不要使用 `python` 进行任何网络请求或者在线 API 调用，这些在线内容的访问将不会成功。{% elif tool['type'] == 'simple_browser' %}\n\n## simple_browser\n\n你可以使用 `simple_browser` 工具。该工具支持以下函数：\n`search(query: str, recency_days: int)`：使用搜索引擎进行查询并显示结果，可以使用 `recency_days` 参数控制搜索内容的时效性。\n`mclick(ids: list[int])`：获取一系列指定 id 的页面内容。每次调用时，须选择3-10个页面。选择多个角度的页面，同时尽可能选择可信任的信息来源。考虑到部分页面是无法加载的，你也可以多打开一些可能有用的页面而不用担心内容过多。\n`open_url(url: str)`：打开指定的 URL。\n\n使用 `【{引用 id}†{引用文本}】` 来引用内容。\n\n操作步骤：1. 使用 `search` 来获得信息列表; 2. 使用 `mclick` 来获取指定 ID 页面的内容; 3. 根据获得的内容进行回复。在回复中应当引用信息来源。\n 如果用户提供了 URL，也可以用 `open_url` 直接打开页面。\n如果初次搜索结果没有找到合适的信息，也可以再次使用 `search` 进行搜索。{% elif tool['type'] == 'cogview' %}\n\n## cogview\n\n如果用户的请求中包含了对图像的描述，你可以使用 `cogview` 来生成图像并展示给用户。你需要向 `cogview` 发送图像描述，规则：\n- 发送给 `cogview` 的消息必须使用英语。用户的中文描述必须完全翻译为英语。\n- 应当尽可能详细地描述图像生成的需求，需求描述约 100 英文单词。\n- 保持用户原始描述的意图。不要虚构内容或者没见过的人物。\n- 如无特殊说明，所在地为中国，持有中国立场并遵循中国社会主义价值观。{% endif %}{% endfor %}{% endif %}{% if item['content'] %}<|{{ item['role'] }}|>{{ item['metadata'] }}\n{{ item['content'] }}{% endif %}{% endfor %}{% if add_generation_prompt %}<|assistant|>{% endif %}",
   "do_lower_case": false,
   "eos_token": "<|endoftext|>",
   "pad_token": "<|endoftext|>",