init model files

README.md

@@ -1,3 +1,186 @@
----
-license: apache-2.0
----
+<div align="center">
+<h1>
+  星辰语义大模型-TeleChat2
+</h1>
+</div>
+
+
+<p align="center">
+	🦉 <a href="https://github.com/Tele-AI/TeleChat2" target="_blank">github</a>️ • 🤗 <a href="https://huggingface.co/Tele-AI" target="_blank">Hugging Face</a> • 🤖 <a href="https://modelscope.cn/organization/TeleAI" target="_blank">ModelScope</a> • 🏔 <a href="https://gitee.com/mindspore/mindformers/tree/dev/research/telechat2" target="_blank">MindSpore</a> • 🐾 <a href="https://gitee.com/Tele-AI/tele-chat2" target="_blank">gitee</a>️ • 💬 <a href="https://github.com/Tele-AI/Telechat/blob/master/images/wechat.jpg" target="_blank">WeChat</a>
+</p>
+
+# 目录
+
+- [模型介绍](#模型介绍)
+- [效果评测](#效果评测)
+- [模型推理](#模型推理)
+- [声明、协议、引用](#声明协议引用)
+
+# 最新动态
+- 2024.11.08 开源 **TeleChat2-3B**、**TeleChat2-7B**、**TeleChat2-35B**，该版本模型均具备 **Function Call** 功能。
+- 2024.10.18 开源TeleChat2-35B模型。
+- 2024.9.20 开源TeleChat2-115B模型，该模型是**首个完全国产算力训练并开源的千亿参数模型**。
+
+# 模型介绍
+
+### 星辰语义大模型-TeleChat2
+
+- 星辰语义大模型**TeleChat2**是由中国电信人工智能研究院研发训练的大语言模型，该系列模型**完全基于国产算力**训练。
+- 本次开源的 **TeleChat2-3B**、**TeleChat2-7B**、**TeleChat2-35B** 模型已支持**工具调用**功能。在 **Function Call** 方面，我们针对性进行了效果优化，在相关榜单评测上相比同尺寸模型均有较好表现。
+- **TeleChat2-115B**模型采用10万亿 Tokens中英文高质量语料进行训练，同步开源对话模型**TeleChat2-115B**的多格式、多平台权重文件。
+- **TeleChat2**在训练数据、训练方法等方面进行了改进，在通用问答和知识类、代码类、数学类榜单上相比**TeleChat1**均有大幅提升。
+    - **TeleChat2**完全基于国产算力和国产深度学习框架进行训练，算力和算法框架更自主可控。优化MP、PP、SP实现方式提升模型性能，优化算子来提升训练速度。
+    - 我们使用大量小模型实验来验证scaling law规律，在不同模型结构、不同数据配比和数据清洗方式中寻找最优设计。
+    - 采用RingAttention及其他序列切分方式，实现长文训练性能提升；通过ntk-aware+attention-scaling的方式保证训练长度切换时的平稳过渡，以此来保证模型在不同长度数据下的训练效果。
+- 在微调数据方面，我们进行了指令复杂性提升与多样性扩充，通过数据合成和人工标注生成高质量数据，并使用拒绝采样生成多样的推理路径；通过研究一套基于base模型反向选择偏好对齐数据方案，基于适配数据最大限度提升模型效果。
+    - 通用能力较TeleChat系列模型提升超过29%，在逻辑推理、总结摘要、长文写作和数学计算上均有大幅提升。
+
+### 模型结构
+
+我们采用标准的 `Decoder-only` 结构设计了 **TeleChat2** 模型，使用 [Rotary Embedding](https://arxiv.org/pdf/2104.09864.pdf)
+的位置编码方法、使用 [SwiGLU](https://arxiv.org/pdf/2002.05202.pdf)
+激活函数来替代GELU激活函数、使用基于 [RMSNorm](https://arxiv.org/abs/1910.07467) 的 Pre-Normalization进行层标准化操作。我们将**TeleChat2**的词嵌入层和输出lm
+head层参数分开，有助于增强训练稳定性和收敛性。我们选择了GQA以节约attention部分的参数量和计算量、提升训练和推理速度。
+
+**TeleChat2**的模型结构配置如下表所示：
+
+|      | layer_num | hidden_size | ffn_hidden_size | head_num | tie_word_embeddings | GQA  |
+| ---- | --------- | ----------- | --------------- | -------- | ------------------- | ---- |
+| 3B   | 24          | 3072      |     6144        | 24       | 否                  | 否   |
+| 7B   | 30          | 4096      | 12288           | 32       | 否                  | 否   |
+| 35B  | 64         | 6144       | 20480           | 48       | 否                  | 否   |
+| 115B | 96        | 8192        | 40960           | 64       | 否                  | 是   |
+
+
+我们开源的 **TeleChat2** 模型：
+
+- 支持deepspeed微调，开源了基于deepspeed的训练代码，支持Zero并行显存优化，同时集成了FlashAttention2
+- 多轮能力支持。开源了多轮数据构建方式，针对多轮模型训练集成了针对多轮的mask loss训练方式，更好的聚焦多轮答案，提升问答效果。
+
+本次发布版本和下载链接见下表
+
+| 模型版本       | 下载链接 |
+| -------------- | -------- |
+| telechat2-3B |   [modelscope](https://modelscope.cn/models/TeleAI/TeleChat2-3B)|
+| telechat2-7B |   [modelscope](https://modelscope.cn/models/TeleAI/TeleChat2-7B)|
+| telechat2-35B |   [modelscope](https://modelscope.cn/models/TeleAI/TeleChat2-35B-Nov)|
+| telechat2-115B |   [modelscope](https://modelscope.cn/models/TeleAI/TeleChat2-115B)|
+
+
+# 效果评测
+
+**TeleChat2** 模型相比同规模模型在评测效果方面也有较好的表现，我们的评测集涵盖了包括MMLU、C-Eval、CMMLU、
+GSM8K、MATH、HumanEval、BBH等数据集，评测能力包括了指令遵循、考试能力、数学计算和推理、代码生成等
+
+## 评测集介绍
+
+### 通用能力
+
+- MMLU 数据集是一个全面的英文评测数据集，涵盖了 57 个学科，包括人文学科、社会科学、自然科学、初等数学、美国历史、计算机科学、法律等等。
+
+- CEVAL 数据集是一个全面的中文评估测试集，包括初中、高中、大学和专业难度级别的多项选择题，涵盖了 52 个不同的学科领域。
+
+- CMMLU 数据集同样是一个全面的中文评估测试集，涵盖了从基础学科到高级专业水平的67个主题。
+
+### 推理和代码能力
+
+- GSM8K 数据集包含了8.5K高质量的小学数学题，能够评估语言模型在数学推理能力上的表现。
+
+- HumanEval 数据集是一个由openai提供的代码能力测试数据集，它由 164 个编程问题组成，要求根据给定的问题和代码模板，生成正确的代码片段。
+
+	- BBH 数据集全名为BIG-Bench Hard（BBH），包含23个具有挑战性的BIG-Bench任务，均为之前的语言模型评估中没有超过平均人类评审者表现的任务。
+
+- MBPP 数据集包含大约1000个众包的Python编程问题，涵盖编程基础知识、标准库功能等。每个问题包括任务描述、代码解决方案和3个自动化测试用例。
+
+### 主观题能力
+
+- [AlignBench](https://github.com/THUDM/AlignBench)是一个多维度全面评估中文大模型对齐水平的评测基准，包含638道单轮主观评测题。
+
+- [MT-bench](https://github.com/lm-sys/FastChat/blob/main/fastchat/llm_judge/README.md)是一个用于评估聊天助手的具有挑战性的多轮开放式问题集，包含80通多轮主观评测题。
+
+### 指令遵循能力
+
+- [IFEval](https://github.com/EleutherAI/lm-evaluation-harness/blob/main/lm_eval/tasks/ifeval/README.md)旨在评估语言模型对指令的精确遵循能力，它包含了500条可精确验证的指令，是Open
+  LLM Leaderboard中使用的核心基准测试之一。
+
+## 评测结果如下
+
+ | Dataset    | Llama-3.1-70B | Qwen1.5-110B | Qwen2-72-instruct | DeepSeek-v2 | TeleChat2-115B |TeleChat2-35B |TeleChat2-7B    |TeleChat2-3B    |
+|:----------:|:-------------:|:------------:|:-----------------:|:-----------:|:--------------:|:--------------:|:--------------:|:----------------:|
+| C-Eval     | -             | -            | 83.8              | 78          | **86.9**       |  85            |  82            |    75            | 
+| MMLU       | **86**        | 80.4         | 82.3              | 77.8        | 80.9           |  82            |  79.6          |    72.9          |  
+| CMMLU      | 69.01         | 87.64        | 87.47             | 81.6        | **89.94**      |  90.18         |  84.6          |    73            | 
+| BBH        | -             | 74.8         | -                 | 79.7        | **89.04**      |   88.6         |  77.3          |    65.99         |
+| GSM8K      | **95.1**      | 85.4         | 91.1              | 92.2        | 92.2           |  91            |  86.8          |    64.7          | 
+| HumanEval  | 80.5          | 52.4         |**86**             | 81.1        | 75             |  73            |  56            |    38            |  
+| MBPP       | **86**        | 58.1         | 80.2              | 72          | 78             |   75           |  62.6          |    47            | 
+| AlignBench | -             | 7.86         | **8.27**          | 7.91        | 8.03           |   7.88         |  6.96          |    5.74          | 
+| MT-bench   | 8.79          | 8.88         | **9.12**          | 8.97        | 8.89           |   8.2          |  7.2           |    5.72          | 
+| IFEval     | **87.5**      | -            | 77.6              | 63.8        | 82.81          |   79.63        |  73.1          |    61.29         | 
+
+# 模型推理
+
+### 模型推理
+
+当前模型推理兼容了单卡和多卡推理，以及针对长文推理做了部分优化工作。
+
+**模型推理方法示范**
+
+```python
+>>> import os
+>>> import torch
+>>> from modelscope import AutoModelForCausalLM, AutoTokenizer, GenerationConfig
+>>> tokenizer = AutoTokenizer.from_pretrained('TeleAI/TeleChat2-7B', trust_remote_code=True)
+>>> model = AutoModelForCausalLM.from_pretrained('TeleAI/TeleChat2-7B', trust_remote_code=True, device_map="auto",
+                                                  torch_dtype=torch.float16)
+>>> prompt = "生抽与老抽的区别？"
+>>> messages = [{"role": "user", "content": prompt}]
+>>> text = tokenizer.apply_chat_template(messages,
+>>>		tokenize=False,
+>>>    		add_generation_prompt=True
+>>>	)
+>>> model_inputs = tokenizer([text], return_tensors="pt").to(model.device)
+>>> generated_ids = model.generate(
+>>>     **model_inputs,
+>>>     max_new_tokens=512
+>>> )
+>>> generated_ids = [
+>>>     output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
+>>> ]
+
+>>> response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
+生抽和老抽是两种不同的酱油，它们在风味、色泽和用途上都有所区别。
+
+1.颜色：生抽的颜色比较淡，而老抽的颜色较深。生抽的颜色呈红褐色或棕红色，而老抽的颜色则呈棕黑色。
+
+2.味道：生抽具有鲜美的咸味和微甜的味浅，而老抽浓郁，颜色较深。根据个人口味和烹饪需求选择不同的酱油类型可以获得更好的口感和菜肴效果。
+```
+
+# 声明、协议、引用
+
+### 声明
+
+我们在此声明，不要使用TeleChat模型及其衍生模型进行任何危害国家社会安全或违法的活动。同时，我们也要求使用者不要将TeleChat模型用于没有安全审查和备案的互联网服务。我们希望所有使用者遵守上述原则，确保科技发展在合法合规的环境下进行。
+
+我们已经尽我们所能，来确保模型训练过程中使用的数据的合规性。然而，尽管我们已经做出了巨大的努力，但由于模型和数据的复杂性，仍有可能存在一些无法预见的问题。因此，如果由于使用TeleChat开源模型而导致的任何问题，包括但不限于数据安全问题、公共舆论风险，或模型被误导、滥用、传播或不当利用所带来的任何风险和问题，我们将不承担任何责任。
+
+### 协议
+
+社区使用 TeleChat 模型需要遵循《[TeleChat模型社区许可协议](./TeleChat模型社区许可协议.pdf)》。TeleChat模型支持商业用途，如果您计划将 TeleChat
+模型或其衍生品用于商业目的，您需要通过以下联系邮箱
+[email protected]，提交《TeleChat模型社区许可协议》要求的申请材料。审核通过后，将特此授予您一个非排他性、全球性、不可转让、不可再许可、可撤销的商用版权许可。
+
+### 引用
+
+如需引用我们的工作，请使用如下 reference:
+
+```
+@misc{wang2024telechat,
+      title={TeleChat Technical Report}, 
+      author={Zihan Wang and Xinzhang Liu and Shixuan Liu and Yitong Yao and Yuyao Huang and Zhongjiang He and Xuelong Li and Yongxiang Li and Zhonghao Che and Zhaoxi Zhang and Yan Wang and Xin Wang and Luwen Pu and Huihan Xu and Ruiyu Fang and Yu Zhao and Jie Zhang and Xiaomeng Huang and Zhilong Lu and Jiaxin Peng and Wenjun Zheng and Shiquan Wang and Bingkai Yang and Xuewei he and Zhuoru Jiang and Qiyi Xie and Yanhan Zhang and Zhongqiu Li and Lingling Shi and Weiwei Fu and Yin Zhang and Zilu Huang and Sishi Xiong and Yuxiang Zhang and Chao Wang and Shuangyong Song},
+      year={2024},
+      eprint={2401.03804},
+      archivePrefix={arXiv},
+      primaryClass={cs.CL}
+}
+```

config.json

@@ -0,0 +1,49 @@
+{
+  "apply_residual_connection_post_layernorm": false,
+  "architectures": [
+    "TeleChat2ForCausalLM"
+  ],
+  "auto_map": {
+	"AutoConfig": "configuration_telechat2.Telechat2Config",
+	"AutoModelForCausalLM": "modeling_telechat2.Telechat2ForCausalLM"
+  },
+  "attention_dropout": 0.0,
+  "attention_softmax_in_fp32": true,
+  "bias_dropout_fusion": true,
+  "bos_token_id": 1,
+  "eos_token_id": 2,
+  "hidden_dropout": 0.0,
+  "hidden_size": 4096,
+  "initializer_range": 0.02,
+  "layer_norm_epsilon": 1e-05,
+  "masked_softmax_fusion": true,
+  "max_position_embeddings": 32768,
+  "model_type": "telechat",
+  "n_head": 32,
+  "n_inner": null,
+  "n_layer": 30,
+  "num_key_value_heads":32,
+  "offset_alibi": 100,
+  "pad_token_id": 3,
+  "pretraining_tp": 2,
+  "skip_bias_add": true,
+  "skip_bias_add_qkv": false,
+  "slow_but_exact": false,
+  "transformers_version": "4.44.2",
+  "torch_dtype": "bfloat16",
+  "unk_token_id": 0,
+  "use_cache": true,
+  "vocab_size": 131072,
+  "ffn_hidden_size": 12288,
+  "flash_attn":true,
+  "tie_word_embeddings":false,
+  "rope_scaling": {
+	  "factor": 1.0,
+	  "rope_type": "dynamic"
+  },
+  "rope_theta": 1000000,
+  "training_seqlen":32768,
+  "base_seqlen":32768,
+  "seq_length": 32768
+}
+

configuration.json

@@ -0,0 +1 @@
+{"framework":"Pytorch","task":"text-generation"}

configuration_telechat2.py

@@ -0,0 +1,94 @@
+# coding=utf-8
+# Copyright 2022 the Big Science Workshop and HuggingFace Inc. team.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+""" Telechat configuration"""
+
+from packaging import version
+from collections import OrderedDict
+from transformers.utils import is_torch_available, logging
+from transformers.configuration_utils import PretrainedConfig
+from typing import TYPE_CHECKING, Any, List, Mapping, Optional
+
+logger = logging.get_logger(__name__)
+
+class Telechat2Config(PretrainedConfig):
+    """
+    Args:
+        vocab_size (`int`, *optional*, defaults to 160256): Vocabulary size of the Telechat model.
+        hidden_size (`int`, *optional*, defaults to 4096): Dimensionality of the embeddings and hidden states.
+        ffn_hidden_size (`int`, *optional*, defaults to 12288): Dimensionality of the feed-forward hidden states.
+        n_layer (`int`, *optional*, defaults to 30): Number of hidden layers in the Transformer
+        n_head (`int`, *optional*, defaults to 32): Number of attention heads for each attention layer.
+        layer_norm_epsilon (`float`, *optional*, defaults to 1e-5): The epsilon to use in the layer normalization layers.
+        initializer_range (`float`, *optional*, defaults to 0.02): The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
+        apply_residual_connection_post_layernorm (`bool`, *optional*, defaults to `False`): If enabled, use the layer norm of the hidden states as the residual in the transformer blocks
+        hidden_dropout (`float`, *optional*, defaults to 0.0): Dropout rate of the dropout function on the bias dropout.
+        attention_dropout (`float`, *optional*, defaults to 0.0): Dropout rate applied to the attention probs
+        use_cache (`bool`, *optional*, defaults to `True`): Whether or not the model should return the last key/values attentions.
+        training_seqlen (`int`, *optional*, defaults to 8192): Sequence length during last finetuning.
+        logn (`bool`, *optional*, defaults to `True`): Whether or not to use logN during extrapolation.
+        embed_layernorm (`bool`, *optional*, defaults to `True`): Whether or not to use embedding layernorm.
+
+    """
+
+    model_type = "telechat"
+    keys_to_ignore_at_inference = ["past_key_values"]
+    attribute_map = {
+        "num_hidden_layers": "n_layer",
+        "num_attention_heads": "n_head",
+    }
+
+    def __init__(
+        self,
+        vocab_size=160256,
+        hidden_size=4096,
+        n_layer=30,
+        n_head=32,
+        layer_norm_epsilon=1e-5,
+        initializer_range=0.02,
+        use_cache=True,
+        bos_token_id=1,
+        eos_token_id=2,
+        apply_residual_connection_post_layernorm=False,
+        hidden_dropout=0.0,
+        attention_dropout=0.0,
+        ffn_hidden_size=12288,
+        training_seqlen = 8192,
+        logn = True,
+        embed_layernorm = False,
+        **kwargs,
+    ):
+        self.vocab_size = vocab_size
+        n_embed = kwargs.pop("n_embed", None)
+        self.hidden_size = hidden_size if n_embed is None else n_embed
+        self.n_layer = n_layer
+        self.n_head = n_head
+        self.layer_norm_epsilon = layer_norm_epsilon
+        self.initializer_range = initializer_range
+        self.use_cache = use_cache
+        self.apply_residual_connection_post_layernorm = apply_residual_connection_post_layernorm
+        self.hidden_dropout = hidden_dropout
+        self.attention_dropout = attention_dropout
+        self.bos_token_id = bos_token_id
+        self.eos_token_id = eos_token_id
+        self.logn = logn
+        self.ffn_hidden_size = ffn_hidden_size
+        self.training_seqlen = training_seqlen
+        self.embed_layernorm = embed_layernorm
+        self.num_key_value_heads= kwargs.pop("num_key_value_heads", None)
+
+
+        super().__init__(bos_token_id=bos_token_id, eos_token_id=eos_token_id, **kwargs)
+

generation_config.json

@@ -0,0 +1,15 @@
+{
+  "max_new_tokens": 1000,
+  "do_sample": false,
+  "use_cache": true,
+  "temperature": 0.3,
+  "top_k": 5,
+  "top_p": 0.85,
+  "repetition_penalty": 1.02,
+  "pad_token_id": 3,
+  "bos_token_id": 1,
+  "eos_token_id": 2,
+  "user_token_id": 4,
+  "bot_token_id": 5,
+  "start_token_id": 1
+}

generation_utils.py

@@ -0,0 +1,162 @@
+from typing import Optional
+from collections import deque
+from queue import Queue
+import copy
+
+
+class History:
+
+    def __init__(self, tokenizer, history):
+        '''
+        init from a list of dict
+        '''
+        # use deque to meet some special situation
+        self.input_history = deque()
+        self.tokenizer = tokenizer
+        if history:
+            self._transfer_from_list(history)
+
+    def _transfer_from_list(self, history):
+        for message in history:
+            content = message.get("content")
+            # the token result may not be equal to the result model gen
+            message.update(self.tokenizer(content))
+            self.input_history.append(message)
+
+    def append(self, message):
+        content = message.get("content")
+        if "input_ids" not in message or "attention_mask" not in message:
+            message.update(self.tokenizer(content))
+        self.input_history.append(message)
+
+    def append_left(self, message):
+        content = message.get("content")
+        if "input_ids" not in message or "attention_mask" not in message:
+            message.update(self.tokenizer(content))
+        self.input_history.appendleft(message)
+
+    def pop(self):
+        x = self.input_history.pop()
+        return x
+
+    def pop_left(self):
+        x = self.pop_left()
+        return x
+
+    def update(self, message):
+        self.input_history.pop()
+        self.append(message)
+
+    def __len__(self):
+        return self.input_history.__len__()
+
+    def __str__(self):
+        return self.input_history.__str__()
+
+    def __copy__(self):
+        new_instance = type(self)(self.tokenizer, [])
+        new_instance.input_history = copy.copy(self.input_history)
+        return new_instance
+
+    def __deepcopy__(self, memodict={}):
+        new_instance = type(self)(self.tokenizer, [])
+        new_instance.input_history = copy.deepcopy(self.input_history)
+        return new_instance
+
+
+class TelechatIterTextStreamer:
+    """
+    With reference to the TextIterStreamers in transformers, we have rewritten this class
+    """
+
+    def __init__(
+            self, tokenizer, history: History = None, skip_prompt: bool = False, timeout: Optional[float] = None,
+            **decode_kwargs
+    ):
+
+        self.tokenizer = tokenizer
+        self.history = history
+        self.skip_prompt = skip_prompt
+        self.timeout = timeout
+        self.decode_kwargs = decode_kwargs
+
+        self.text_queue = Queue()
+        self.cache_time = 0
+        self.text_until = ""
+        self.token_until = []
+        self.stop_signal = None
+        self.next_tokens_are_prompt = True
+
+        self.history.append({"role": "bot", "content": self.text_until})
+
+    def put(self, value):
+        """
+        put printable text into queue
+        """
+        if len(value.shape) > 1 and value.shape[0] > 1:
+            raise ValueError("TextStreamer only supports batch size 1")
+        elif len(value.shape) > 1:
+            value = value[0]
+
+        if self.skip_prompt and self.next_tokens_are_prompt:
+            self.next_tokens_are_prompt = False
+            return
+
+        if value[-1] == self.tokenizer.eos_token_id:
+            return
+
+        # there may be some smart way to decode.
+        self.token_until.extend(value.tolist())
+        text = self.tokenizer.decode(self.token_until, **self.decode_kwargs)
+
+
+        if self._is_printable(text) or self.cache_time >= 6:
+            output_text = text[len(self.text_until):]
+            self.text_until = text
+
+        else:
+            self.cache_time+=1
+            return
+
+        self.on_finalized_text(output_text)
+
+    def end(self):
+        """Flushes any remaining cache and prints a newline to stdout."""
+        # Flush the cache, if it exists
+        text = self.tokenizer.decode(self.token_until, **self.decode_kwargs)
+        output_text = text[len(self.text_until):]
+        self.text_until = text
+        self.on_finalized_text(output_text, stream_end=True)
+        self.clear_cache()
+
+    def clear_cache(self):
+        self.cache_time = 0
+        self.token_until = []
+        self.text_until = ""
+        self.history = None
+        self.next_tokens_are_prompt = True
+
+    def on_finalized_text(self, text: str, stream_end: bool = False):
+        """Put the text tuple in the queue."""
+        self.history.update({"role": "bot", "content": self.text_until, "input_ids": self.token_until,
+                             "attention_mask": [1] * len(self.token_until)})
+        self.text_queue.put((text, self.history), timeout=self.timeout)
+        if stream_end:
+            self.text_queue.put((self.stop_signal, self.history), timeout=self.timeout)
+
+    @staticmethod
+    def _is_printable(cp):
+        """Checks whether tokens can be decoded or not"""
+        if "�" in cp:
+            return False
+        return True
+
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        value_now, history_until = self.text_queue.get(timeout=self.timeout)
+        if value_now == self.stop_signal:
+            raise StopIteration()
+        else:
+            return value_now, history_until

modeling_telechat2.py

@@ -0,0 +1,854 @@
+# coding=utf-8
+# Copyright 2022 HuggingFace Inc. team and BigScience workshop.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Copyright (c) 2023, NVIDIA CORPORATION.  All rights reserved.
+
+# Copyright (c) 2021 EleutherAI
+# This file is based on code by the authors denoted below and has been modified from its original version.
+#
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+"""PyTorch TELECHAT model."""
+
+import warnings
+from typing import Optional, Tuple, Union, List, Dict
+from threading import Thread
+
+import torch
+import math
+import copy
+from torch import nn
+import torch.utils.checkpoint
+from torch.nn import functional as F
+from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, LayerNorm, MSELoss
+from transformers.modeling_outputs import (
+    BaseModelOutputWithPastAndCrossAttentions,
+    CausalLMOutputWithCrossAttentions
+)
+from transformers.modeling_utils import PreTrainedModel
+from transformers.utils import logging
+from transformers import GenerationConfig
+
+from .configuration_telechat2 import Telechat2Config
+
+
+logger = logging.get_logger(__name__)
+
+_CHECKPOINT_FOR_DOC = "telechat"
+_CONFIG_FOR_DOC = "Telechat2Config"
+
+TELECHAT_PRETRAINED_MODEL_ARCHIVE_LIST = []
+
+try:
+    from einops import rearrange
+except ImportError:
+    rearrange = None
+
+use_flash_attn = True
+try:
+    from flash_attn.flash_attn_interface import flash_attn_unpadded_func
+except ImportError:
+    try:
+        from flash_attn.flash_attn_interface import flash_attn_varlen_func as flash_attn_unpadded_func
+    except ImportError:
+        flash_attn_unpadded_func = None
+
+
+class RotaryEmbedding(torch.nn.Module):
+    # Extracted from: https://github.com/EleutherAI/gpt-neox
+    def __init__(self, dim, config):
+        super().__init__()
+        self.config = config
+        self.dim = dim
+        self.base = config.rope_theta
+        self.inv_freq = 1. / (self.base ** (torch.arange(0, dim, 2).float().half() / dim))
+        self.max_seq_len_cached = None
+        self.cos_cached = None
+        self.sin_cached = None
+        self.precision = config.torch_dtype
+
+    def get_mscale(self, scale=1):
+        if scale <= 1:
+            return 1.0
+        return 0.1 * math.log(scale) + 1.0
+
+    def get_ntk_alpha(self, true_seq_len):
+        context_value = math.log(true_seq_len / self.config.base_seqlen, 2) + 1
+        # ntk_alpha = 2 ** context_value - 1
+        ntk_alpha = 2 ** math.ceil(context_value) - 1
+        ntk_alpha = max(ntk_alpha, 1)
+        return ntk_alpha
+
+    def forward(self, x, seq_dim=0, seq_len=None):
+        if seq_len is None:
+            seq_len = x.shape[seq_dim]
+        seq_len = max(seq_len, self.config.training_seqlen)
+        ntk_alpha = self.get_ntk_alpha(seq_len)
+        self.mscale = float(self.get_mscale(seq_len / self.config.training_seqlen))
+        if True:
+            base = self.base * ntk_alpha ** (self.dim / (self.dim - 2))
+            self.inv_freq = 1.0 / (base ** (torch.arange(0, self.dim, 2, device=x.device).float() / self.dim))
+            self.max_seq_len_cached = seq_len
+            t = torch.arange(self.max_seq_len_cached, device=x.device, dtype=self.inv_freq.dtype)
+            freqs = torch.einsum('i,j->ij', t, self.inv_freq)
+            # Different from paper, but it uses a different permutation in order to obtain the same calculation
+            emb = torch.cat((freqs, freqs), dim=-1).to(x.device)
+            if self.precision == torch.bfloat16:
+                emb = emb.float()
+            # [sx, 1 (b * np), hn]
+            self.cos_cached = self.mscale * emb.cos()[:, None, :].half()
+            self.sin_cached = self.mscale * emb.sin()[:, None, :].half()
+            if self.precision == torch.bfloat16:
+                self.cos_cached = self.cos_cached.bfloat16()
+                self.sin_cached = self.sin_cached.bfloat16()
+        return self.cos_cached[:seq_len, ...], self.sin_cached[:seq_len, ...]
+
+
+# rotary pos emb helpers:
+def rotate_half(x):
+    x1, x2 = x[..., :x.shape[-1] // 2], x[..., x.shape[-1] // 2:]
+    return torch.cat((-x2, x1), dim=x1.ndim - 1)  # dim=-1 triggers a bug in earlier torch versions
+
+
+def apply_rotary_pos_emb_torch(q, k, cos, sin, offset: int = 0):  # jitting fails with bf16
+    cos, sin = cos[offset:q.shape[0] + offset, ...], sin[offset:q.shape[0] + offset, ...]
+    return (q * cos) + (rotate_half(q) * sin), (k * cos) + (rotate_half(k) * sin)
+
+
+class MixedFusedRMSNorm(nn.Module):
+    # Extracted from https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/modeling_llama.py
+    def __init__(self, hidden_size, eps=1e-6):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.variance_epsilon = eps
+
+    def forward(self, hidden_states):
+        input_dtype = hidden_states.dtype
+        hidden_states = hidden_states.to(torch.float32)
+        variance = hidden_states.pow(2).mean(-1, keepdim=True)
+        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
+        return self.weight * hidden_states.to(input_dtype)
+
+
+class FlashSelfAttention(torch.nn.Module):
+    # Extracted from https://github.com/microsoft/Megatron-DeepSpeed/blob/main/megatron/model/transformer.py
+    """Implement the scaled dot product attention with softmax.
+    Arguments
+    ---------
+        softmax_scale: The temperature to use for the softmax attention.
+                      (default: 1/sqrt(d_keys) where d_keys is computed at
+                      runtime)
+        attention_dropout: The dropout rate to apply to the attention
+                           (default: 0.0)
+    """
+
+    def __init__(self, causal=False, softmax_scale=None, attention_dropout=0.0,
+                 device=None, dtype=None):
+        super().__init__()
+        assert flash_attn_unpadded_func is not None, ('Please install FlashAttention first, '
+                                                      'e.g., with pip install flash-attn')
+        assert rearrange is not None, 'Please install einops first, e.g., with pip install einops'
+        self.causal = causal
+        self.softmax_scale = softmax_scale
+        self.dropout_p = attention_dropout
+
+    def forward(self, q, k, v):
+        """Implements the multihead softmax attention.
+        Arguments
+        ---------
+            q, k, v: The tensor containing the query, key, and value. (B, S, H, D)
+        """
+        assert all((i.dtype in [torch.float16, torch.bfloat16] for i in (q, k, v)))
+        assert all((i.is_cuda for i in (q, k, v)))
+
+        batch_size, seqlen_q = q.shape[0], q.shape[1]
+        seqlen_k = k.shape[1]
+
+        q, k, v = [rearrange(x, 'b s ... -> (b s) ...') for x in [q, k, v]]
+        cu_seqlens_q = torch.arange(0, (batch_size + 1) * seqlen_q, step=seqlen_q, dtype=torch.int32,
+                                    device=q.device)
+        self.training = False
+        if self.training:
+            # during training q,k,v always have same seqlen
+            assert seqlen_k == seqlen_q
+
+            is_causal = self.causal
+            cu_seqlens_k = cu_seqlens_q
+            dropout_p = self.dropout_p
+        else:
+            # turn off FA causal mask after first inference autoregressive iteration
+            # only on first autoregressive step q,k,v have same seqlen
+            is_causal = seqlen_q == seqlen_k
+            cu_seqlens_k = torch.arange(0, (batch_size + 1) * seqlen_k, step=seqlen_k, dtype=torch.int32,
+                                        device=q.device)
+            dropout_p = 0
+
+        output = flash_attn_unpadded_func(
+            q, k, v, cu_seqlens_q, cu_seqlens_k, seqlen_q, seqlen_k,
+            dropout_p=dropout_p,
+            softmax_scale=self.softmax_scale, causal=is_causal
+        )
+
+        output = rearrange(output, '(b s) ... -> b s ...', b=batch_size)
+        return output
+
+
+def _make_causal_mask(
+        input_ids_shape: torch.Size, device: torch.device, past_key_values_length: int
+) -> torch.BoolTensor:
+    """
+    Make causal mask used for self-attention.
+    """
+    batch_size, target_length = input_ids_shape
+    mask = torch.empty((target_length, target_length + past_key_values_length), dtype=torch.bool, device=device)
+    # ONNX doesn't support `torch.Tensor.triu` properly, thus we use this workaround
+    seq_ids = torch.arange(target_length, device=device)
+    mask[:, past_key_values_length:] = seq_ids[:, None] < seq_ids[None, :]
+
+    if past_key_values_length > 0:
+        mask[:, :past_key_values_length] = False
+
+    expanded_mask = mask[None, None, :, :].expand(batch_size, 1, target_length, target_length + past_key_values_length)
+    return expanded_mask
+
+
+def _expand_mask(mask: torch.Tensor, tgt_length: int) -> torch.BoolTensor:
+    """
+    Expands attention_mask from `[batch_size, src_length]` to `[batch_size, 1, tgt_length, src_length]`.
+    """
+    batch_size, src_length = mask.shape
+    tgt_length = tgt_length if tgt_length is not None else src_length
+
+    expanded_mask = ~(mask[:, None, None, :].to(torch.bool))
+    return expanded_mask.expand(batch_size, 1, tgt_length, src_length)
+
+
+def dropout_add(x: torch.Tensor, residual: torch.Tensor, prob: float, training: bool) -> torch.Tensor:
+    """
+    Dropout add function
+
+    Args:
+        x (`torch.tensor`, *required*):
+            input tensor
+        residual (`torch.tensor`, *required*):
+            residual tensor
+        prob (`float`, *required*):
+            dropout probability
+        training (`bool`, *required*):
+            training mode
+    """
+    out = F.dropout(x, p=prob, training=training)
+    out = residual + out
+    return out
+
+
+def telechat_gelu_forward(x: torch.Tensor) -> torch.Tensor:
+    """
+    Custom bias GELU function. Adapted from Megatron-DeepSpeed code. Here we use a simple implementation (inference) to
+    make the model jitable.
+
+    Args:
+        x (`torch.tensor`, *required*):
+            input hidden states
+    """
+    return x * 0.5 * (1.0 + torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)))
+
+
+def telechat_gelu_back(g: torch.Tensor, x: torch.Tensor) -> torch.Tensor:
+    """
+    gradient of tanh approximation of gelu gradient of actual gelu is: 0.5 * (1. + torch.erf(x * 0.70710678)) +
+    0.3989423 * x * torch.exp(-0.5 * x * x)
+
+    Args:
+        g (`torch.tensor`, *required*):
+            gradient output tensor
+        x (`torch.tensor`, *required*):
+            input tensor
+    """
+    x = x[0]  # x is a tuple of 1 element, needs to unpack it first
+    tanh_out = torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x))
+    # sqrt(2/pi) * 3 * 0.044715 -> 0.1070322243
+    ff = 0.5 * x * ((1 - tanh_out * tanh_out) * (0.79788456 + 0.1070322243 * x * x)) + 0.5 * (1 + tanh_out)
+    return ff * g
+
+
+class GeLUFunction(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, input: torch.Tensor) -> torch.Tensor:
+        ctx.save_for_backward(input)
+        return telechat_gelu_forward(input)
+
+    @staticmethod
+    def backward(ctx, grad_output: torch.Tensor) -> torch.Tensor:
+        input = ctx.saved_tensors
+        tmp = telechat_gelu_back(grad_output, input)
+        return tmp
+
+
+class TelechatGelu(nn.Module):
+    """
+    TelechatBiasGelu wrapper function that make use of the simple function on inference mode to make the model
+    torchscriptable and use the autograd function in training mode to get the accurate results of the gradients Partly
+    copied from Megatron-DeepSpeed code and adapted for our needs
+
+    See here why autograd functions are not torchscriptable: https://github.com/pytorch/pytorch/issues/22329
+    """
+
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if self.training:
+            return GeLUFunction.apply(x)
+        else:
+            return telechat_gelu_forward(x)
+
+
+class TelechatAttention(nn.Module):
+    def __init__(self, config: Telechat2Config, layer_idx):
+        super().__init__()
+        self.kv_cache = None
+        self.layer_idx = layer_idx
+
+        self.hidden_size = config.hidden_size
+        self.num_heads = config.n_head
+        self.head_dim = self.hidden_size // self.num_heads
+        self.split_size = self.hidden_size
+        self.hidden_dropout = config.hidden_dropout
+        self.config = config
+
+        if self.head_dim * self.num_heads != self.hidden_size:
+            raise ValueError(
+                f"`hidden_size` must be divisible by num_heads (got `hidden_size`: {self.hidden_size} and `num_heads`:"
+                f" {self.num_heads})."
+            )
+
+        # Layer-wise attention scaling
+        self.inv_norm_factor = 1.0 / math.sqrt(self.head_dim)
+        self.beta = 1.0
+
+        self.num_key_value_heads = config.num_key_value_heads if config.num_key_value_heads else self.num_heads
+        self.kv_projection_size = self.head_dim * self.num_key_value_heads
+        self.num_key_value_groups = self.num_heads // self.num_key_value_heads
+        self.query = nn.Linear(self.hidden_size, self.hidden_size, bias=False)
+        self.key_value = nn.Linear(self.hidden_size, self.kv_projection_size * 2, bias=False)
+        self.dense = nn.Linear(self.hidden_size, self.hidden_size)
+        self.attention_dropout = nn.Dropout(config.attention_dropout)
+        self.rotary_emb = RotaryEmbedding(self.head_dim, config=config)
+
+        self.core_attention_flash = FlashSelfAttention(
+            causal=True, attention_dropout=config.attention_dropout
+        )
+
+        self.last_key_layer = None
+        # logn_list = [math.log(i, 4096) if i > 4096 else 1 for i in range(1, 32768)]
+        # self.logn_tensor = torch.tensor(logn_list)[None, :, None, None].half().cuda()
+
+    def repeat_kv(self, hidden_states, n_rep):
+        slen, batch, num_key_value_heads_per_partition, head_dim = hidden_states.shape
+        if n_rep == 1:
+            return hidden_states
+        hidden_states = hidden_states[:, :, :, None, :].expand(slen, batch, num_key_value_heads_per_partition, n_rep,
+                                                               head_dim)
+        return hidden_states.reshape(slen, batch, num_key_value_heads_per_partition * n_rep, head_dim)
+
+    def split_tensor_along_last_dim(self,
+                                    tensor: torch.Tensor,
+                                    num_partitions: int,
+                                    contiguous_split_chunks: bool = False,
+                                    ):
+
+        # Get the size and dimension.
+        last_dim = tensor.dim() - 1
+        last_dim_size = tensor.size()[last_dim] // num_partitions
+        # Split.
+        tensor_list = torch.split(tensor, last_dim_size, dim=last_dim)
+        # Note: torch.split does not create contiguous tensors by default.
+        if contiguous_split_chunks:
+            return tuple(chunk.contiguous() for chunk in tensor_list)
+
+        return tensor_list
+
+    def _merge_heads(self, x: torch.Tensor) -> torch.Tensor:
+        batch_size_and_num_heads, seq_length, _ = x.shape
+        batch_size = batch_size_and_num_heads // self.num_heads
+        x = x.view(batch_size, self.num_heads, seq_length, self.head_dim)
+        x = x.permute(0, 2, 1, 3)
+        return x.reshape(batch_size, seq_length, self.num_heads * self.head_dim)
+
+    def forward(
+            self,
+            hidden_states: torch.Tensor,
+            residual: torch.Tensor,
+            attention_mask: torch.Tensor,
+            layer_past: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+            use_cache: bool = False,
+            output_attentions: bool = False,
+    ):
+        hidden_states = hidden_states.transpose(1, 0)
+        query_layer = self.query(hidden_states)
+        new_tensor_shape = query_layer.size()[:-1] + \
+                           (self.num_heads,
+                            self.head_dim)
+        query_layer = query_layer.view(*new_tensor_shape)
+
+        mixed_kv_layer = self.key_value(hidden_states)
+        new_tensor_shape = mixed_kv_layer.size()[:-1] + \
+                           (self.num_key_value_heads,
+                            2 * self.head_dim)
+        mixed_kv_layer = mixed_kv_layer.view(*new_tensor_shape)
+        (key_layer, value_layer) = self.split_tensor_along_last_dim(mixed_kv_layer, 2)
+
+        output_size = (query_layer.size(1),
+                       query_layer.size(2),
+                       query_layer.size(0),
+                       key_layer.size(0),
+                       key_layer.size(2)
+                       )
+
+        query_layer = query_layer.view(output_size[2], output_size[0] * output_size[1], -1)
+        key_layer = key_layer.view(output_size[3], output_size[0] * output_size[4], -1)
+
+        apply_rotary_fn = apply_rotary_pos_emb_torch
+
+        seq_len = key_layer.shape[0]
+        offset = 0
+
+        if use_cache and layer_past != None:
+            past_key, past_value = layer_past
+            offset = past_key.shape[0]
+            seq_len += offset
+
+        cos, sin = self.rotary_emb(value_layer, seq_len=seq_len)
+
+        query_layer, key_layer = apply_rotary_fn(query_layer, key_layer, cos, sin, offset=offset)
+        if use_cache:
+            if layer_past != None:
+                past_key, past_value = layer_past
+                key_layer = torch.cat((past_key, key_layer[-1, ...].unsqueeze(0)), dim=0)
+                value_layer = torch.cat((past_value, value_layer[-1, ...].unsqueeze(0)), dim=0)
+            layer_past = key_layer, value_layer
+
+        s_value, bz, kv_head, dim = value_layer.shape
+        s_key = key_layer.shape[0]
+        s_query = query_layer.shape[0]
+        q_head = output_size[1]
+
+        query_layer = query_layer.reshape((s_query, bz, q_head, dim))
+        key_layer = key_layer.reshape((s_key, bz, kv_head, dim))
+
+        key_layer = self.repeat_kv(key_layer, self.num_key_value_groups)
+        value_layer = self.repeat_kv(value_layer, self.num_key_value_groups)
+
+        if self.config.flash_attn:
+            q, k, v = [rearrange(x, 's b ... -> b s ...').contiguous() for x in
+                       (query_layer, key_layer, value_layer)]
+            context_layer = self.core_attention_flash(q, k, v)
+            context_layer = rearrange(context_layer, 'b s h d -> b s (h d)').contiguous()
+        else:
+            ##[sq, b, np, hn] -> [sq, b * np, hn]
+            query_layer = query_layer.reshape(s_query, bz * self.num_heads, dim)
+            # [sk, b, np, hn] -> [sk, b * np, hn]
+            key_layer = key_layer.reshape(s_key, bz * self.num_heads, dim)
+            matmul_result = self.inv_norm_factor * torch.einsum('bik,bkj->bij', query_layer.transpose(0, 1),
+                                                                key_layer.transpose(0, 1).transpose(1, 2))
+
+            attention_scores = matmul_result.view(bz, self.num_heads, s_query, s_key)
+
+            input_dtype = attention_scores.dtype
+            if input_dtype == torch.float16:
+                attention_scores = attention_scores.to(torch.float)
+            attn_weights = torch.masked_fill(attention_scores, attention_mask, torch.finfo(attention_scores.dtype).min)
+            attention_probs = F.softmax(attn_weights, dim=-1).to(input_dtype)  ##dtype = torch.float32
+            attention_probs = self.attention_dropout(attention_probs)
+            attention_probs_reshaped = attention_probs.view(bz * self.num_heads, s_query, s_key)
+
+            value_layer = value_layer.reshape(s_key, bz * self.num_heads, dim)
+            context_layer = torch.bmm(attention_probs_reshaped, value_layer.transpose(0, 1))
+            context_layer = self._merge_heads(context_layer)
+        output_tensor = self.dense(context_layer)
+
+        output_tensor = dropout_add(output_tensor, residual, self.hidden_dropout, self.training)
+        present = None
+        outputs = (output_tensor, present)
+        if output_attentions:
+            outputs += (attention_probs,)
+
+        return output_tensor, layer_past
+
+
+class TelechatMLP(nn.Module):
+    def __init__(self, config: Telechat2Config):
+        super().__init__()
+        hidden_size = config.hidden_size
+        self.gate_proj = nn.Linear(hidden_size, config.ffn_hidden_size, bias=False)
+        self.up_proj = nn.Linear(hidden_size, config.ffn_hidden_size, bias=False)
+        self.down_proj = nn.Linear(config.ffn_hidden_size, hidden_size, bias=True)
+        self.hidden_dropout = config.hidden_dropout
+
+    def forward(self, hidden_states: torch.Tensor, residual: torch.Tensor) -> torch.Tensor:
+        intermediate_output = self.down_proj(F.silu(self.gate_proj(hidden_states)) * self.up_proj(hidden_states))
+        output = dropout_add(intermediate_output, residual, self.hidden_dropout, self.training)
+        return output
+
+
+class TelechatBlock(nn.Module):
+    def __init__(self, config: Telechat2Config, layer_idx):
+        super().__init__()
+        hidden_size = config.hidden_size
+
+        self.input_layernorm = MixedFusedRMSNorm(hidden_size, eps=config.layer_norm_epsilon)
+        self.num_heads = config.n_head
+        self.layer_idx = layer_idx
+        self.self_attention = TelechatAttention(config, layer_idx)
+        self.post_attention_layernorm = MixedFusedRMSNorm(hidden_size, eps=config.layer_norm_epsilon)
+
+        self.mlp = TelechatMLP(config)
+
+        self.apply_residual_connection_post_layernorm = config.apply_residual_connection_post_layernorm
+        self.hidden_dropout = config.hidden_dropout
+
+    def forward(
+            self,
+            hidden_states: torch.Tensor,
+            attention_mask: torch.Tensor,
+            layer_past: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+            use_cache: bool = False,
+            output_attentions: bool = False,
+    ):
+        layernorm_output = self.input_layernorm(hidden_states)
+        if self.apply_residual_connection_post_layernorm:
+            residual = layernorm_output
+        else:
+            residual = hidden_states
+
+        attn_outputs = self.self_attention(
+            layernorm_output,
+            residual,
+            layer_past=layer_past,
+            attention_mask=attention_mask,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+        )
+
+        attention_output = attn_outputs[0]
+        outputs = attn_outputs[1:]
+        layernorm_output = self.post_attention_layernorm(attention_output)
+
+        if self.apply_residual_connection_post_layernorm:
+            residual = layernorm_output
+        else:
+            residual = attention_output
+        output = self.mlp(layernorm_output, residual)
+
+        if use_cache:
+            outputs = (output,) + outputs
+        else:
+            outputs = (output,) + outputs[1:]
+
+        return outputs
+
+
+class TelechatPreTrainedModel(PreTrainedModel):
+    config_class = Telechat2Config
+    base_model_prefix = "transformer"
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["TelechatBlock"]
+    _skip_keys_device_placement = "past_key_values"
+
+    def __init__(self, *inputs, **kwargs):
+        super().__init__(*inputs, **kwargs)
+
+    def _init_weights(self, module: nn.Module):
+        """Initialize the weights."""
+        if isinstance(module, nn.Linear):
+            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
+            if module.bias is not None:
+                module.bias.data.zero_()
+
+        elif isinstance(module, nn.Embedding):
+            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
+            if module.padding_idx is not None:
+                module.weight.data[module.padding_idx].zero_()
+
+        elif isinstance(module, LayerNorm):
+            module.bias.data.zero_()
+            module.weight.data.fill_(1.0)
+
+    def _set_gradient_checkpointing(self, module: nn.Module, value: bool = False):
+        if isinstance(module, TelechatModel):
+            module.gradient_checkpointing = value
+
+
+class TelechatModel(TelechatPreTrainedModel):
+    def __init__(self, config: Telechat2Config):
+        super().__init__(config)
+
+        self.embed_dim = config.hidden_size
+        self.num_heads = config.n_head
+        self.config = config
+        self.word_embeddings = nn.Embedding(config.vocab_size, self.embed_dim)
+        if self.config.embed_layernorm:
+            self.word_embeddings_layernorm = MixedFusedRMSNorm(self.embed_dim, eps=config.layer_norm_epsilon)
+
+        self.h = nn.ModuleList([TelechatBlock(config, _) for _ in range(config.num_hidden_layers)])
+        self.ln_f = MixedFusedRMSNorm(self.embed_dim, eps=config.layer_norm_epsilon)
+        self.gradient_checkpointing = False
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.word_embeddings
+
+    def _prepare_attn_mask(
+            self, attention_mask: torch.Tensor, input_shape: Tuple[int, int], past_key_values_length: int
+    ) -> torch.BoolTensor:
+        combined_attention_mask = None
+        device = attention_mask.device
+        _, src_length = input_shape
+
+        if src_length > 1:
+            combined_attention_mask = _make_causal_mask(
+                input_shape, device=device, past_key_values_length=past_key_values_length
+            )
+        expanded_attn_mask = _expand_mask(attention_mask, tgt_length=src_length)
+        combined_attention_mask = (
+            expanded_attn_mask if combined_attention_mask is None else expanded_attn_mask | combined_attention_mask
+        )
+
+        return combined_attention_mask
+
+    def set_input_embeddings(self, new_embeddings: torch.Tensor):
+        self.word_embeddings = new_embeddings
+
+    def forward(
+            self,
+            input_ids: Optional[torch.LongTensor] = None,
+            past_key_values: Optional[Tuple[Tuple[torch.Tensor, torch.Tensor], ...]] = None,
+            attention_mask: Optional[torch.Tensor] = None,
+            inputs_embeds: 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,
+            **deprecated_arguments,
+    ) -> Union[Tuple[torch.Tensor, ...], BaseModelOutputWithPastAndCrossAttentions]:
+
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        use_cache = use_cache if use_cache is not None else self.config.use_cache
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        if input_ids is not None:
+            batch_size, seq_length = input_ids.shape
+        elif inputs_embeds is not None:
+            batch_size, seq_length, _ = inputs_embeds.shape
+
+        if past_key_values is None:
+            past_key_values = tuple([None] * len(self.h))
+        # input_ids = torch.load("Megatron-LM-0624-3B/tensors/input_ids.pt").to(input_ids.device)
+        if inputs_embeds is None:
+            inputs_embeds = self.word_embeddings(input_ids)
+        hidden_states = inputs_embeds
+        # print(f"[INFO_Telechat]: inputs_embeds={inputs_embeds}")
+        if self.config.embed_layernorm:
+            hidden_states = self.word_embeddings_layernorm(inputs_embeds)
+
+        presents = () if use_cache else None
+        all_self_attentions = () if output_attentions else None
+        all_hidden_states = () if output_hidden_states else None
+
+        if self.gradient_checkpointing and self.training:
+            if use_cache:
+                use_cache = False
+
+        seq_length_with_past = seq_length
+        past_key_values_length = 0
+        if past_key_values[0] is not None:
+            past_key_values_length = past_key_values[0][0].shape[2]
+            seq_length_with_past = seq_length_with_past + past_key_values_length
+        if attention_mask is None:
+            attention_mask = torch.ones((batch_size, seq_length_with_past), device=hidden_states.device)
+        else:
+            attention_mask = attention_mask.to(hidden_states.device)
+        causal_mask = self._prepare_attn_mask(
+            attention_mask,
+            input_shape=(batch_size, seq_length),
+            past_key_values_length=past_key_values_length,
+        )
+
+        # print(f"[INFO_Telechat]: word_embeddings_layernorm={hidden_states}")
+        for i, (block, layer_past) in enumerate(zip(self.h, past_key_values)):
+            if output_hidden_states:
+                all_hidden_states = all_hidden_states + (hidden_states,)
+
+            if self.gradient_checkpointing and self.training:
+
+                def create_custom_forward(module):
+                    def custom_forward(*inputs):
+                        # None for past_key_value
+                        return module(*inputs, use_cache=use_cache, output_attentions=output_attentions)
+
+                    return custom_forward
+
+                outputs = torch.utils.checkpoint.checkpoint(
+                    create_custom_forward(block),
+                    hidden_states,
+                    causal_mask,
+                    layer_past,
+                )
+            else:
+                outputs = block(
+                    hidden_states,
+                    layer_past=layer_past,
+                    attention_mask=causal_mask,
+                    use_cache=use_cache,
+                    output_attentions=output_attentions,
+                )
+
+            # print(f"[INFO_Telechat]: outputs{i}={outputs}")
+            hidden_states = outputs[0]
+            if use_cache is True:
+                presents = presents + (outputs[1],)
+
+            if output_attentions:
+                all_self_attentions = all_self_attentions + (outputs[2 if use_cache else 1],)
+        hidden_states = self.ln_f(hidden_states)
+        # print(f"[INFO_Telechat]: hidden_states={hidden_states}")
+        # ref = torch.load("Megatron-LM-0624-3B/tensors/final_layernorm.pt")
+        # print(hidden_states.squeeze()[2048:])
+        # print(ref.squeeze())
+        # print(torch.max(hidden_states.squeeze()[2048:] - ref.squeeze().to(hidden_states.device)))
+        # exit()
+        # print(ref.shape,hidden_states.shape)
+        # print(hidden_states)
+        # exit()
+        if output_hidden_states:
+            all_hidden_states = all_hidden_states + (hidden_states,)
+        if not return_dict:
+            return tuple(v for v in [hidden_states, presents, all_hidden_states, all_self_attentions] if v is not None)
+        return BaseModelOutputWithPastAndCrossAttentions(
+            last_hidden_state=hidden_states,
+            past_key_values=presents,
+            hidden_states=all_hidden_states,
+            attentions=all_self_attentions,
+        )
+
+
+class Telechat2ForCausalLM(TelechatPreTrainedModel):
+    # _tied_weights_keys = ["lm_head.weight"]
+    _keys_to_ignore_on_load_missing = [r"lm_head.weight"]
+
+    def __init__(self, config: Telechat2Config):
+        super().__init__(config)
+        self.transformer = TelechatModel(config)
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        self.post_init()
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def set_output_embeddings(self, new_embeddings: torch.Tensor):
+        self.lm_head = new_embeddings
+
+    def prepare_inputs_for_generation(
+            self,
+            input_ids: torch.LongTensor,
+            past_key_values: Optional[torch.Tensor] = None,
+            attention_mask: Optional[torch.Tensor] = None,
+            inputs_embeds: Optional[torch.Tensor] = None,
+            **kwargs,
+    ) -> dict:
+        if past_key_values:
+            input_ids = input_ids[:, -1].unsqueeze(-1)
+        if inputs_embeds is not None and past_key_values is None:
+            model_inputs = {"inputs_embeds": inputs_embeds}
+        else:
+            model_inputs = {"input_ids": input_ids}
+
+        model_inputs.update(
+            {
+                "past_key_values": past_key_values,
+                "use_cache": kwargs.get("use_cache"),
+                "attention_mask": attention_mask,
+            }
+        )
+        return model_inputs
+
+    def forward(
+            self,
+            input_ids: Optional[torch.LongTensor] = None,
+            past_key_values: Optional[Tuple[Tuple[torch.Tensor, torch.Tensor], ...]] = None,
+            attention_mask: Optional[torch.Tensor] = None,
+            inputs_embeds: Optional[torch.Tensor] = None,
+            labels: 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,
+            **deprecated_arguments,
+    ) -> Union[Tuple[torch.Tensor], CausalLMOutputWithCrossAttentions]:
+
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        transformer_outputs = self.transformer(
+            input_ids,
+            past_key_values=past_key_values,
+            attention_mask=attention_mask,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+        hidden_states = transformer_outputs[0]
+        lm_logits = self.lm_head(hidden_states)
+
+        loss = None
+        if labels is not None:
+            labels = labels.to(lm_logits.device)
+            shift_logits = lm_logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+            batch_size, seq_length, vocab_size = shift_logits.shape
+            loss_fct = CrossEntropyLoss()
+            loss = loss_fct(
+                shift_logits.view(batch_size * seq_length, vocab_size), shift_labels.view(batch_size * seq_length)
+            )
+
+        if not return_dict:
+            output = (lm_logits,) + transformer_outputs[1:]
+            return ((loss,) + output) if loss is not None else output
+
+        return CausalLMOutputWithCrossAttentions(
+            loss=loss,
+            logits=lm_logits,
+            past_key_values=transformer_outputs.past_key_values,
+            hidden_states=transformer_outputs.hidden_states,
+            attentions=transformer_outputs.attentions,
+        )

pytorch_model.bin.index.json

@@ -0,0 +1,310 @@
+{
+  "metadata": {
+    "total_size": 15234703360
+  },
+  "weight_map": {
+    "lm_head.weight": "pytorch_model_00004-of-00004.bin",
+    "transformer.h.0.input_layernorm.weight": "pytorch_model_00001-of-00004.bin",
+    "transformer.h.0.mlp.down_proj.bias": "pytorch_model_00001-of-00004.bin",
+    "transformer.h.0.mlp.down_proj.weight": "pytorch_model_00001-of-00004.bin",
+    "transformer.h.0.mlp.gate_proj.weight": "pytorch_model_00001-of-00004.bin",
+    "transformer.h.0.mlp.up_proj.weight": "pytorch_model_00001-of-00004.bin",
+    "transformer.h.0.post_attention_layernorm.weight": "pytorch_model_00001-of-00004.bin",
+    "transformer.h.0.self_attention.dense.bias": "pytorch_model_00001-of-00004.bin",
+    "transformer.h.0.self_attention.dense.weight": "pytorch_model_00001-of-00004.bin",
+    "transformer.h.0.self_attention.key_value.weight": "pytorch_model_00001-of-00004.bin",
+    "transformer.h.0.self_attention.query.weight": "pytorch_model_00001-of-00004.bin",
+    "transformer.h.1.input_layernorm.weight": "pytorch_model_00001-of-00004.bin",
+    "transformer.h.1.mlp.down_proj.bias": "pytorch_model_00001-of-00004.bin",
+    "transformer.h.1.mlp.down_proj.weight": "pytorch_model_00001-of-00004.bin",
+    "transformer.h.1.mlp.gate_proj.weight": "pytorch_model_00001-of-00004.bin",
+    "transformer.h.1.mlp.up_proj.weight": "pytorch_model_00001-of-00004.bin",
+    "transformer.h.1.post_attention_layernorm.weight": "pytorch_model_00001-of-00004.bin",
+    "transformer.h.1.self_attention.dense.bias": "pytorch_model_00001-of-00004.bin",
+    "transformer.h.1.self_attention.dense.weight": "pytorch_model_00001-of-00004.bin",
+    "transformer.h.1.self_attention.key_value.weight": "pytorch_model_00001-of-00004.bin",
+    "transformer.h.1.self_attention.query.weight": "pytorch_model_00001-of-00004.bin",
+    "transformer.h.10.input_layernorm.weight": "pytorch_model_00002-of-00004.bin",
+    "transformer.h.10.mlp.down_proj.bias": "pytorch_model_00002-of-00004.bin",
+    "transformer.h.10.mlp.down_proj.weight": "pytorch_model_00002-of-00004.bin",
+    "transformer.h.10.mlp.gate_proj.weight": "pytorch_model_00002-of-00004.bin",
+    "transformer.h.10.mlp.up_proj.weight": "pytorch_model_00002-of-00004.bin",
+    "transformer.h.10.post_attention_layernorm.weight": "pytorch_model_00002-of-00004.bin",
+    "transformer.h.10.self_attention.dense.bias": "pytorch_model_00002-of-00004.bin",
+    "transformer.h.10.self_attention.dense.weight": "pytorch_model_00002-of-00004.bin",
+    "transformer.h.10.self_attention.key_value.weight": "pytorch_model_00002-of-00004.bin",
+    "transformer.h.10.self_attention.query.weight": "pytorch_model_00002-of-00004.bin",
+    "transformer.h.11.input_layernorm.weight": "pytorch_model_00002-of-00004.bin",
+    "transformer.h.11.mlp.down_proj.bias": "pytorch_model_00002-of-00004.bin",
+    "transformer.h.11.mlp.down_proj.weight": "pytorch_model_00002-of-00004.bin",
+    "transformer.h.11.mlp.gate_proj.weight": "pytorch_model_00002-of-00004.bin",
+    "transformer.h.11.mlp.up_proj.weight": "pytorch_model_00002-of-00004.bin",
+    "transformer.h.11.post_attention_layernorm.weight": "pytorch_model_00002-of-00004.bin",
+    "transformer.h.11.self_attention.dense.bias": "pytorch_model_00002-of-00004.bin",
+    "transformer.h.11.self_attention.dense.weight": "pytorch_model_00002-of-00004.bin",
+    "transformer.h.11.self_attention.key_value.weight": "pytorch_model_00002-of-00004.bin",
+    "transformer.h.11.self_attention.query.weight": "pytorch_model_00002-of-00004.bin",
+    "transformer.h.12.input_layernorm.weight": "pytorch_model_00002-of-00004.bin",
+    "transformer.h.12.mlp.down_proj.bias": "pytorch_model_00002-of-00004.bin",
+    "transformer.h.12.mlp.down_proj.weight": "pytorch_model_00002-of-00004.bin",
+    "transformer.h.12.mlp.gate_proj.weight": "pytorch_model_00002-of-00004.bin",
+    "transformer.h.12.mlp.up_proj.weight": "pytorch_model_00002-of-00004.bin",
+    "transformer.h.12.post_attention_layernorm.weight": "pytorch_model_00002-of-00004.bin",
+    "transformer.h.12.self_attention.dense.bias": "pytorch_model_00002-of-00004.bin",
+    "transformer.h.12.self_attention.dense.weight": "pytorch_model_00002-of-00004.bin",
+    "transformer.h.12.self_attention.key_value.weight": "pytorch_model_00002-of-00004.bin",
+    "transformer.h.12.self_attention.query.weight": "pytorch_model_00002-of-00004.bin",
+    "transformer.h.13.input_layernorm.weight": "pytorch_model_00002-of-00004.bin",
+    "transformer.h.13.mlp.down_proj.bias": "pytorch_model_00002-of-00004.bin",
+    "transformer.h.13.mlp.down_proj.weight": "pytorch_model_00002-of-00004.bin",
+    "transformer.h.13.mlp.gate_proj.weight": "pytorch_model_00002-of-00004.bin",
+    "transformer.h.13.mlp.up_proj.weight": "pytorch_model_00002-of-00004.bin",
+    "transformer.h.13.post_attention_layernorm.weight": "pytorch_model_00002-of-00004.bin",
+    "transformer.h.13.self_attention.dense.bias": "pytorch_model_00002-of-00004.bin",
+    "transformer.h.13.self_attention.dense.weight": "pytorch_model_00002-of-00004.bin",
+    "transformer.h.13.self_attention.key_value.weight": "pytorch_model_00002-of-00004.bin",
+    "transformer.h.13.self_attention.query.weight": "pytorch_model_00002-of-00004.bin",
+    "transformer.h.14.input_layernorm.weight": "pytorch_model_00002-of-00004.bin",
+    "transformer.h.14.mlp.down_proj.bias": "pytorch_model_00002-of-00004.bin",
+    "transformer.h.14.mlp.down_proj.weight": "pytorch_model_00002-of-00004.bin",
+    "transformer.h.14.mlp.gate_proj.weight": "pytorch_model_00002-of-00004.bin",
+    "transformer.h.14.mlp.up_proj.weight": "pytorch_model_00002-of-00004.bin",
+    "transformer.h.14.post_attention_layernorm.weight": "pytorch_model_00002-of-00004.bin",
+    "transformer.h.14.self_attention.dense.bias": "pytorch_model_00002-of-00004.bin",
+    "transformer.h.14.self_attention.dense.weight": "pytorch_model_00002-of-00004.bin",
+    "transformer.h.14.self_attention.key_value.weight": "pytorch_model_00002-of-00004.bin",
+    "transformer.h.14.self_attention.query.weight": "pytorch_model_00002-of-00004.bin",
+    "transformer.h.15.input_layernorm.weight": "pytorch_model_00003-of-00004.bin",
+    "transformer.h.15.mlp.down_proj.bias": "pytorch_model_00003-of-00004.bin",
+    "transformer.h.15.mlp.down_proj.weight": "pytorch_model_00003-of-00004.bin",
+    "transformer.h.15.mlp.gate_proj.weight": "pytorch_model_00003-of-00004.bin",
+    "transformer.h.15.mlp.up_proj.weight": "pytorch_model_00003-of-00004.bin",
+    "transformer.h.15.post_attention_layernorm.weight": "pytorch_model_00003-of-00004.bin",
+    "transformer.h.15.self_attention.dense.bias": "pytorch_model_00003-of-00004.bin",
+    "transformer.h.15.self_attention.dense.weight": "pytorch_model_00003-of-00004.bin",
+    "transformer.h.15.self_attention.key_value.weight": "pytorch_model_00003-of-00004.bin",
+    "transformer.h.15.self_attention.query.weight": "pytorch_model_00003-of-00004.bin",
+    "transformer.h.16.input_layernorm.weight": "pytorch_model_00003-of-00004.bin",
+    "transformer.h.16.mlp.down_proj.bias": "pytorch_model_00003-of-00004.bin",
+    "transformer.h.16.mlp.down_proj.weight": "pytorch_model_00003-of-00004.bin",
+    "transformer.h.16.mlp.gate_proj.weight": "pytorch_model_00003-of-00004.bin",
+    "transformer.h.16.mlp.up_proj.weight": "pytorch_model_00003-of-00004.bin",
+    "transformer.h.16.post_attention_layernorm.weight": "pytorch_model_00003-of-00004.bin",
+    "transformer.h.16.self_attention.dense.bias": "pytorch_model_00003-of-00004.bin",
+    "transformer.h.16.self_attention.dense.weight": "pytorch_model_00003-of-00004.bin",
+    "transformer.h.16.self_attention.key_value.weight": "pytorch_model_00003-of-00004.bin",
+    "transformer.h.16.self_attention.query.weight": "pytorch_model_00003-of-00004.bin",
+    "transformer.h.17.input_layernorm.weight": "pytorch_model_00003-of-00004.bin",
+    "transformer.h.17.mlp.down_proj.bias": "pytorch_model_00003-of-00004.bin",
+    "transformer.h.17.mlp.down_proj.weight": "pytorch_model_00003-of-00004.bin",
+    "transformer.h.17.mlp.gate_proj.weight": "pytorch_model_00003-of-00004.bin",
+    "transformer.h.17.mlp.up_proj.weight": "pytorch_model_00003-of-00004.bin",
+    "transformer.h.17.post_attention_layernorm.weight": "pytorch_model_00003-of-00004.bin",
+    "transformer.h.17.self_attention.dense.bias": "pytorch_model_00003-of-00004.bin",
+    "transformer.h.17.self_attention.dense.weight": "pytorch_model_00003-of-00004.bin",
+    "transformer.h.17.self_attention.key_value.weight": "pytorch_model_00003-of-00004.bin",
+    "transformer.h.17.self_attention.query.weight": "pytorch_model_00003-of-00004.bin",
+    "transformer.h.18.input_layernorm.weight": "pytorch_model_00003-of-00004.bin",
+    "transformer.h.18.mlp.down_proj.bias": "pytorch_model_00003-of-00004.bin",
+    "transformer.h.18.mlp.down_proj.weight": "pytorch_model_00003-of-00004.bin",
+    "transformer.h.18.mlp.gate_proj.weight": "pytorch_model_00003-of-00004.bin",
+    "transformer.h.18.mlp.up_proj.weight": "pytorch_model_00003-of-00004.bin",
+    "transformer.h.18.post_attention_layernorm.weight": "pytorch_model_00003-of-00004.bin",
+    "transformer.h.18.self_attention.dense.bias": "pytorch_model_00003-of-00004.bin",
+    "transformer.h.18.self_attention.dense.weight": "pytorch_model_00003-of-00004.bin",
+    "transformer.h.18.self_attention.key_value.weight": "pytorch_model_00003-of-00004.bin",
+    "transformer.h.18.self_attention.query.weight": "pytorch_model_00003-of-00004.bin",
+    "transformer.h.19.input_layernorm.weight": "pytorch_model_00003-of-00004.bin",
+    "transformer.h.19.mlp.down_proj.bias": "pytorch_model_00003-of-00004.bin",
+    "transformer.h.19.mlp.down_proj.weight": "pytorch_model_00003-of-00004.bin",
+    "transformer.h.19.mlp.gate_proj.weight": "pytorch_model_00003-of-00004.bin",
+    "transformer.h.19.mlp.up_proj.weight": "pytorch_model_00003-of-00004.bin",
+    "transformer.h.19.post_attention_layernorm.weight": "pytorch_model_00003-of-00004.bin",
+    "transformer.h.19.self_attention.dense.bias": "pytorch_model_00003-of-00004.bin",
+    "transformer.h.19.self_attention.dense.weight": "pytorch_model_00003-of-00004.bin",
+    "transformer.h.19.self_attention.key_value.weight": "pytorch_model_00003-of-00004.bin",
+    "transformer.h.19.self_attention.query.weight": "pytorch_model_00003-of-00004.bin",
+    "transformer.h.2.input_layernorm.weight": "pytorch_model_00001-of-00004.bin",
+    "transformer.h.2.mlp.down_proj.bias": "pytorch_model_00001-of-00004.bin",
+    "transformer.h.2.mlp.down_proj.weight": "pytorch_model_00001-of-00004.bin",
+    "transformer.h.2.mlp.gate_proj.weight": "pytorch_model_00001-of-00004.bin",
+    "transformer.h.2.mlp.up_proj.weight": "pytorch_model_00001-of-00004.bin",
+    "transformer.h.2.post_attention_layernorm.weight": "pytorch_model_00001-of-00004.bin",
+    "transformer.h.2.self_attention.dense.bias": "pytorch_model_00001-of-00004.bin",
+    "transformer.h.2.self_attention.dense.weight": "pytorch_model_00001-of-00004.bin",
+    "transformer.h.2.self_attention.key_value.weight": "pytorch_model_00001-of-00004.bin",
+    "transformer.h.2.self_attention.query.weight": "pytorch_model_00001-of-00004.bin",
+    "transformer.h.20.input_layernorm.weight": "pytorch_model_00003-of-00004.bin",
+    "transformer.h.20.mlp.down_proj.bias": "pytorch_model_00003-of-00004.bin",
+    "transformer.h.20.mlp.down_proj.weight": "pytorch_model_00003-of-00004.bin",
+    "transformer.h.20.mlp.gate_proj.weight": "pytorch_model_00003-of-00004.bin",
+    "transformer.h.20.mlp.up_proj.weight": "pytorch_model_00003-of-00004.bin",
+    "transformer.h.20.post_attention_layernorm.weight": "pytorch_model_00003-of-00004.bin",
+    "transformer.h.20.self_attention.dense.bias": "pytorch_model_00003-of-00004.bin",
+    "transformer.h.20.self_attention.dense.weight": "pytorch_model_00003-of-00004.bin",
+    "transformer.h.20.self_attention.key_value.weight": "pytorch_model_00003-of-00004.bin",
+    "transformer.h.20.self_attention.query.weight": "pytorch_model_00003-of-00004.bin",
+    "transformer.h.21.input_layernorm.weight": "pytorch_model_00003-of-00004.bin",
+    "transformer.h.21.mlp.down_proj.bias": "pytorch_model_00003-of-00004.bin",
+    "transformer.h.21.mlp.down_proj.weight": "pytorch_model_00003-of-00004.bin",
+    "transformer.h.21.mlp.gate_proj.weight": "pytorch_model_00003-of-00004.bin",
+    "transformer.h.21.mlp.up_proj.weight": "pytorch_model_00003-of-00004.bin",
+    "transformer.h.21.post_attention_layernorm.weight": "pytorch_model_00003-of-00004.bin",
+    "transformer.h.21.self_attention.dense.bias": "pytorch_model_00003-of-00004.bin",
+    "transformer.h.21.self_attention.dense.weight": "pytorch_model_00003-of-00004.bin",
+    "transformer.h.21.self_attention.key_value.weight": "pytorch_model_00003-of-00004.bin",
+    "transformer.h.21.self_attention.query.weight": "pytorch_model_00003-of-00004.bin",
+    "transformer.h.22.input_layernorm.weight": "pytorch_model_00003-of-00004.bin",
+    "transformer.h.22.mlp.down_proj.bias": "pytorch_model_00003-of-00004.bin",
+    "transformer.h.22.mlp.down_proj.weight": "pytorch_model_00003-of-00004.bin",
+    "transformer.h.22.mlp.gate_proj.weight": "pytorch_model_00003-of-00004.bin",
+    "transformer.h.22.mlp.up_proj.weight": "pytorch_model_00003-of-00004.bin",
+    "transformer.h.22.post_attention_layernorm.weight": "pytorch_model_00003-of-00004.bin",
+    "transformer.h.22.self_attention.dense.bias": "pytorch_model_00003-of-00004.bin",
+    "transformer.h.22.self_attention.dense.weight": "pytorch_model_00003-of-00004.bin",
+    "transformer.h.22.self_attention.key_value.weight": "pytorch_model_00003-of-00004.bin",
+    "transformer.h.22.self_attention.query.weight": "pytorch_model_00003-of-00004.bin",
+    "transformer.h.23.input_layernorm.weight": "pytorch_model_00004-of-00004.bin",
+    "transformer.h.23.mlp.down_proj.bias": "pytorch_model_00004-of-00004.bin",
+    "transformer.h.23.mlp.down_proj.weight": "pytorch_model_00004-of-00004.bin",
+    "transformer.h.23.mlp.gate_proj.weight": "pytorch_model_00004-of-00004.bin",
+    "transformer.h.23.mlp.up_proj.weight": "pytorch_model_00004-of-00004.bin",
+    "transformer.h.23.post_attention_layernorm.weight": "pytorch_model_00004-of-00004.bin",
+    "transformer.h.23.self_attention.dense.bias": "pytorch_model_00004-of-00004.bin",
+    "transformer.h.23.self_attention.dense.weight": "pytorch_model_00004-of-00004.bin",
+    "transformer.h.23.self_attention.key_value.weight": "pytorch_model_00004-of-00004.bin",
+    "transformer.h.23.self_attention.query.weight": "pytorch_model_00004-of-00004.bin",
+    "transformer.h.24.input_layernorm.weight": "pytorch_model_00004-of-00004.bin",
+    "transformer.h.24.mlp.down_proj.bias": "pytorch_model_00004-of-00004.bin",
+    "transformer.h.24.mlp.down_proj.weight": "pytorch_model_00004-of-00004.bin",
+    "transformer.h.24.mlp.gate_proj.weight": "pytorch_model_00004-of-00004.bin",
+    "transformer.h.24.mlp.up_proj.weight": "pytorch_model_00004-of-00004.bin",
+    "transformer.h.24.post_attention_layernorm.weight": "pytorch_model_00004-of-00004.bin",
+    "transformer.h.24.self_attention.dense.bias": "pytorch_model_00004-of-00004.bin",
+    "transformer.h.24.self_attention.dense.weight": "pytorch_model_00004-of-00004.bin",
+    "transformer.h.24.self_attention.key_value.weight": "pytorch_model_00004-of-00004.bin",
+    "transformer.h.24.self_attention.query.weight": "pytorch_model_00004-of-00004.bin",
+    "transformer.h.25.input_layernorm.weight": "pytorch_model_00004-of-00004.bin",
+    "transformer.h.25.mlp.down_proj.bias": "pytorch_model_00004-of-00004.bin",
+    "transformer.h.25.mlp.down_proj.weight": "pytorch_model_00004-of-00004.bin",
+    "transformer.h.25.mlp.gate_proj.weight": "pytorch_model_00004-of-00004.bin",
+    "transformer.h.25.mlp.up_proj.weight": "pytorch_model_00004-of-00004.bin",
+    "transformer.h.25.post_attention_layernorm.weight": "pytorch_model_00004-of-00004.bin",
+    "transformer.h.25.self_attention.dense.bias": "pytorch_model_00004-of-00004.bin",
+    "transformer.h.25.self_attention.dense.weight": "pytorch_model_00004-of-00004.bin",
+    "transformer.h.25.self_attention.key_value.weight": "pytorch_model_00004-of-00004.bin",
+    "transformer.h.25.self_attention.query.weight": "pytorch_model_00004-of-00004.bin",
+    "transformer.h.26.input_layernorm.weight": "pytorch_model_00004-of-00004.bin",
+    "transformer.h.26.mlp.down_proj.bias": "pytorch_model_00004-of-00004.bin",
+    "transformer.h.26.mlp.down_proj.weight": "pytorch_model_00004-of-00004.bin",
+    "transformer.h.26.mlp.gate_proj.weight": "pytorch_model_00004-of-00004.bin",
+    "transformer.h.26.mlp.up_proj.weight": "pytorch_model_00004-of-00004.bin",
+    "transformer.h.26.post_attention_layernorm.weight": "pytorch_model_00004-of-00004.bin",
+    "transformer.h.26.self_attention.dense.bias": "pytorch_model_00004-of-00004.bin",
+    "transformer.h.26.self_attention.dense.weight": "pytorch_model_00004-of-00004.bin",
+    "transformer.h.26.self_attention.key_value.weight": "pytorch_model_00004-of-00004.bin",
+    "transformer.h.26.self_attention.query.weight": "pytorch_model_00004-of-00004.bin",
+    "transformer.h.27.input_layernorm.weight": "pytorch_model_00004-of-00004.bin",
+    "transformer.h.27.mlp.down_proj.bias": "pytorch_model_00004-of-00004.bin",
+    "transformer.h.27.mlp.down_proj.weight": "pytorch_model_00004-of-00004.bin",
+    "transformer.h.27.mlp.gate_proj.weight": "pytorch_model_00004-of-00004.bin",
+    "transformer.h.27.mlp.up_proj.weight": "pytorch_model_00004-of-00004.bin",
+    "transformer.h.27.post_attention_layernorm.weight": "pytorch_model_00004-of-00004.bin",
+    "transformer.h.27.self_attention.dense.bias": "pytorch_model_00004-of-00004.bin",
+    "transformer.h.27.self_attention.dense.weight": "pytorch_model_00004-of-00004.bin",
+    "transformer.h.27.self_attention.key_value.weight": "pytorch_model_00004-of-00004.bin",
+    "transformer.h.27.self_attention.query.weight": "pytorch_model_00004-of-00004.bin",
+    "transformer.h.28.input_layernorm.weight": "pytorch_model_00004-of-00004.bin",
+    "transformer.h.28.mlp.down_proj.bias": "pytorch_model_00004-of-00004.bin",
+    "transformer.h.28.mlp.down_proj.weight": "pytorch_model_00004-of-00004.bin",
+    "transformer.h.28.mlp.gate_proj.weight": "pytorch_model_00004-of-00004.bin",
+    "transformer.h.28.mlp.up_proj.weight": "pytorch_model_00004-of-00004.bin",
+    "transformer.h.28.post_attention_layernorm.weight": "pytorch_model_00004-of-00004.bin",
+    "transformer.h.28.self_attention.dense.bias": "pytorch_model_00004-of-00004.bin",
+    "transformer.h.28.self_attention.dense.weight": "pytorch_model_00004-of-00004.bin",
+    "transformer.h.28.self_attention.key_value.weight": "pytorch_model_00004-of-00004.bin",
+    "transformer.h.28.self_attention.query.weight": "pytorch_model_00004-of-00004.bin",
+    "transformer.h.29.input_layernorm.weight": "pytorch_model_00004-of-00004.bin",
+    "transformer.h.29.mlp.down_proj.bias": "pytorch_model_00004-of-00004.bin",
+    "transformer.h.29.mlp.down_proj.weight": "pytorch_model_00004-of-00004.bin",
+    "transformer.h.29.mlp.gate_proj.weight": "pytorch_model_00004-of-00004.bin",
+    "transformer.h.29.mlp.up_proj.weight": "pytorch_model_00004-of-00004.bin",
+    "transformer.h.29.post_attention_layernorm.weight": "pytorch_model_00004-of-00004.bin",
+    "transformer.h.29.self_attention.dense.bias": "pytorch_model_00004-of-00004.bin",
+    "transformer.h.29.self_attention.dense.weight": "pytorch_model_00004-of-00004.bin",
+    "transformer.h.29.self_attention.key_value.weight": "pytorch_model_00004-of-00004.bin",
+    "transformer.h.29.self_attention.query.weight": "pytorch_model_00004-of-00004.bin",
+    "transformer.h.3.input_layernorm.weight": "pytorch_model_00001-of-00004.bin",
+    "transformer.h.3.mlp.down_proj.bias": "pytorch_model_00001-of-00004.bin",
+    "transformer.h.3.mlp.down_proj.weight": "pytorch_model_00001-of-00004.bin",
+    "transformer.h.3.mlp.gate_proj.weight": "pytorch_model_00001-of-00004.bin",
+    "transformer.h.3.mlp.up_proj.weight": "pytorch_model_00001-of-00004.bin",
+    "transformer.h.3.post_attention_layernorm.weight": "pytorch_model_00001-of-00004.bin",
+    "transformer.h.3.self_attention.dense.bias": "pytorch_model_00001-of-00004.bin",
+    "transformer.h.3.self_attention.dense.weight": "pytorch_model_00001-of-00004.bin",
+    "transformer.h.3.self_attention.key_value.weight": "pytorch_model_00001-of-00004.bin",
+    "transformer.h.3.self_attention.query.weight": "pytorch_model_00001-of-00004.bin",
+    "transformer.h.4.input_layernorm.weight": "pytorch_model_00001-of-00004.bin",
+    "transformer.h.4.mlp.down_proj.bias": "pytorch_model_00001-of-00004.bin",
+    "transformer.h.4.mlp.down_proj.weight": "pytorch_model_00001-of-00004.bin",
+    "transformer.h.4.mlp.gate_proj.weight": "pytorch_model_00001-of-00004.bin",
+    "transformer.h.4.mlp.up_proj.weight": "pytorch_model_00001-of-00004.bin",
+    "transformer.h.4.post_attention_layernorm.weight": "pytorch_model_00001-of-00004.bin",
+    "transformer.h.4.self_attention.dense.bias": "pytorch_model_00001-of-00004.bin",
+    "transformer.h.4.self_attention.dense.weight": "pytorch_model_00001-of-00004.bin",
+    "transformer.h.4.self_attention.key_value.weight": "pytorch_model_00001-of-00004.bin",
+    "transformer.h.4.self_attention.query.weight": "pytorch_model_00001-of-00004.bin",
+    "transformer.h.5.input_layernorm.weight": "pytorch_model_00001-of-00004.bin",
+    "transformer.h.5.mlp.down_proj.bias": "pytorch_model_00001-of-00004.bin",
+    "transformer.h.5.mlp.down_proj.weight": "pytorch_model_00001-of-00004.bin",
+    "transformer.h.5.mlp.gate_proj.weight": "pytorch_model_00001-of-00004.bin",
+    "transformer.h.5.mlp.up_proj.weight": "pytorch_model_00001-of-00004.bin",
+    "transformer.h.5.post_attention_layernorm.weight": "pytorch_model_00001-of-00004.bin",
+    "transformer.h.5.self_attention.dense.bias": "pytorch_model_00001-of-00004.bin",
+    "transformer.h.5.self_attention.dense.weight": "pytorch_model_00001-of-00004.bin",
+    "transformer.h.5.self_attention.key_value.weight": "pytorch_model_00001-of-00004.bin",
+    "transformer.h.5.self_attention.query.weight": "pytorch_model_00001-of-00004.bin",
+    "transformer.h.6.input_layernorm.weight": "pytorch_model_00001-of-00004.bin",
+    "transformer.h.6.mlp.down_proj.bias": "pytorch_model_00001-of-00004.bin",
+    "transformer.h.6.mlp.down_proj.weight": "pytorch_model_00001-of-00004.bin",
+    "transformer.h.6.mlp.gate_proj.weight": "pytorch_model_00001-of-00004.bin",
+    "transformer.h.6.mlp.up_proj.weight": "pytorch_model_00001-of-00004.bin",
+    "transformer.h.6.post_attention_layernorm.weight": "pytorch_model_00001-of-00004.bin",
+    "transformer.h.6.self_attention.dense.bias": "pytorch_model_00001-of-00004.bin",
+    "transformer.h.6.self_attention.dense.weight": "pytorch_model_00001-of-00004.bin",
+    "transformer.h.6.self_attention.key_value.weight": "pytorch_model_00001-of-00004.bin",
+    "transformer.h.6.self_attention.query.weight": "pytorch_model_00001-of-00004.bin",
+    "transformer.h.7.input_layernorm.weight": "pytorch_model_00002-of-00004.bin",
+    "transformer.h.7.mlp.down_proj.bias": "pytorch_model_00002-of-00004.bin",
+    "transformer.h.7.mlp.down_proj.weight": "pytorch_model_00002-of-00004.bin",
+    "transformer.h.7.mlp.gate_proj.weight": "pytorch_model_00002-of-00004.bin",
+    "transformer.h.7.mlp.up_proj.weight": "pytorch_model_00002-of-00004.bin",
+    "transformer.h.7.post_attention_layernorm.weight": "pytorch_model_00002-of-00004.bin",
+    "transformer.h.7.self_attention.dense.bias": "pytorch_model_00002-of-00004.bin",
+    "transformer.h.7.self_attention.dense.weight": "pytorch_model_00002-of-00004.bin",
+    "transformer.h.7.self_attention.key_value.weight": "pytorch_model_00002-of-00004.bin",
+    "transformer.h.7.self_attention.query.weight": "pytorch_model_00002-of-00004.bin",
+    "transformer.h.8.input_layernorm.weight": "pytorch_model_00002-of-00004.bin",
+    "transformer.h.8.mlp.down_proj.bias": "pytorch_model_00002-of-00004.bin",
+    "transformer.h.8.mlp.down_proj.weight": "pytorch_model_00002-of-00004.bin",
+    "transformer.h.8.mlp.gate_proj.weight": "pytorch_model_00002-of-00004.bin",
+    "transformer.h.8.mlp.up_proj.weight": "pytorch_model_00002-of-00004.bin",
+    "transformer.h.8.post_attention_layernorm.weight": "pytorch_model_00002-of-00004.bin",
+    "transformer.h.8.self_attention.dense.bias": "pytorch_model_00002-of-00004.bin",
+    "transformer.h.8.self_attention.dense.weight": "pytorch_model_00002-of-00004.bin",
+    "transformer.h.8.self_attention.key_value.weight": "pytorch_model_00002-of-00004.bin",
+    "transformer.h.8.self_attention.query.weight": "pytorch_model_00002-of-00004.bin",
+    "transformer.h.9.input_layernorm.weight": "pytorch_model_00002-of-00004.bin",
+    "transformer.h.9.mlp.down_proj.bias": "pytorch_model_00002-of-00004.bin",
+    "transformer.h.9.mlp.down_proj.weight": "pytorch_model_00002-of-00004.bin",
+    "transformer.h.9.mlp.gate_proj.weight": "pytorch_model_00002-of-00004.bin",
+    "transformer.h.9.mlp.up_proj.weight": "pytorch_model_00002-of-00004.bin",
+    "transformer.h.9.post_attention_layernorm.weight": "pytorch_model_00002-of-00004.bin",
+    "transformer.h.9.self_attention.dense.bias": "pytorch_model_00002-of-00004.bin",
+    "transformer.h.9.self_attention.dense.weight": "pytorch_model_00002-of-00004.bin",
+    "transformer.h.9.self_attention.key_value.weight": "pytorch_model_00002-of-00004.bin",
+    "transformer.h.9.self_attention.query.weight": "pytorch_model_00002-of-00004.bin",
+    "transformer.ln_f.weight": "pytorch_model_00004-of-00004.bin",
+    "transformer.word_embeddings.weight": "pytorch_model_00001-of-00004.bin"
+  }
+}

tokenization_telechat2.py

@@ -0,0 +1,221 @@
+import os
+from shutil import copyfile
+from typing import Any, Dict, List, Optional, Tuple
+import sentencepiece as spm
+from transformers.tokenization_utils import AddedToken, PreTrainedTokenizer
+from transformers.utils import logging
+
+logger = logging.get_logger(__name__)
+
+VOCAB_FILES_NAMES = {"vocab_file": "tokenizer.model"}
+
+# TODO: when we get download url from huggingface, refresh the map
+PRETRAINED_VOCAB_FILES_MAP = {
+    "vocab_file": {},
+    "tokenizer_file": {},
+}
+
+
+class Telechat2Tokenizer(PreTrainedTokenizer):
+    vocab_files_names = VOCAB_FILES_NAMES
+    pretrained_vocab_files_map = PRETRAINED_VOCAB_FILES_MAP
+    model_input_names = ["input_ids", "attention_mask"]
+
+    def __init__(
+            self,
+            vocab_file,
+            unk_token="<unk>",
+            bos_token="<_start>",
+            eos_token="<_end>",
+            pad_token="<_pad>",
+            sp_model_kwargs: Optional[Dict[str, Any]] = None,
+            add_bos_token=True,
+            add_eos_token=False,
+            clean_up_tokenization_spaces=False,
+            **kwargs,
+    ):
+        self.sp_model_kwargs = {} if sp_model_kwargs is None else sp_model_kwargs
+        bos_token = AddedToken(bos_token, lstrip=False, rstrip=False) if isinstance(bos_token, str) else bos_token
+        eos_token = AddedToken(eos_token, lstrip=False, rstrip=False) if isinstance(eos_token, str) else eos_token
+        pad_token = AddedToken(pad_token, lstrip=False, rstrip=False) if isinstance(pad_token, str) else pad_token
+        self.sp_model = spm.SentencePieceProcessor(**self.sp_model_kwargs)
+        self.sp_model.Load(vocab_file)
+        super().__init__(
+            bos_token=bos_token,
+            eos_token=eos_token,
+            pad_token=pad_token,
+            add_bos_token=add_bos_token,
+            add_eos_token=add_eos_token,
+            sp_model_kwargs=self.sp_model_kwargs,
+            clean_up_tokenization_spaces=clean_up_tokenization_spaces,
+            **kwargs,
+        )
+        self.vocab_file = vocab_file
+        self.add_bos_token = add_bos_token
+        self.add_eos_token = add_eos_token
+
+    def __getstate__(self):
+        state = self.__dict__.copy()
+        state["sp_model"] = None
+        return state
+
+    def __setstate__(self, d):
+        self.__dict__ = d
+        self.sp_model = spm.SentencePieceProcessor(**self.sp_model_kwargs)
+        self.sp_model.Load(self.vocab_file)
+
+    @property
+    def vocab_size(self):
+        """Returns vocab size"""
+        return self.sp_model.get_piece_size()
+
+    def get_vocab(self):
+        """Returns vocab as a dict"""
+        vocab = {self.convert_ids_to_tokens(i): i for i in range(self.vocab_size)}
+        vocab.update(self.added_tokens_encoder)
+        return vocab
+
+    @property
+    def vocab(self):
+        return self.get_vocab()
+
+    def _tokenize(self, text):
+        """Returns a tokenized string."""
+        return self.sp_model.encode(text, out_type=str)
+
+    def _convert_token_to_id(self, token):
+        """Converts a token (str) in an id using the vocab."""
+        return self.sp_model.piece_to_id(token)
+
+    def _convert_id_to_token(self, index):
+        """Converts an index (integer) in a token (str) using the vocab."""
+        token = self.sp_model.IdToPiece(index)
+        return token
+
+    def convert_tokens_to_string(self, tokens):
+        """Converts a sequence of tokens (string) in a single string."""
+        current_sub_tokens = []
+        out_string = ""
+        # prev_is_special = False
+        for i, token in enumerate(tokens):
+            # make sure that special tokens are not decoded using sentencepiece model
+            if token in self.all_special_tokens:
+                # if not prev_is_special and i != 0:
+                #     out_string += " "
+                out_string += self.sp_model.decode(current_sub_tokens) + token
+                # prev_is_special = True
+                current_sub_tokens = []
+            else:
+                current_sub_tokens.append(token)
+                # prev_is_special = False
+        out_string += self.sp_model.decode(current_sub_tokens)
+        return out_string
+
+    def save_vocabulary(self, save_directory, filename_prefix: Optional[str] = None) -> Tuple[str]:
+        """
+        Save the vocabulary and special tokens file to a directory.
+
+        Args:
+            save_directory (`str`):
+                The directory in which to save the vocabulary.
+
+        Returns:
+            `Tuple(str)`: Paths to the files saved.
+        """
+        if not os.path.isdir(save_directory):
+            logger.error(f"Vocabulary path ({save_directory}) should be a directory")
+            return
+        out_vocab_file = os.path.join(
+            save_directory, (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["vocab_file"]
+        )
+
+        if os.path.abspath(self.vocab_file) != os.path.abspath(out_vocab_file) and os.path.isfile(self.vocab_file):
+            copyfile(self.vocab_file, out_vocab_file)
+        elif not os.path.isfile(self.vocab_file):
+            with open(out_vocab_file, "wb") as fi:
+                content_spiece_model = self.sp_model.serialized_model_proto()
+                fi.write(content_spiece_model)
+
+        return (out_vocab_file,)
+
+    def build_inputs_with_special_tokens(self, token_ids_0, token_ids_1=None):
+        bos_token_id = [self.bos_token_id] if self.add_bos_token else []
+        eos_token_id = [self.eos_token_id] if self.add_eos_token else []
+
+        output = bos_token_id + token_ids_0 + eos_token_id
+
+        if token_ids_1 is not None:
+            output = output + bos_token_id + token_ids_1 + eos_token_id
+
+        return output
+
+    def get_special_tokens_mask(
+            self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None,
+            already_has_special_tokens: bool = False
+    ) -> List[int]:
+        """
+        Retrieve sequence ids from a token list that has no special tokens added. This method is called when adding
+        special tokens using the tokenizer `prepare_for_model` method.
+
+        Args:
+            token_ids_0 (`List[int]`):
+                List of IDs.
+            token_ids_1 (`List[int]`, *optional*):
+                Optional second list of IDs for sequence pairs.
+            already_has_special_tokens (`bool`, *optional*, defaults to `False`):
+                Whether or not the token list is already formatted with special tokens for the model.
+
+        Returns:
+            `List[int]`: A list of integers in the range [0, 1]: 1 for a special token, 0 for a sequence token.
+        """
+        if already_has_special_tokens:
+            return super().get_special_tokens_mask(
+                token_ids_0=token_ids_0, token_ids_1=token_ids_1, already_has_special_tokens=True
+            )
+
+        bos_token_id = [1] if self.add_bos_token else []
+        eos_token_id = [1] if self.add_eos_token else []
+
+        if token_ids_1 is None:
+            return bos_token_id + ([0] * len(token_ids_0)) + eos_token_id
+        return (
+                bos_token_id
+                + ([0] * len(token_ids_0))
+                + eos_token_id
+                + bos_token_id
+                + ([0] * len(token_ids_1))
+                + eos_token_id
+        )
+
+    def create_token_type_ids_from_sequences(
+            self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None
+    ) -> List[int]:
+        """
+        Creates a mask from the two sequences passed to be used in a sequence-pair classification task. An ALBERT
+        sequence pair mask has the following format:
+
+        ```
+        0 0 0 0 0 0 0 0 0 0 0 1 1 1 1 1 1 1 1 1
+        | first sequence    | second sequence |
+        ```
+
+        if token_ids_1 is None, only returns the first portion of the mask (0s).
+
+        Args:
+            token_ids_0 (`List[int]`):
+                List of ids.
+            token_ids_1 (`List[int]`, *optional*):
+                Optional second list of IDs for sequence pairs.
+
+        Returns:
+            `List[int]`: List of [token type IDs](../glossary#token-type-ids) according to the given sequence(s).
+        """
+        bos_token_id = [self.bos_token_id] if self.add_bos_token else []
+        eos_token_id = [self.eos_token_id] if self.add_eos_token else []
+
+        output = [0] * len(bos_token_id + token_ids_0 + eos_token_id)
+
+        if token_ids_1 is not None:
+            output += [1] * len(bos_token_id + token_ids_1 + eos_token_id)
+
+        return output

tokenizer.model

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a7a5b465bbc9465b214e0962076c1170783a8ee88fb01454b0c33609bd3cf954
+size 2197499

tokenizer_config.json

@@ -0,0 +1,114 @@
+{
+  "tokenizer_class": "Telechat2Tokenizer",
+  "auto_map": {
+    "AutoTokenizer": [
+      "tokenization_telechat2.Telechat2Tokenizer",
+      null
+    ]
+  },
+  "added_tokens_decoder": {
+    "1": {
+      "content": "<_start>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "2": {
+      "content": "<_end>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "3": {
+      "content": "<_pad>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "4": {
+      "content": "<_user>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "5": {
+      "content": "<_bot>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "6": {
+      "content": "<_system>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "9": {
+      "content": "<tool_call>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "10": {
+      "content": "</tool_call>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "11": {
+      "content": "<tool_response>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "12": {
+      "content": "</tool_response>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "additional_special_tokens": [
+    "<_start>",
+    "<_end>",
+    "<_pad>",
+    "<_user>",
+    "<_bot>",
+    "<_system>",
+    "<tool_call>",
+    "</tool_call>",
+    "<tool_response>",
+    "</tool_response>"
+  ],
+  "add_bos_token": false,
+  "add_eos_token": false,
+  "use_fast": false,
+  "clean_up_tokenization_spaces": false,
+  "split_special_tokens": false,
+  "model_max_length": 100000000,
+  "sp_model_kwargs": {},
+  "bos_token": "<_start>",
+  "eos_token": "<_end>",
+  "pad_token": "<_pad>",
+  "chat_template": "{%- if tools %}\n    {%- if messages[0]['role'] == 'system' %}\n        {{-'<_system>'+messages[0]['content'] }}\n    {%- else %}\n        {{- '<_system>'+'你是中国电信星辰语义大模型，英文名是TeleChat，你是由中电信人工智能科技有限公司和中国电信人工智能研究院（TeleAI）研发的人工智能助手。' }}\n    {%- endif %}\n    {{- '\\n\\n# 可用工具\\n你可以调用<tools></tools>标签中包含的一个或多个工具来辅助你回答问题,以下是可用工具详情：\\n<tools>\\n' }}\n    {%- for tool in tools %}\n        {{- tool | tojson }}\n        {{-'\\n'}}\n    {%- endfor %}\n    {{- '</tools>\\n\\n# 调用方法\\n你需要遵循工具的要求，使用json格式返回工具名称及参数，并用<tool_call></tool_call>包含。下方是一个调用模板：\\n<tool_call>\\n{\\\"name\\\": <function-name>, \\\"arguments\\\": <args-json-object>}\\n</tool_call>\\n\\n' }}\n{%- else %}\n    {%- if messages[0]['role'] == 'system' %}\n        {{- '<_system>' + messages[0]['content'] + '\\n' }}\n    {%- else %}\n        {{- '<_system>'+'你是中国电信星辰语义大模型，英文名是TeleChat，你是由中电信人工智能科技有限公司和中国电信人工智能研究院（TeleAI）研发的人工智能助手。\\n' }}\n    {%- endif %}\n{%- endif %}\n{%- for message in messages %}\n    {%- if (message.role == 'user') %}\n        {{- '<_user>' + message.content }}\n    {%- elif message.role == 'bot' or message.role == 'assistant' %}\n        {{- '<_bot>' }}\n        {%- if message.content %}\n            {{- message.content }}\n        {%- endif %}\n        {%- for tool_call in message.tool_calls %}\n            {%- if tool_call.function is defined %}\n                {%- set tool_call = tool_call.function %}\n            {%- endif %}\n            {%- if loop.index0 == 0 %}\n                {{-'<tool_call>'}}\n            {%- else %}\n                {{-'\\n<tool_call>'}}\n            {%- endif %}\n            {{- '\\n{\"name\": \"' }}{{ tool_call.name }}\n            {{- '\", \"arguments\": ' }}\n            {{- tool_call.arguments | tojson }}\n            {{- '}\\n</tool_call>' }}\n        {%- endfor %}\n        {{- '<_end>\\n' }}\n    {%- elif message.role == 'tool' %}\n        {%- if (loop.index0 == 0) or (messages[loop.index0 - 1].role != 'tool') %}\n            {{- '<_user>'+'<tool_response>\\n' }}\n        {%- else %}\n            {{- '\\n<tool_response>\\n' }}\n        {%- endif %}\n        {{- message.content }}\n        {{- '\\n</tool_response>' }}\n    {%- endif %}\n{%- endfor %}\n{%- if add_generation_prompt %}\n    {{- '<_bot>' }}\n{%- endif %}"
+}