Update for huggingface hub

demo.py

@@ -0,0 +1,20 @@
+from lyraChatGLM import LyraChatGLM6B
+
+model_path = "./models/1-gpu-fp16.h5"
+tokenizer_path = "./models"
+data_type = "fp16"
+int8_mode = 0
+max_output_length = 150
+arch = "Ampere" # Ampere or Volta
+
+model = LyraChatGLM6B(model_path, tokenizer_path, data_type, int8_mode, arch)
+prompt = "今天天气大概 25度，有点小雨，吹着风，我想去户外散步，应该穿什么样的衣服裤子鞋子搭配。"
+test_batch_size = 256
+
+prompts = [prompt, ]
+
+
+# If you want to get different output in same batch, you can set do_sample to True
+output_texts = model.generate(prompts, output_length=max_output_length,top_k=30, top_p=0.85, temperature=0.35, repetition_penalty=1.2, do_sample=False)
+
+print(output_texts)

lyraChatGLM/__init__.py

@@ -0,0 +1 @@
+from .lyra_glm import LyraChatGLM6B

lyraChatGLM/config.py

@@ -0,0 +1,31 @@
+import dataclasses
+from typing import Optional
+
+
+@dataclasses.dataclass
+class ChatGLM6BParam:
+    num_heads: int = 32
+    size_per_head: int = 128
+    inter_size: int = 16384
+    num_layers: int = 28
+    vocab_size: int = 130528
+    start_id: Optional[int] = 130004
+    end_id: Optional[int] = 130005
+    tensor_para_size: int = 1
+    pipeline_para_size: int = 1
+    remove_padding: bool = True
+    shared_contexts_ratio: float = 1.0
+    layernorm_eps: float = 1e-5
+    weights_data_type: str = "fp16"
+
+    def __post_init__(self):
+        if not 0.0 <= self.shared_contexts_ratio <= 1.0:
+            raise ValueError(
+                f'Got an invalid value of shared_context_ratio '
+                f'{self.shared_contexts_ratio} - range: [0.0, 1.0]')
+
+    def asdict(self):
+        return dataclasses.asdict(self)
+
+
+CHATGLM_6B_PARAM = ChatGLM6BParam()

lyraChatGLM/lyra_glm.py

@@ -0,0 +1,174 @@
+from __future__ import annotations
+
+import configparser
+import pathlib
+import typing
+
+import torch
+import transformers
+
+from .config import CHATGLM_6B_PARAM
+from .model import ChatGLM6BModel
+
+
+class LyraChatGLM6B:
+    def __init__(self, model_path, tokenizer_path=None, dtype='fp16', int8_mode=0, arch="Ampere") -> None:
+        self.model_path = model_path
+        self.tokenizer_path = tokenizer_path
+        self.dtype = dtype
+        self.arch=arch
+        if dtype != 'int8':
+            int8_mode = 0
+        self.int8_mode = int8_mode
+
+        self.model, self.tokenizer = self.load_model_and_tokenizer()
+        if not (arch in ["Ampere", "Volta"]):
+            raise ValueError("Only support GPU device Ampere(A100,A10) or Volta(V100)")
+
+        print("Got model and tokenizer")
+
+    def load_model_and_tokenizer(self):
+        if self.tokenizer_path is None:
+            tokenizer_path = self.model_path
+        else:
+            tokenizer_path = self.tokenizer_path
+
+        print(f'Loading tokenizer from {pathlib.Path(tokenizer_path).parent}')
+        tokenizer = transformers.AutoTokenizer.from_pretrained(tokenizer_path, trust_remote_code=True)
+
+        checkpoint_path = pathlib.Path(self.model_path)
+
+        config_path = checkpoint_path.parent / 'config.ini'
+
+        if config_path.exists():
+            # Read model params from config.
+            cfg = configparser.ConfigParser()
+            cfg.read(config_path)
+            model_name = 'glm6b'
+            inference_data_type = self.dtype
+            if inference_data_type == None:
+                inference_data_type = cfg.get(model_name, "weight_data_type")
+            model_args = dict(
+                head_num=cfg.getint(model_name, 'head_num'),
+                size_per_head=cfg.getint(model_name, "size_per_head"),
+                layer_num=cfg.getint(model_name, "num_layer"),
+                tensor_para_size=cfg.getint(model_name, "tensor_para_size"),
+                vocab_size=cfg.getint(model_name, "vocab_size"),
+                start_id=cfg.getint(model_name, "start_id"),
+                end_id=cfg.getint(model_name, "end_id"),
+                weights_data_type=cfg.get(model_name, "weight_data_type"),
+                layernorm_eps=cfg.getfloat(model_name, 'layernorm_eps'),
+                inference_data_type=inference_data_type)
+        else:
+            inference_data_type = self.dtype
+            if inference_data_type == None:
+                inference_data_type = CHATGLM_6B_PARAM.weights_data_type
+            model_args = dict(head_num=CHATGLM_6B_PARAM.num_heads,
+                              size_per_head=CHATGLM_6B_PARAM.size_per_head,
+                              vocab_size=CHATGLM_6B_PARAM.vocab_size,
+                              start_id=CHATGLM_6B_PARAM.start_id or tokenizer.bos_token_id,
+                              end_id=CHATGLM_6B_PARAM.end_id or tokenizer.eos_token_id,
+                              layer_num=CHATGLM_6B_PARAM.num_layers,
+                              tensor_para_size=CHATGLM_6B_PARAM.tensor_para_size,
+                              weights_data_type=CHATGLM_6B_PARAM.weights_data_type,
+                              layernorm_eps=CHATGLM_6B_PARAM.layernorm_eps,
+                              inference_data_type=inference_data_type,
+                              )
+
+        # update common parameters
+        model_args.update(dict(
+            rotary_embedding_dim=64,
+            max_seq_len=0,  # for position seq embedding
+            pipeline_para_size=CHATGLM_6B_PARAM.pipeline_para_size,
+            shared_contexts_ratio=CHATGLM_6B_PARAM.shared_contexts_ratio,
+            int8_mode=self.int8_mode
+        ))
+
+        print('[INFO] Load Our Highly Optimized LyraChatGLM6B model')
+        for k, v in model_args.items():
+            print(f' - {k.ljust(25, ".")}: {v}')
+
+        # Check sanity and consistency between the model and tokenizer.
+        checklist = ['head_num', 'size_per_head', 'vocab_size', 'layer_num',
+                     'tensor_para_size', 'tensor_para_size', 'weights_data_type']
+        if None in [model_args[k] for k in checklist]:
+            none_params = [p for p in checklist if model_args[p] is None]
+            print(f'[WARNING] Found None parameters {none_params}. They must '
+                  f'be provided either by config file or CLI arguments.')
+        if model_args['start_id'] != tokenizer.bos_token_id:
+            print('[WARNING] Given start_id is not matched with the bos token '
+                  'id of the pretrained tokenizer.')
+        if model_args['end_id'] not in (tokenizer.pad_token_id, tokenizer.eos_token_id):
+            print('[WARNING] Given end_id is not matched with neither pad '
+                  'token id nor eos token id of the pretrained tokenizer.')
+
+        print(f'Loading tokenizer from {self.model_path}')
+        model = ChatGLM6BModel(arch=self.arch,**model_args)
+        if not model.load(ckpt_path=self.model_path):
+            print('[WARNING] Skip model loading since no checkpoints are found')
+
+        return model, tokenizer
+
+    def generate(self, prompts: typing.List[str] | str,
+                 output_length: int = 512,
+                 beam_width: int = 1,
+                 top_k: typing.Optional[torch.IntTensor] = 1,
+                 top_p: typing.Optional[torch.FloatTensor] = 1.0,
+                 beam_search_diversity_rate: typing.Optional[torch.FloatTensor] = 0.0,
+                 temperature: typing.Optional[torch.FloatTensor] = 1.0,
+                 len_penalty: typing.Optional[torch.FloatTensor] = 0.0,
+                 repetition_penalty: typing.Optional[torch.FloatTensor] = 1.0,
+                 presence_penalty: typing.Optional[torch.FloatTensor] = None,
+                 min_length: typing.Optional[torch.IntTensor] = None,
+                 bad_words_list: typing.Optional[torch.IntTensor] = None,
+                 do_sample: bool = False,
+                 return_output_length: bool = False,
+                 return_cum_log_probs: int = 0):
+        #
+        if isinstance(prompts, str):
+            prompts = [prompts, ]
+
+        inputs = prompts
+
+        batch_size = len(inputs)
+        ones_int = torch.ones(size=[batch_size], dtype=torch.int32)
+        ones_float = torch.ones(size=[batch_size], dtype=torch.float32)
+
+        input_token_ids = self.tokenizer(prompts, return_tensors="pt", padding=True).input_ids.int()
+        input_lengths = torch.IntTensor([len(ids) for ids in input_token_ids])
+        mask_positions = torch.IntTensor([seq.index(130001) for seq in input_token_ids.tolist()])
+
+        random_seed = None
+        if do_sample:
+            random_seed = torch.randint(0, 262144, (batch_size,), dtype=torch.long)
+
+        outputs = self.model(start_ids=input_token_ids,
+                             start_lengths=input_lengths,
+                             mask_positions=mask_positions,
+                             output_len=output_length,
+                             beam_width=beam_width,
+                             top_k=top_k*ones_int,
+                             top_p=top_p*ones_float,
+                             beam_search_diversity_rate=beam_search_diversity_rate*ones_float,
+                             temperature=temperature*ones_float,
+                             len_penalty=len_penalty*ones_float,
+                             repetition_penalty=repetition_penalty*ones_float,
+                             presence_penalty=presence_penalty,
+                             min_length=min_length,
+                             random_seed=random_seed,
+                             bad_words_list=bad_words_list,
+                             return_output_length=return_output_length,
+                             return_cum_log_probs=return_cum_log_probs)
+
+        if return_cum_log_probs > 0:
+            outputs = outputs[0]  # output_token_ids.
+
+        # Slice the generated token ids of the 1st beam result.
+        # output = input tokens + generated tokens.
+        output_token_ids = [out[0, length:].cpu()
+                            for out, length in zip(outputs, input_lengths)]
+
+        output_texts = self.tokenizer.batch_decode(
+            output_token_ids, skip_special_tokens=False)
+
+        return output_texts

lyraChatGLM/model.py

@@ -0,0 +1,625 @@
+import os
+import h5py
+import pathlib
+import typing
+
+import numpy as np
+import torch
+import torch.distributed as dist
+import torch.nn as nn
+
+str_type_map = {"fp32": torch.float32, "fp16": torch.float16, "bf16": torch.bfloat16}
+
+
+class ChatGLM6BWeights:
+    def __init__(
+            self, head_num, size_per_head, layer_num, vocab_size, max_seq_len, tensor_para_size, pipeline_para_size,
+            weights_data_type: typing.Union[str, np.dtype],
+            inference_data_type: str, has_adapters: bool = False, adapter_inter_size: int = 0, gpt_with_moe: bool = False,
+            has_positional_encoding: bool = False, has_pre_decoder_layernorm: bool = False,
+            has_post_decoder_layernorm: bool = True, int8_mode: int = 0, inter_size: int = 0):
+        assert(head_num % tensor_para_size == 0)
+        if int8_mode == 1:
+            torch_infer_dtype = str_type_map[inference_data_type]
+            assert torch_infer_dtype == torch.float16 or torch_infer_dtype == torch.bfloat16, "Weight only quant only supported for infer type fp16 or bf16."
+            quant = torch.ops.fastertransformer.symmetric_quantize_last_axis_of_batched_matrix
+            self.weight_transpose_calibrate_quantize = lambda x: quant(x, torch.int8)
+        else:
+            assert int8_mode == 0, "Invalid int8 mode for GPT. Must be 0 or 1"
+
+        self.head_num = head_num
+        self.size_per_head = size_per_head
+        self.layer_num = layer_num
+        self.vocab_size = vocab_size
+        self.max_seq_len = max_seq_len
+        self.tensor_para_size = tensor_para_size
+        self.pipeline_para_size = pipeline_para_size
+        self.layers_per_device = layer_num // pipeline_para_size
+
+        self.has_adapters = has_adapters
+        self.adapter_inter_size = adapter_inter_size
+        self.gpt_with_moe = gpt_with_moe
+        self.has_positional_encoding = has_positional_encoding
+        self.has_pre_decoder_layernorm = has_pre_decoder_layernorm
+        self.has_post_decoder_layernorm = has_post_decoder_layernorm
+
+        local_head_num = head_num // tensor_para_size
+        global_head_num = head_num
+        local_hidden_units = local_head_num * size_per_head
+        global_hidden_units = global_head_num * size_per_head
+        local_inter_size = local_hidden_units * 4
+        if inter_size != 0:
+            assert inter_size % tensor_para_size == 0, f"inter_size({inter_size}) \% tensor_para_size({tensor_para_size}) must be 0"
+            local_inter_size = inter_size // tensor_para_size
+        local_adapter_inter_size = self.adapter_inter_size // tensor_para_size
+
+        self.local_head_num = local_head_num
+        self.global_head_num = global_head_num
+        self.local_hidden_units = local_hidden_units
+        self.global_hidden_units = global_hidden_units
+        self.local_inter_size = local_inter_size
+
+        self.int8_mode = int8_mode
+        self.share_embed = False
+
+        if isinstance(weights_data_type, str):
+            try:
+                weights_data_type = {
+                    "fp16": np.float16,
+                    "fp32": np.float32,
+                    "float16": np.float16,
+                    "float32": np.float32,
+                }[weights_data_type]
+            except KeyError:
+                raise ValueError(f"Don't know how to interpret weights_data_type: {weights_data_type}")
+
+        assert weights_data_type in [np.float32, np.float16]
+        self.weights_data_type = weights_data_type
+        self.inference_data_type = inference_data_type
+
+        self.w = []
+        self.int8_w = []
+        self.scale = []
+
+        # Transformer blocks
+        self.w.extend([torch.zeros(global_hidden_units, dtype=str_type_map[self.inference_data_type])]
+                      * layer_num)   # self_layernorm_gamma
+        self.w.extend([torch.zeros(global_hidden_units, dtype=str_type_map[self.inference_data_type])]
+                      * layer_num)   # self_layernorm_beta
+        self.w.extend([torch.zeros(global_hidden_units, local_hidden_units * 3,
+                      dtype=str_type_map[self.inference_data_type])] * layer_num)   # self_kernel
+        self.w.extend([torch.zeros(local_hidden_units * 3, dtype=str_type_map[self.inference_data_type])]
+                      * layer_num)   # self_bias
+        self.w.extend(
+            [torch.zeros(local_hidden_units, global_hidden_units, dtype=str_type_map[self.inference_data_type])] *
+            layer_num)  # self_output_kernel
+        self.w.extend([torch.zeros(global_hidden_units, dtype=str_type_map[self.inference_data_type])]
+                      * layer_num)   # self_output_bias
+        self.w.extend([torch.zeros(global_hidden_units, dtype=str_type_map[self.inference_data_type])]
+                      * layer_num)   # ffn_layernorm_gamma
+        self.w.extend([torch.zeros(global_hidden_units, dtype=str_type_map[self.inference_data_type])]
+                      * layer_num)   # ffn_layernorm_beta
+        self.w.extend(
+            [torch.zeros(global_hidden_units, local_inter_size, dtype=str_type_map[self.inference_data_type])] *
+            layer_num)  # ffn_kernel1
+        self.w.extend([torch.zeros(local_inter_size, dtype=str_type_map[self.inference_data_type])]
+                      * layer_num)   # ffn_bias1
+        self.w.extend(
+            [torch.zeros(local_inter_size, global_hidden_units, dtype=str_type_map[self.inference_data_type])] *
+            layer_num)  # ffn_kernel2
+        self.w.extend([torch.zeros(global_hidden_units, dtype=str_type_map[self.inference_data_type])]
+                      * layer_num)   # ffn_bias2
+
+        optional_adapter_offset = 0
+
+        # After Transformer blocks
+        if self.has_pre_decoder_layernorm:
+            self.w.append(torch.zeros(global_hidden_units, dtype=str_type_map[
+                self.inference_data_type]))   # embedding layernorm gamma
+            self.w.append(torch.zeros(global_hidden_units, dtype=str_type_map[
+                self.inference_data_type]))   # embedding layernorm beta
+            optional_adapter_offset += 2
+        if self.has_post_decoder_layernorm:
+            self.w.append(torch.zeros(global_hidden_units, dtype=str_type_map[
+                self.inference_data_type]))   # final layernorm gamma
+            self.w.append(torch.zeros(global_hidden_units, dtype=str_type_map[
+                self.inference_data_type]))   # final layernorm beta
+            optional_adapter_offset += 2
+        if self.has_positional_encoding:
+            self.w.append(torch.zeros(max_seq_len, global_hidden_units, dtype=str_type_map[
+                self.inference_data_type]))   # position_encoding_table
+            optional_adapter_offset += 1
+
+        self.pre_embed_idx = len(self.w)
+        self.w.append(torch.zeros(vocab_size, global_hidden_units,
+                      dtype=str_type_map[self.inference_data_type]))   # embedding_table
+        self.post_embed_idx = len(self.w)
+        self.w.append(torch.zeros(vocab_size, global_hidden_units, dtype=str_type_map[
+            self.inference_data_type]))   # post embedding_kernel
+        self.adapter_offset = 2 + optional_adapter_offset
+
+        self.w.extend([torch.empty(0, dtype=str_type_map[self.inference_data_type])] * layer_num)   # gating_weight
+        self.adapter_offset += layer_num
+
+        # adapters
+        if self.has_adapters:
+            self.w.extend([torch.zeros(global_hidden_units, local_adapter_inter_size,
+                          dtype=str_type_map[self.inference_data_type])] * layer_num)   # adaptor1_kernel1
+            self.w.extend([torch.zeros(local_adapter_inter_size, dtype=str_type_map[
+                self.inference_data_type])] * layer_num)   # adaptor1_bias1
+            self.w.extend([torch.zeros(local_adapter_inter_size, global_hidden_units,
+                          dtype=str_type_map[self.inference_data_type])] * layer_num)   # adaptor1_kernel2
+            self.w.extend([torch.zeros(global_hidden_units, dtype=str_type_map[
+                self.inference_data_type])] * layer_num)   # adaptor1_bias2
+            self.w.extend([torch.zeros(global_hidden_units, local_adapter_inter_size,
+                          dtype=str_type_map[self.inference_data_type])] * layer_num)   # adaptor2_kernel1
+            self.w.extend([torch.zeros(local_adapter_inter_size, dtype=str_type_map[
+                self.inference_data_type])] * layer_num)   # adaptor2_bias1
+            self.w.extend([torch.zeros(local_adapter_inter_size, global_hidden_units,
+                          dtype=str_type_map[self.inference_data_type])] * layer_num)   # adaptor2_kernel2
+            self.w.extend([torch.zeros(global_hidden_units, dtype=str_type_map[
+                self.inference_data_type])] * layer_num)   # adaptor2_bias2
+
+        # Initialization
+        self._map(lambda w: torch.nn.init.normal_(w, mean=0., std=1.))
+
+        if (self.int8_mode != 0):
+            self.int8_w.extend([torch.zeros(global_hidden_units, local_hidden_units *
+                               3, dtype=torch.int8)] * layer_num)   # self_int8_kernel
+            self.scale.extend([torch.zeros(local_hidden_units * 3, dtype=torch.float)] * layer_num)   # self_scale
+            self.int8_w.extend([torch.zeros(local_hidden_units, global_hidden_units, dtype=torch.int8)]
+                               * layer_num)   # self_output_int8_kernel
+            self.scale.extend([torch.zeros(global_hidden_units, dtype=torch.float)] * layer_num)   # self_output_scale
+            self.int8_w.extend([torch.zeros(global_hidden_units, local_inter_size,
+                               dtype=torch.int8)] * layer_num)   # ffn_int8_kernel1
+            self.scale.extend([torch.zeros(local_inter_size, dtype=torch.float)] * layer_num)   # ffn_scale1
+            self.int8_w.extend([torch.zeros(local_inter_size, global_hidden_units,
+                               dtype=torch.int8)] * layer_num)   # ffn_int8_kernel2
+            self.scale.extend([torch.zeros(global_hidden_units, dtype=torch.float)] * layer_num)   # ffn_scale2
+
+            if self.has_adapters:
+                self.int8_w.extend([torch.zeros(global_hidden_units, local_adapter_inter_size,
+                                   dtype=torch.int8)] * layer_num)   # adaptor1_int8_kernel1
+                self.scale.extend([torch.zeros(local_adapter_inter_size, dtype=torch.float)]
+                                  * layer_num)   # adaptor1_scale1
+                self.int8_w.extend([torch.zeros(local_adapter_inter_size, global_hidden_units,
+                                   dtype=torch.int8)] * layer_num)   # adaptor1_int8_kernel2
+                self.scale.extend([torch.zeros(global_hidden_units, dtype=torch.float)] * layer_num)   # adaptor1_scale2
+                self.int8_w.extend([torch.zeros(global_hidden_units, local_adapter_inter_size,
+                                   dtype=torch.int8)] * layer_num)   # adaptor2_int8_kernel1
+                self.scale.extend([torch.zeros(local_adapter_inter_size, dtype=torch.float)]
+                                  * layer_num)   # adaptor2_scale1
+                self.int8_w.extend([torch.zeros(local_adapter_inter_size, global_hidden_units,
+                                   dtype=torch.int8)] * layer_num)   # adaptor2_int8_kernel2
+                self.scale.extend([torch.zeros(global_hidden_units, dtype=torch.float)] * layer_num)   # adaptor2_scale2
+
+    def __getitem__(self, idx):
+        return self.w[idx]
+
+    def __setitem__(self, idx, val):
+        self.w[idx] = val
+
+    def __len__(self):
+        return len(self.w)
+
+    def _map(self, func):
+        assert(self.pre_embed_idx < self.post_embed_idx,
+               "Pre decoder embedding index should be lower than post decoder embedding index.")
+        for i in range(len(self.w)):
+            if isinstance(self.w[i], list):
+                for j in range(len(self.w[i])):
+                    self.w[i][j] = func(self.w[i][j])
+            else:
+                if self.share_embed and i == self.post_embed_idx:
+                    # If sharing the pre and post embedding, any mapping to
+                    # the pre decoder weight will give the same output to the
+                    # post decoder weight, so we just copy here.
+                    self.w[self.post_embed_idx] = self.w[self.pre_embed_idx]
+                else:
+                    self.w[i] = func(self.w[i])
+
+    def _map_int8(self, func):
+        for i in range(len(self.int8_w)):
+            if isinstance(self.int8_w[i], list):
+                for j in range(len(self.int8_w[i])):
+                    self.int8_w[i][j] = func(self.int8_w[i][j])
+
+            else:
+                self.int8_w[i] = func(self.int8_w[i])
+        for i in range(len(self.scale)):
+            if isinstance(self.scale[i], list):
+                for j in range(len(self.scale[i])):
+                    self.scale[i][j] = func(self.scale[i][j])
+
+            else:
+                self.scale[i] = func(self.scale[i])
+
+    def _map_int8_scales(self, func):
+        for i in range(len(self.scale)):
+            if isinstance(self.scale[i], list):
+                for j in range(len(self.scale[i])):
+                    self.scale[i][j] = func(self.scale[i][j])
+
+            else:
+                self.scale[i] = func(self.scale[i])
+
+    def load(self, ckpt_path, tp_rank, pipeline_para_rank):
+        if not os.path.exists(ckpt_path):
+            raise FileNotFoundError(f"Failed to find {ckpt_path}")
+        w = []
+
+        type_map = {np.float32: torch.float32, np.float16: torch.float16}
+        # Load
+
+        def is_load(i): return i >= self.layers_per_device * \
+            pipeline_para_rank and i < self.layers_per_device * (pipeline_para_rank + 1)
+
+        h5f = h5py.File(ckpt_path, "r")
+
+        def load_to_torch(key, is_load: bool):
+            if is_load:
+                npdata = h5f[key]["weights"][:]
+                return torch.from_numpy(npdata).to(str_type_map[self.inference_data_type])
+            else:
+                return torch.empty(0).to(str_type_map[self.inference_data_type])
+        w.extend([load_to_torch(f"model.layers.{i}.input_layernorm.weight", is_load(i))
+                 for i in range(self.layer_num)])
+        w.extend([load_to_torch(f"model.layers.{i}.input_layernorm.bias", is_load(i))
+                 for i in range(self.layer_num)])
+        w.extend(
+            [load_to_torch(
+                f"model.layers.{i}.attention.query_key_value.weight.{tp_rank}", is_load(i))
+             for i in range(self.layer_num)])
+        w.extend([
+            load_to_torch(
+                f"model.layers.{i}.attention.query_key_value.bias.{tp_rank}", is_load(i))
+            for i in range(self.layer_num)])
+        w.extend([load_to_torch(f"model.layers.{i}.attention.dense.weight.{tp_rank}",
+                 is_load(i)) for i in range(self.layer_num)])
+        w.extend([load_to_torch(f"model.layers.{i}.attention.dense.bias", is_load(i))
+                 for i in range(self.layer_num)])
+        w.extend([load_to_torch(f"model.layers.{i}.post_attention_layernorm.weight",
+                 is_load(i)) for i in range(self.layer_num)])
+        w.extend([load_to_torch(f"model.layers.{i}.post_attention_layernorm.bias",
+                 is_load(i)) for i in range(self.layer_num)])
+        w.extend(
+            [load_to_torch(f"model.layers.{i}.mlp.dense_h_to_4h.weight.{tp_rank}", is_load(i))
+             for i in range(self.layer_num)])
+        w.extend(
+            [load_to_torch(f"model.layers.{i}.mlp.dense_h_to_4h.bias.{tp_rank}", is_load(i))
+             for i in range(self.layer_num)])
+        w.extend(
+            [load_to_torch(f"model.layers.{i}.mlp.dense_4h_to_h.weight.{tp_rank}", is_load(i))
+             for i in range(self.layer_num)])
+        w.extend([load_to_torch(f"model.layers.{i}.mlp.dense_4h_to_h.bias", is_load(i)) for i in range(self.layer_num)])
+
+        if self.has_pre_decoder_layernorm:
+            w.append(load_to_torch(f"model.pre_decoder_layernorm.weight", True))
+            w.append(load_to_torch(f"model.pre_decoder_layernorm.bias", True))
+
+        if self.has_post_decoder_layernorm:
+            w.append(load_to_torch(f"model.final_layernorm.weight", True))
+            w.append(load_to_torch(f"model.final_layernorm.bias", True))
+
+        if self.has_positional_encoding:
+            wpe = load_to_torch(f"model.wpe", True).reshape(-1, self.global_hidden_units)
+            assert self.max_seq_len <= wpe.size(0), (
+                f"max_seq_len ({self.max_seq_len} must not exceed "
+                f"the value of maximum sequence length during training ({wpe.size(0)})."
+            )
+            w.append(wpe)
+        w.append(load_to_torch(f"model.wte", True))
+        self.share_embed = True
+        w.append(torch.empty(0).to(str_type_map[self.inference_data_type]))
+
+        gate_list = []
+        for i in range(self.layer_num):
+            gate_list.append(load_to_torch(f"model.layers.{i}.mlp.moe.gate.wg.weight", False))
+        w.extend(gate_list)
+
+        if self.has_adapters:
+            w.extend(
+                [load_to_torch(
+                    f"model.layers.{i}.after_attention_adapter.dense_h_to_4h.weight.{tp_rank}", is_load(i))
+                 for i in range(self.layer_num)])
+            w.extend([
+                load_to_torch(
+                    f"model.layers.{i}.after_attention_adapter.dense_h_to_4h.bias.{tp_rank}", is_load(i))
+                for i in range(self.layer_num)])
+            w.extend(
+                [load_to_torch(
+                    f"model.layers.{i}.after_attention_adapter.dense_4h_to_h.weight.{tp_rank}", is_load(i))
+                 for i in range(self.layer_num)])
+            w.extend(
+                [load_to_torch(f"model.layers.{i}.after_attention_adapter.dense_4h_to_h.bias", is_load(i))
+                 for i in range(self.layer_num)])
+            w.extend(
+                [load_to_torch(f"model.layers.{i}.after_ffn_adapter.dense_h_to_4h.weight.{tp_rank}", is_load(i))
+                 for i in range(self.layer_num)])
+            w.extend(
+                [load_to_torch(f"model.layers.{i}.after_ffn_adapter.dense_h_to_4h.bias.{tp_rank}", is_load(i))
+                 for i in range(self.layer_num)])
+            w.extend(
+                [load_to_torch(f"model.layers.{i}.after_ffn_adapter.dense_4h_to_h.weight.{tp_rank}", is_load(i))
+                 for i in range(self.layer_num)])
+            w.extend([load_to_torch(
+                f"model.layers.{i}.after_ffn_adapter.dense_4h_to_h.bias", is_load(i)) for i in range(self.layer_num)])
+
+        assert len(self.w) == len(w)
+
+        # Reshape
+        try:
+            for i in range(len(w)):
+                if w[i].nelement() == self.w[i].nelement():
+                    self.w[i] = w[i].reshape(self.w[i].shape)
+                else:
+                    self.w[i] = w[i]
+
+        except RuntimeError:
+            raise RuntimeError(
+                f"head_num, size_per_head, vocab_size, and max_seq_len must be the same as the ones during training "
+                f"(idx: {i} expected shape: {self.w[i].shape} got shape: {w[i].shape})."
+            )
+
+        # transpose calibrate quantize the kernel
+        layer_num = self.layer_num
+        if self.int8_mode != 0:
+            for i in range(layer_num):
+                self.int8_w[i + 0 * layer_num], self.scale[i + 0 *
+                                                           layer_num] = self.weight_transpose_calibrate_quantize(self.w[2 * layer_num + i])
+                self.int8_w[i + 1 * layer_num], self.scale[i + 1 *
+                                                           layer_num] = self.weight_transpose_calibrate_quantize(self.w[4 * layer_num + i])
+                self.int8_w[i + 2 * layer_num], self.scale[i + 2 *
+                                                           layer_num] = self.weight_transpose_calibrate_quantize(self.w[8 * layer_num + i])
+                self.int8_w[i + 3 * layer_num], self.scale[i + 3 *
+                                                           layer_num] = self.weight_transpose_calibrate_quantize(self.w[10 * layer_num + i])
+
+                # We clear the original weights since they are no longer needed
+                if self.int8_mode == 1:
+                    self.w[2 * layer_num + i] = torch.empty(0).to(str_type_map[self.inference_data_type])
+                    self.w[4 * layer_num + i] = torch.empty(0).to(str_type_map[self.inference_data_type])
+                    self.w[8 * layer_num + i] = torch.empty(0).to(str_type_map[self.inference_data_type])
+                    self.w[10 * layer_num + i] = torch.empty(0).to(str_type_map[self.inference_data_type])
+
+                if self.has_adapters:
+                    self.int8_w[i + 4 * layer_num], self.scale[i + 4 * layer_num] = self.weight_transpose_calibrate_quantize(
+                        self.w[12 * layer_num + i + self.adapter_offset])
+                    self.int8_w[i + 5 * layer_num], self.scale[i + 5 * layer_num] = self.weight_transpose_calibrate_quantize(
+                        self.w[14 * layer_num + i + self.adapter_offset])
+                    self.int8_w[i + 6 * layer_num], self.scale[i + 6 * layer_num] = self.weight_transpose_calibrate_quantize(
+                        self.w[16 * layer_num + i + self.adapter_offset])
+                    self.int8_w[i + 7 * layer_num], self.scale[i + 7 * layer_num] = self.weight_transpose_calibrate_quantize(
+                        self.w[18 * layer_num + i + self.adapter_offset])
+
+                    # Similar to above:
+                    if self.int8_mode == 1:
+                        self.w[12 * layer_num + i + self.adapter_offset] = torch.empty(
+                            0).to(str_type_map[self.inference_data_type])
+                        self.w[14 * layer_num + i + self.adapter_offset] = torch.empty(
+                            0).to(str_type_map[self.inference_data_type])
+                        self.w[16 * layer_num + i + self.adapter_offset] = torch.empty(
+                            0).to(str_type_map[self.inference_data_type])
+                        self.w[18 * layer_num + i + self.adapter_offset] = torch.empty(
+                            0).to(str_type_map[self.inference_data_type])
+        return True
+
+
+class ChatGLM6BModel(nn.Module):
+    def __init__(self,
+                 head_num, size_per_head,
+                 vocab_size,
+                 rotary_embedding_dim,
+                 start_id, end_id, layer_num,
+                 arch,
+                 max_seq_len: int,
+                 tensor_para_size: int,
+                 pipeline_para_size: int,
+                 inference_data_type: str,
+                 inter_size: int = 0,
+                 # glm_variant_params
+                 layernorm_eps: float = 1e-5,
+                 layernorm_type: typing.Literal['pre_layernorm', 'post_layernorm'] = "pre_layernorm",
+                 activation_type: str = "Gelu",
+                 gpt_with_moe: bool = False,
+                 expert_num: int = 0,
+                 moe_k: int = 0,
+                 moe_layer_index: typing.List = [],
+                 has_positional_encoding: bool = False,
+                 has_pre_decoder_layernorm: bool = False,
+                 has_post_decoder_layernorm: bool = True,
+                 has_adapters: bool = False,
+                 adapter_inter_size: int = 0,
+                 use_attention_linear_bias: bool = False,
+                 int8_mode: int = 0,
+                 weights_data_type: typing.Union[str, np.dtype] = np.float32,
+                 shared_contexts_ratio: float = 1.0):
+        super().__init__()
+        self.head_num = head_num
+        self.size_per_head = size_per_head
+        self.vocab_size = vocab_size
+        self.rotary_embedding_dim = rotary_embedding_dim
+        self.start_id = start_id
+        self.end_id = end_id
+        self.layer_num = layer_num
+        self.inter_size = inter_size if inter_size != 0 else 4 * self.head_num * self.size_per_head
+        self.arch = arch
+        # gpt_variant_params
+        self.layernorm_eps = layernorm_eps
+        self.layernorm_type = layernorm_type
+        self.activation_type = activation_type
+        self.gpt_with_moe = gpt_with_moe
+        self.expert_num = expert_num
+        self.moe_k = moe_k
+        self.moe_layer_index = moe_layer_index
+        self.has_positional_encoding = has_positional_encoding
+        self.has_pre_decoder_layernorm = has_pre_decoder_layernorm
+        self.has_post_decoder_layernorm = has_post_decoder_layernorm
+        self.has_adapters = has_adapters
+        self.adapter_inter_size = adapter_inter_size
+        self.use_attention_linear_bias = use_attention_linear_bias
+
+        # multi-gpu params
+        self.tensor_para_size = tensor_para_size
+        self.pipeline_para_size = pipeline_para_size
+        self.use_sparse_gemm = False
+        self.build_model = False
+        self.int8_mode = int8_mode
+        self.weights_data_type = weights_data_type
+        self.shared_contexts_ratio = shared_contexts_ratio
+
+        assert torch.cuda.is_available(), "CUDA is required for this model."
+
+        assert head_num % tensor_para_size == 0, "head_num must be a multiple of tensor_para_size."
+        assert layer_num % pipeline_para_size == 0, "layer_num must be a multiple of pipeline_para_size."
+
+        # Load the C++ model into Pytorch model.
+        if arch == "Ampere":
+            lib_path = pathlib.Path(__file__).parent / "ftlib" / "libth_transformer_sm80.so"
+        elif arch == "Volta":
+            lib_path = pathlib.Path(__file__).parent / "ftlib" / "libth_transformer_sm70.so"
+        torch.classes.load_library(os.path.abspath(lib_path))
+
+        # Prepare weights
+        self.weights = ChatGLM6BWeights(head_num, size_per_head, layer_num, vocab_size,
+                                        max_seq_len, tensor_para_size, pipeline_para_size,
+                                        weights_data_type=weights_data_type,
+                                        inference_data_type=inference_data_type,
+                                        gpt_with_moe=self.gpt_with_moe,
+                                        has_positional_encoding=self.has_positional_encoding,
+                                        has_pre_decoder_layernorm=self.has_pre_decoder_layernorm,
+                                        has_post_decoder_layernorm=self.has_post_decoder_layernorm,
+                                        has_adapters=self.has_adapters,
+                                        adapter_inter_size=self.adapter_inter_size,
+                                        int8_mode=int8_mode,
+                                        inter_size=inter_size)
+
+        # Prepare for tensor/pipeline parallel
+        try:
+            dist.init_process_group(backend='mpi')
+        except:
+            print("[INFO] WARNING: Have initialized the process group")
+        self.rank = dist.get_rank()
+        self.device_count = torch.cuda.device_count()
+        self.device = self.rank % self.device_count
+        torch.cuda.set_device(self.device)
+
+        world_size = dist.get_world_size()
+        assert world_size == tensor_para_size * pipeline_para_size, "tensor_para_size * pipeline_para_size must be equal to world_size."
+
+        self.tensor_para_rank = self.rank % self.tensor_para_size
+        self.pipeline_para_rank = self.rank // self.tensor_para_size
+
+    def load(self, ckpt_path):
+        is_load = self.weights.load(ckpt_path, tp_rank=self.tensor_para_rank,
+                                    pipeline_para_rank=self.pipeline_para_rank)
+        self.cuda()
+        torch.cuda.empty_cache()  # clean cache for model weight preprocessing
+        return is_load
+
+    def sparse(self):
+        if not self.use_sparse_gemm:
+            self.use_sparse_gemm = True
+
+    def cuda(self):
+        self.weights._map(lambda w: w.cuda(self.device))
+        if self.int8_mode != 0:
+            self.weights._map_int8(lambda w: w.cuda(self.device))
+
+        if self.build_model:
+            del self.model
+            self.build_model = False
+
+        self.model = torch.classes.FasterTransformer.GlmOp(
+            self.head_num, self.size_per_head, self.inter_size,
+            self.layer_num,
+            self.expert_num,
+            self.moe_k,
+            self.moe_layer_index,
+            self.vocab_size,
+            self.rotary_embedding_dim,
+            self.start_id, self.end_id,
+            self.tensor_para_size, self.pipeline_para_size, self.int8_mode,
+            # GLM variant parameters
+            self.layernorm_eps,
+            self.layernorm_type,
+            self.activation_type,
+            self.has_positional_encoding,
+            self.has_pre_decoder_layernorm,
+            self.has_post_decoder_layernorm,
+            self.has_adapters,
+            self.adapter_inter_size,
+            self.use_attention_linear_bias,
+            self.weights.w,
+            self.weights.int8_w,
+            self.weights.scale,
+            self.shared_contexts_ratio)
+        self.build_model = True
+
+    def forward(self,
+                start_ids: torch.IntTensor,
+                start_lengths: torch.IntTensor,
+                mask_positions: torch.IntTensor,
+                output_len: int,
+                beam_width: int = 1,
+                top_k: typing.Optional[torch.IntTensor] = None,
+                top_p: typing.Optional[torch.FloatTensor] = None,
+                beam_search_diversity_rate: typing.Optional[torch.FloatTensor] = None,
+                temperature: typing.Optional[torch.FloatTensor] = None,
+                len_penalty: typing.Optional[torch.FloatTensor] = None,
+                repetition_penalty: typing.Optional[torch.FloatTensor] = None,
+                presence_penalty: typing.Optional[torch.FloatTensor] = None,
+                min_length: typing.Optional[torch.IntTensor] = None,
+                random_seed: typing.Optional[torch.LongTensor] = None,
+                bad_words_list: typing.Optional[torch.IntTensor] = None,
+                return_output_length: bool = False,
+                return_cum_log_probs: int = 0):
+        if not self.build_model:
+            # for the cases we don't load model
+            self.cuda()
+            torch.cuda.empty_cache()  # clean cache for model weight preprocessing
+        input_len = start_ids.size(1)
+        assert input_len > 0, "input len must be larger than zero. For an unconditional case, use start_id as the first token."
+
+        # Inputs to device
+        start_ids = start_ids.cuda(self.device)
+        start_lengths = start_lengths.cuda(self.device)
+        mask_positions = mask_positions.cuda(self.device)
+
+        # outputs: output_ids, output_lengths, output_cum_log_probs (optional)
+        outputs = self.model.forward(start_ids,
+                                     start_lengths,
+                                     mask_positions,
+                                     output_len,
+                                     beam_width,  # optional, can be None
+                                     top_k,  # optional, can be None
+                                     top_p,  # optional, can be None
+                                     beam_search_diversity_rate,  # optional, can be None
+                                     temperature,  # optional, can be None
+                                     len_penalty,  # optional, can be None
+                                     repetition_penalty,  # optional, can be None
+                                     presence_penalty,  # optional, can be None
+                                     min_length,  # optional, can be None
+                                     random_seed,  # optional, can be None
+                                     bad_words_list,  # optional, can be None
+                                     return_cum_log_probs)  # optional, can be None
+        if return_cum_log_probs == 0:
+            output_ids, output_lengths = outputs
+        else:
+            output_ids, output_lengths, output_cum_log_probs = outputs
+        if return_output_length:
+            if return_cum_log_probs > 0:
+                return output_ids, output_lengths, output_cum_log_probs
+            else:
+                return output_ids, output_lengths
+        else:
+            return output_ids
+
+    def set_input_tensor(self, input_tensor):
+        """Set input tensor to be used instead of forward()'s input.
+
+        When doing pipeline parallelism the input from the previous
+        stage comes from communication, not from the input, so the
+        model's forward_step_func won't have it. This function is thus
+        used by internal code to bypass the input provided by the
+        forward_step_func"""
+        self.input_tensor = input_tensor

models/config.ini

@@ -0,0 +1,13 @@
+[glm6b]
+model_name = chatglm-6b
+head_num = 32
+size_per_head = 128
+inter_size = 16384
+max_pos_seq_len = 2048
+num_layer = 28
+vocab_size = 130528
+start_id = 130004
+end_id = 130005
+weight_data_type = fp16
+tensor_para_size = 1
+layernorm_eps = 1e-5

models/tokenization_chatglm.py

@@ -0,0 +1,443 @@
+"""Tokenization classes for ChatGLM."""
+from typing import List, Optional, Union
+import os
+
+from transformers.tokenization_utils import PreTrainedTokenizer
+from transformers.utils import logging, PaddingStrategy
+from transformers.tokenization_utils_base import EncodedInput, BatchEncoding
+from typing import Dict
+import sentencepiece as spm
+import numpy as np
+
+logger = logging.get_logger(__name__)
+
+PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES = {
+    "THUDM/chatglm-6b": 2048,
+}
+
+
+class TextTokenizer:
+    def __init__(self, model_path):
+        self.sp = spm.SentencePieceProcessor()
+        self.sp.Load(model_path)
+        self.num_tokens = self.sp.vocab_size()
+
+    def encode(self, text):
+        return self.sp.EncodeAsIds(text)
+
+    def decode(self, ids: List[int]):
+        return self.sp.DecodeIds(ids)
+
+    def tokenize(self, text):
+        return self.sp.EncodeAsPieces(text)
+
+    def convert_tokens_to_string(self, tokens):
+        return self.sp.DecodePieces(tokens)
+
+    def convert_tokens_to_ids(self, tokens):
+        return [self.sp.PieceToId(token) for token in tokens]
+
+    def convert_token_to_id(self, token):
+        return self.sp.PieceToId(token)
+
+    def convert_id_to_token(self, idx):
+        return self.sp.IdToPiece(idx)
+
+    def __len__(self):
+        return self.num_tokens
+
+
+class SPTokenizer:
+    def __init__(
+            self,
+            vocab_file,
+            num_image_tokens=20000,
+            max_blank_length=80,
+            byte_fallback=True,
+    ):
+        assert vocab_file is not None
+        self.vocab_file = vocab_file
+        self.num_image_tokens = num_image_tokens
+        self.special_tokens = ["[MASK]", "[gMASK]", "[sMASK]", "<unused_0>", "<sop>", "<eop>", "<ENC>", "<dBLOCK>"]
+        self.max_blank_length = max_blank_length
+        self.byte_fallback = byte_fallback
+        self.text_tokenizer = TextTokenizer(vocab_file)
+
+    def _get_text_tokenizer(self):
+        return self.text_tokenizer
+
+    @staticmethod
+    def get_blank_token(length: int):
+        assert length >= 2
+        return f"<|blank_{length}|>"
+
+    @staticmethod
+    def get_tab_token():
+        return f"<|tab|>"
+
+    @property
+    def num_text_tokens(self):
+        return self.text_tokenizer.num_tokens
+
+    @property
+    def num_tokens(self):
+        return self.num_image_tokens + self.num_text_tokens
+
+    @staticmethod
+    def _encode_whitespaces(text: str, max_len: int = 80):
+        text = text.replace("\t", SPTokenizer.get_tab_token())
+        for i in range(max_len, 1, -1):
+            text = text.replace(" " * i, SPTokenizer.get_blank_token(i))
+        return text
+
+    def _preprocess(self, text: str, linebreak=True, whitespaces=True):
+        if linebreak:
+            text = text.replace("\n", "<n>")
+        if whitespaces:
+            text = self._encode_whitespaces(text, max_len=self.max_blank_length)
+        return text
+
+    def encode(
+            self, text: str, linebreak=True, whitespaces=True, add_dummy_prefix=True
+    ) -> List[int]:
+        """
+        @param text: Text to encode.
+        @param linebreak: Whether to encode newline (\n) in text.
+        @param whitespaces: Whether to encode multiple whitespaces or tab in text, useful for source code encoding.
+        @param special_tokens: Whether to encode special token ([MASK], [gMASK], etc.) in text.
+        @param add_dummy_prefix: Whether to add dummy blank space in the beginning.
+        """
+        text = self._preprocess(text, linebreak, whitespaces)
+        if not add_dummy_prefix:
+            text = "<n>" + text
+        tmp = self._get_text_tokenizer().encode(text)
+        tokens = [x + self.num_image_tokens for x in tmp]
+        return tokens if add_dummy_prefix else tokens[2:]
+
+    def postprocess(self, text):
+        text = text.replace("<n>", "\n")
+        text = text.replace(SPTokenizer.get_tab_token(), "\t")
+        for i in range(2, self.max_blank_length + 1):
+            text = text.replace(self.get_blank_token(i), " " * i)
+        return text
+
+    def decode(self, text_ids: List[int]) -> str:
+        ids = [int(_id) - self.num_image_tokens for _id in text_ids]
+        ids = [_id for _id in ids if _id >= 0]
+        text = self._get_text_tokenizer().decode(ids)
+        text = self.postprocess(text)
+        return text
+
+    def decode_tokens(self, tokens: List[str]) -> str:
+        text = self._get_text_tokenizer().convert_tokens_to_string(tokens)
+        text = self.postprocess(text)
+        return text
+
+    def tokenize(
+            self, text: str, linebreak=True, whitespaces=True, add_dummy_prefix=True
+    ) -> List[str]:
+        """
+        @param text: Text to encode.
+        @param linebreak: Whether to encode newline (\n) in text.
+        @param whitespaces: Whether to encode multiple whitespaces or tab in text, useful for source code encoding.
+        @param special_tokens: Whether to encode special token ([MASK], [gMASK], etc.) in text.
+        @param add_dummy_prefix: Whether to add dummy blank space in the beginning.
+        """
+        text = self._preprocess(text, linebreak, whitespaces)
+        if not add_dummy_prefix:
+            text = "<n>" + text
+        tokens = self._get_text_tokenizer().tokenize(text)
+        return tokens if add_dummy_prefix else tokens[2:]
+
+    def __getitem__(self, x: Union[int, str]):
+        if isinstance(x, int):
+            if x < self.num_image_tokens:
+                return "<image_{}>".format(x)
+            else:
+                return self.text_tokenizer.convert_id_to_token(x - self.num_image_tokens)
+        elif isinstance(x, str):
+            if x.startswith("<image_") and x.endswith(">") and x[7:-1].isdigit():
+                return int(x[7:-1])
+            else:
+                return self.text_tokenizer.convert_token_to_id(x) + self.num_image_tokens
+        else:
+            raise ValueError("The key should be str or int.")
+
+
+class ChatGLMTokenizer(PreTrainedTokenizer):
+    """
+    Construct a ChatGLM tokenizer. Based on byte-level Byte-Pair-Encoding.
+
+    Args:
+        vocab_file (`str`):
+            Path to the vocabulary file.
+    """
+
+    vocab_files_names = {"vocab_file": "ice_text.model"}
+    max_model_input_sizes = PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
+    model_input_names = ["input_ids", "attention_mask", "position_ids"]
+
+    def __init__(
+            self,
+            vocab_file,
+            do_lower_case=False,
+            remove_space=False,
+            bos_token='<sop>',
+            eos_token='<eop>',
+            end_token='</s>',
+            mask_token='[MASK]',
+            gmask_token='[gMASK]',
+            padding_side="left",
+            pad_token="<pad>",
+            unk_token="<unk>",
+            num_image_tokens=20000,
+            **kwargs
+    ) -> None:
+        super().__init__(
+            do_lower_case=do_lower_case,
+            remove_space=remove_space,
+            padding_side=padding_side,
+            bos_token=bos_token,
+            eos_token=eos_token,
+            end_token=end_token,
+            mask_token=mask_token,
+            gmask_token=gmask_token,
+            pad_token=pad_token,
+            unk_token=unk_token,
+            num_image_tokens=num_image_tokens,
+            **kwargs
+        )
+
+        self.do_lower_case = do_lower_case
+        self.remove_space = remove_space
+        self.vocab_file = vocab_file
+
+        self.bos_token = bos_token
+        self.eos_token = eos_token
+        self.end_token = end_token
+        self.mask_token = mask_token
+        self.gmask_token = gmask_token
+
+        self.sp_tokenizer = SPTokenizer(vocab_file, num_image_tokens=num_image_tokens)
+
+        """ Initialisation """
+
+    @property
+    def gmask_token_id(self) -> Optional[int]:
+        if self.gmask_token is None:
+            return None
+        return self.convert_tokens_to_ids(self.gmask_token)
+
+    @property
+    def end_token_id(self) -> Optional[int]:
+        """
+        `Optional[int]`: Id of the end of context token in the vocabulary. Returns `None` if the token has not been
+        set.
+        """
+        if self.end_token is None:
+            return None
+        return self.convert_tokens_to_ids(self.end_token)
+
+    @property
+    def vocab_size(self):
+        """ Returns vocab size """
+        return self.sp_tokenizer.num_tokens
+
+    def get_vocab(self):
+        """ Returns vocab as a dict """
+        vocab = {self._convert_id_to_token(i): i for i in range(self.vocab_size)}
+        vocab.update(self.added_tokens_encoder)
+        return vocab
+
+    def preprocess_text(self, inputs):
+        if self.remove_space:
+            outputs = " ".join(inputs.strip().split())
+        else:
+            outputs = inputs
+
+        if self.do_lower_case:
+            outputs = outputs.lower()
+
+        return outputs
+
+    def _tokenize(self, text, **kwargs):
+        """ Returns a tokenized string. """
+        text = self.preprocess_text(text)
+
+        seq = self.sp_tokenizer.tokenize(text)
+
+        return seq
+
+    def convert_tokens_to_string(self, tokens: List[str]) -> str:
+        return self.sp_tokenizer.decode_tokens(tokens)
+
+    def _decode(
+            self,
+            token_ids: Union[int, List[int]],
+            **kwargs
+    ) -> str:
+        if isinstance(token_ids, int):
+            token_ids = [token_ids]
+        if len(token_ids) == 0:
+            return ""
+        if self.pad_token_id in token_ids:  # remove pad
+            token_ids = list(filter((self.pad_token_id).__ne__, token_ids))
+        return super()._decode(token_ids, **kwargs)
+
+    def _convert_token_to_id(self, token):
+        """ Converts a token (str) in an id using the vocab. """
+        return self.sp_tokenizer[token]
+
+    def _convert_id_to_token(self, index):
+        """Converts an index (integer) in a token (str) using the vocab."""
+        return self.sp_tokenizer[index]
+
+    def save_vocabulary(self, save_directory, filename_prefix=None):
+        """
+        Save the vocabulary and special tokens file to a directory.
+
+        Args:
+            save_directory (`str`):
+                The directory in which to save the vocabulary.
+            filename_prefix (`str`, *optional*):
+                An optional prefix to add to the named of the saved files.
+
+        Returns:
+            `Tuple(str)`: Paths to the files saved.
+        """
+        if os.path.isdir(save_directory):
+            vocab_file = os.path.join(
+                save_directory, self.vocab_files_names["vocab_file"]
+            )
+        else:
+            vocab_file = save_directory
+
+        with open(self.vocab_file, 'rb') as fin:
+            proto_str = fin.read()
+
+        with open(vocab_file, "wb") as writer:
+            writer.write(proto_str)
+
+        return (vocab_file,)
+
+    def build_inputs_with_special_tokens(
+            self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None
+    ) -> List[int]:
+        """
+        Build model inputs from a sequence or a pair of sequence for sequence classification tasks by concatenating and
+        adding special tokens. A BERT sequence has the following format:
+
+        - single sequence: `[CLS] X [SEP]`
+        - pair of sequences: `[CLS] A [SEP] B [SEP]`
+
+        Args:
+            token_ids_0 (`List[int]`):
+                List of IDs to which the special tokens will be added.
+            token_ids_1 (`List[int]`, *optional*):
+                Optional second list of IDs for sequence pairs.
+
+        Returns:
+            `List[int]`: List of [input IDs](../glossary#input-ids) with the appropriate special tokens.
+        """
+        gmask_id = self.sp_tokenizer[self.gmask_token]
+        eos_id = self.sp_tokenizer[self.eos_token]
+        token_ids_0 = token_ids_0 + [gmask_id, self.sp_tokenizer[self.bos_token]]
+        if token_ids_1 is not None:
+            token_ids_0 = token_ids_0 + token_ids_1 + [eos_id]
+        return token_ids_0
+
+    def _pad(
+            self,
+            encoded_inputs: Union[Dict[str, EncodedInput], BatchEncoding],
+            max_length: Optional[int] = None,
+            padding_strategy: PaddingStrategy = PaddingStrategy.DO_NOT_PAD,
+            pad_to_multiple_of: Optional[int] = None,
+            return_attention_mask: Optional[bool] = None,
+    ) -> dict:
+        """
+        Pad encoded inputs (on left/right and up to predefined length or max length in the batch)
+
+        Args:
+            encoded_inputs:
+                Dictionary of tokenized inputs (`List[int]`) or batch of tokenized inputs (`List[List[int]]`).
+            max_length: maximum length of the returned list and optionally padding length (see below).
+                Will truncate by taking into account the special tokens.
+            padding_strategy: PaddingStrategy to use for padding.
+
+                - PaddingStrategy.LONGEST Pad to the longest sequence in the batch
+                - PaddingStrategy.MAX_LENGTH: Pad to the max length (default)
+                - PaddingStrategy.DO_NOT_PAD: Do not pad
+                The tokenizer padding sides are defined in self.padding_side:
+
+                    - 'left': pads on the left of the sequences
+                    - 'right': pads on the right of the sequences
+            pad_to_multiple_of: (optional) Integer if set will pad the sequence to a multiple of the provided value.
+                This is especially useful to enable the use of Tensor Core on NVIDIA hardware with compute capability
+                `>= 7.5` (Volta).
+            return_attention_mask:
+                (optional) Set to False to avoid returning attention mask (default: set to model specifics)
+        """
+        # Load from model defaults
+        bos_token_id = self.sp_tokenizer[self.bos_token]
+        mask_token_id = self.sp_tokenizer[self.mask_token]
+        gmask_token_id = self.sp_tokenizer[self.gmask_token]
+        assert self.padding_side == "left"
+
+        required_input = encoded_inputs[self.model_input_names[0]]
+        seq_length = len(required_input)
+
+        if padding_strategy == PaddingStrategy.LONGEST:
+            max_length = len(required_input)
+
+        if max_length is not None and pad_to_multiple_of is not None and (max_length % pad_to_multiple_of != 0):
+            max_length = ((max_length // pad_to_multiple_of) + 1) * pad_to_multiple_of
+
+        needs_to_be_padded = padding_strategy != PaddingStrategy.DO_NOT_PAD and len(required_input) != max_length
+
+        # Initialize attention mask if not present.
+        if max_length is not None:
+            if "attention_mask" not in encoded_inputs:
+                if bos_token_id in required_input:
+                    context_length = required_input.index(bos_token_id)
+                else:
+                    context_length = seq_length
+                attention_mask = np.ones((1, seq_length, seq_length))
+                attention_mask = np.tril(attention_mask)
+                attention_mask[:, :, :context_length] = 1
+                attention_mask = np.bool_(attention_mask < 0.5)
+                encoded_inputs["attention_mask"] = attention_mask
+
+            if "position_ids" not in encoded_inputs:
+                if bos_token_id in required_input:
+                    context_length = required_input.index(bos_token_id)
+                else:
+                    context_length = seq_length
+                position_ids = np.arange(seq_length, dtype=np.int64)
+                mask_token = mask_token_id if mask_token_id in required_input else gmask_token_id
+                if mask_token in required_input:
+                    mask_position = required_input.index(mask_token)
+                    position_ids[context_length:] = mask_position
+                block_position_ids = np.concatenate(
+                    [np.zeros(context_length, dtype=np.int64),
+                     np.arange(1, seq_length - context_length + 1, dtype=np.int64)])
+                encoded_inputs["position_ids"] = np.stack([position_ids, block_position_ids], axis=0)
+
+        if needs_to_be_padded:
+            difference = max_length - len(required_input)
+
+            if "attention_mask" in encoded_inputs:
+                encoded_inputs["attention_mask"] = np.pad(encoded_inputs["attention_mask"],
+                                                          pad_width=[(0, 0), (difference, 0), (difference, 0)],
+                                                          mode='constant', constant_values=True)
+            if "token_type_ids" in encoded_inputs:
+                encoded_inputs["token_type_ids"] = [self.pad_token_type_id] * difference + encoded_inputs[
+                    "token_type_ids"
+                ]
+            if "special_tokens_mask" in encoded_inputs:
+                encoded_inputs["special_tokens_mask"] = [1] * difference + encoded_inputs["special_tokens_mask"]
+            if "position_ids" in encoded_inputs:
+                encoded_inputs["position_ids"] = np.pad(encoded_inputs["position_ids"],
+                                                        pad_width=[(0, 0), (difference, 0)])
+            encoded_inputs[self.model_input_names[0]] = [self.pad_token_id] * difference + required_input
+
+        return encoded_inputs

models/tokenizer_config.json

@@ -0,0 +1,20 @@
+{
+  "name_or_path": "THUDM/chatglm-6b",
+  "bos_token": "<sop>",
+  "eos_token": "<eop>",
+  "end_token": "</s>",
+  "gmask_token": "[gMASK]",
+  "mask_token": "[MASK]",
+  "pad_token": "<pad>",
+  "unk_token": "<unk>",
+  "remove_space": false,
+  "do_lower_case": false,
+  "tokenizer_class": "ChatGLMTokenizer",
+  "num_image_tokens": 0,
+  "auto_map": {
+    "AutoTokenizer": [
+      "tokenization_chatglm.ChatGLMTokenizer",
+      null
+      ]
+  }
+}

requirements.txt

@@ -0,0 +1,8 @@
+icetk
+cpm_kernels
+transformers
+huggingface_hub
+numpy
+setuptools
+torch
+protobuf==3.20.3