NewAcceleration (#19)

- remove files (2b20dbd049a0be820cd120a1447d6499b2d34d8a)
- Update for huggingface hub (b061fc933c657f766d8322eff24fc1f8ff06ea8d)

CHANGELOG.md

@@ -1,3 +0,0 @@
-## v1.0
-
-- Add accelerated ChatGLM-6B model (from: https://huggingface.co/THUDM/chatglm-6b)

Dockerfile

@@ -1,11 +0,0 @@
-FROM nvcr.io/nvidia/pytorch:23.02-py3
-
-WORKDIR /workdir
-
-COPY requirements.txt /workdir/
-
-# since installing icetk will install protobuf 3.18.3, and we need protobuf==3.20.3
-RUN pip install -r requirements.txt && \
-    pip install protobuf==3.20.3 
-
-

README.md

@@ -1,120 +0,0 @@
----
-license: creativeml-openrail-m
-language:
-- en
-tags:
-- LLM
-- tensorRT
-- ChatGLM
----
-## Model Card for lyraChatGLM
-
-lyraChatGLM is currently the **fastest ChatGLM-6B** available. To the best of our knowledge, it is the **first accelerated version of ChatGLM-6B**.
-
-The inference speed of lyraChatGLM has achieved **10x** acceleration upon the ealry original version. We are still working hard to further improve the performance.
-
-Among its main features are:
-
-- weights: original ChatGLM-6B weights released by THUDM.
-- device: lyraChatGLM is mainly based on TensorRT compiled for SM=80 (A100, for example).
-- batch_size: compiled with dynamic batch size, max batch_size = 8
-
-## Speed
-
-### test environment
-
-- device: Nvidia A100 40G
-- batch size: 8
-
-**Since early chatGLM version didn't suport batch inference, `original` in below table was measured on batch_size=1**
-
-
-**According to [this discussion](https://huggingface.co/TMElyralab/lyraChatGLM/discussions/6), this bug has been fixed and the speed on batch_size=8 reachs up to 137 tokens/s. We will evaluate and update the latest performance.**
-
-|version|speed|
-|:-:|:-:|
-|original|30 tokens/s|
-|lyraChatGLM|310 tokens/s|
-
-
-## Model Sources
-
-- **Repository:** https://huggingface.co/THUDM/chatglm-6b
-
-## Try Demo in 2 fast steps
-
-``` bash
-#step 1
-git clone https://huggingface.co/TMElyralab/lyraChatGLM
-cd lyraChatGLM
-
-#step 2
-docker run --gpus=1 --rm --net=host -v ${PWD}:/workdir yibolu96/lyra-chatglm-env:0.0.1 python3 /workdir/demo.py
-```
-
-## Uses
-
-```python
-from transformers import AutoTokenizer
-from lyraChatGLM import GLM6B, FasterChatGLM
-import os
-
-current_workdir = os.path.dirname(__file__)
-
-MAX_OUT_LEN = 100
-chatglm6b_dir = os.path.join(current_workdir, "models")
-tokenizer = AutoTokenizer.from_pretrained(chatglm6b_dir, trust_remote_code=True)
-input_str = ["为什么我们需要对深度学习模型加速？", ]
-inputs = tokenizer(input_str, return_tensors="pt", padding=True)
-input_ids = inputs.input_ids.to('cuda:0')
-
-plan_path = os.path.join(current_workdir, "models/glm6b-bs8.ftm")
-
-# kernel for chat model.
-kernel = GLM6B(plan_path=plan_path,
-               batch_size=1,
-               num_beams=1,
-               use_cache=True,
-               num_heads=32,
-               emb_size_per_heads=128,
-               decoder_layers=28,
-               vocab_size=150528,
-               max_seq_len=MAX_OUT_LEN)
-
-chat = FasterChatGLM(model_dir=chatglm6b_dir, kernel=kernel).half().cuda()
-
-# generate
-sample_output = chat.generate(inputs=input_ids, max_length=MAX_OUT_LEN)
-# de-tokenize model output to text
-res = tokenizer.decode(sample_output[0], skip_special_tokens=True)
-print(res)
-```
-## Demo output
-
-### input
-为什么我们需要对深度学习模型加速? 。
-
-### output
-为什么我们需要对深度学习模型加速? 深度学习模型的训练需要大量计算资源,特别是在训练模型时,需要大量的内存、GPU(图形处理器)和其他计算资源。因此,训练深度学习模型需要一定的时间,并且如果模型不能快速训练,则可能会导致训练进度缓慢或无法训练。
-
-以下是一些原因我们需要对深度学习模型加速:
-
-1. 训练深度神经网络需要大量的计算资源,特别是在训练深度神经网络时,需要更多的计算资源,因此需要更快的训练速度。
-
-### TODO：
-
-We plan to implement a FasterTransformer version to publish a much faster release. Stay tuned!
-
-## Citation
-``` bibtex
-@Misc{lyraChatGLM2023,
-  author =       {Kangjian Wu, Zhengtao Wang, Yibo Lu, Bin Wu},
-  title =        {lyraChatGLM: Accelerating ChatGLM by 10x+},
-  howpublished = {\url{https://huggingface.co/TMElyralab/lyraChatGLM}},
-  year =         {2023}
-}
-```
-
-## Report bug
-- start a discussion to report any bugs!--> https://huggingface.co/TMElyralab/lyraChatGLM/discussions
-- report bug with a `[bug]` mark in the title.

demo.py

@@ -1,35 +1,20 @@
-# coding=utf-8
+from lyraChatGLM import LyraChatGLM6B
 
-from transformers import AutoTokenizer
-from lyraChatGLM import GLM6B, FasterChatGLM
-import os
+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
 
-current_workdir = os.path.dirname(__file__)
+model = LyraChatGLM6B(model_path, tokenizer_path, data_type, int8_mode, arch)
+prompt = "今天天气大概 25度，有点小雨，吹着风，我想去户外散步，应该穿什么样的衣服裤子鞋子搭配。"
+test_batch_size = 256
 
-MAX_OUT_LEN = 100
-chatglm6b_dir = os.path.join(current_workdir, "models")
-tokenizer = AutoTokenizer.from_pretrained(chatglm6b_dir, trust_remote_code=True)
-input_str = ["为什么我们需要对深度学习模型加速？", ]
-inputs = tokenizer(input_str, return_tensors="pt", padding=True)
-input_ids = inputs.input_ids.to('cuda:0')
+prompts = [prompt, ]
 
-plan_path = os.path.join(current_workdir, "models/glm6b-bs8.ftm")
 
-# kernel for chat model.
-kernel = GLM6B(plan_path=plan_path,
-               batch_size=1,
-               num_beams=1,
-               use_cache=True,
-               num_heads=32,
-               emb_size_per_heads=128,
-               decoder_layers=28,
-               vocab_size=150528,
-               max_seq_len=MAX_OUT_LEN)
+# 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)
 
-chat = FasterChatGLM(model_dir=chatglm6b_dir, kernel=kernel).half().cuda()
-
-# generate
-sample_output = chat.generate(inputs=input_ids, max_length=MAX_OUT_LEN)
-# de-tokenize model output to text
-res = tokenizer.decode(sample_output[0], skip_special_tokens=True)
-print(res)
+print(output_texts)

lyraChatGLM/__init__.py

@@ -1,10 +1 @@
-import os
-import ctypes
-
-current_workdir = os.path.dirname(__file__)
-ctypes.cdll.LoadLibrary(os.path.join(current_workdir, "libnvinfer_plugin.so"))
-os.environ["TORCH_USE_RTLD_GLOBAL"]="YES"
-
-import torch
-from .glm import GLM6B
-from .model import FasterChatGLM
+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

@@ -1,131 +1,625 @@
+import os
+import h5py
+import pathlib
+import typing
+
+import numpy as np
 import torch
-from transformers.modeling_outputs import CausalLMOutputWithPast
-from transformers.modeling_utils import PreTrainedModel
-from transformers import AutoConfig
-from typing import Dict, List, Tuple, Union, Optional
-
-
-class FasterChatGLM(PreTrainedModel):
-    def __init__(self, model_dir, kernel, *inputs, **kwargs):
-        config = AutoConfig.from_pretrained(model_dir, trust_remote_code=True)
-        config.n_head = config.num_attention_heads
-        config.n_embd = config.hidden_size
-        config.n_layer = config.num_layers
-        super().__init__(config, *inputs, **kwargs)
-        self.kernel = kernel
-        self.fake_reg = torch.nn.Linear(2, 2)
-        self.position_encoding_2d = True
-
-    def forward(self, input_ids, position_ids, attention_mask, past_key_values, *args, **kwargs):
-        inputs_values = [input_ids, position_ids, attention_mask]
-        if past_key_values is not None:
-            inputs_values = inputs_values + past_key_values
-
-        computed = self.kernel.infer(inputs_values)
-        logits = computed[0]
-        if len(computed) == 1:
-            present_key_values = None
-        else:
-            present_key_values = computed[1:]
-
-        return CausalLMOutputWithPast(logits=logits, past_key_values=present_key_values)
-
-    def get_masks_and_position_ids(self, seq, mask_position, context_length, device, gmask=False):
-        attention_mask = torch.ones((1, context_length, context_length), device=device)
-        attention_mask.tril_()
-        attention_mask[..., :context_length - 1] = 1
-        attention_mask.unsqueeze_(1)
-        attention_mask = (attention_mask < 0.5).bool()
-
-        if self.position_encoding_2d:
-            seq_length = seq.index(150004)
-            position_ids = torch.arange(context_length, dtype=torch.long, device=device)
-            if not gmask:
-                position_ids[seq_length:] = mask_position
-            block_position_ids = torch.cat((
-                torch.zeros(seq_length, dtype=torch.long, device=device),
-                torch.arange(context_length - seq_length, dtype=torch.long, device=device) + 1
-            ))
-            position_ids = torch.stack((position_ids, block_position_ids), dim=0)
+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:
-            position_ids = torch.arange(context_length, dtype=torch.long, device=device)
-            if not gmask:
-                position_ids[context_length - 1:] = mask_position
+            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
 
-        position_ids = position_ids.unsqueeze(0)
+        self.w.extend([torch.empty(0, dtype=str_type_map[self.inference_data_type])] * layer_num)   # gating_weight
+        self.adapter_offset += layer_num
 
-        return attention_mask, position_ids
+        # 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
 
-    def prepare_one_sample(self, input_id, mask_token, past, past_key_values, use_gmask):
+        # Initialization
+        self._map(lambda w: torch.nn.init.normal_(w, mean=0., std=1.))
 
-        seq = input_id.tolist()
-        mask_position = seq.index(mask_token)
+        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 mask_token not in seq:
-            raise ValueError("You have to add either [MASK] or [gMASK] in your input")
+            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
 
-        # only last token for input_ids if past is not None
-        if past is not None or past_key_values is not None:
-            context_length = seq.index(150004)
-            last_token = input_id[-1].unsqueeze(-1).unsqueeze(0)  # 2 dim
-            proc_input_id = last_token
-            if self.position_encoding_2d:
-                position_ids = torch.tensor([[[mask_position], [len(seq) - context_length]]], dtype=torch.long,
-                                            device=input_id.device)
+    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:
-                position_ids = torch.tensor([[mask_position]], dtype=torch.long, device=input_id.device)
+                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])
 
-            attention_mask = torch.zeros(1, 1, 1, 1, device=input_id.device)
-        else:
-            proc_input_id = input_id.unsqueeze(0)
-            attention_mask, position_ids = self.get_masks_and_position_ids(
-                seq=seq,
-                mask_position=mask_position,
-                context_length=len(seq),
-                device=input_id.device,
-                gmask=use_gmask
+    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})."
             )
 
-        return (proc_input_id.to(torch.int32), position_ids.to(torch.int32),
-                attention_mask.to(torch.bool))
-
-    def prepare_inputs_for_generation(
-            self,
-            input_ids: torch.LongTensor,
-            past: Optional[torch.Tensor] = None,
-            past_key_values: Optional[torch.Tensor] = None,
-            attention_mask: Optional[torch.Tensor] = None,
-            use_cache: bool = None,
-            **kwargs
-    ) -> dict:
-
-        MASK, gMASK = 150000, 150001
-        mask_token = MASK if MASK in input_ids else gMASK
-        use_gmask = False if MASK in input_ids else gMASK
-
-        batch_input_ids, batch_position_ids, batch_attention_mask = [], [], []
-        for input_id in input_ids:
-            proc_input_id, position_id, attention_mask = self.prepare_one_sample(
-                input_id, mask_token, past, past_key_values, use_gmask)
-            batch_input_ids.append(proc_input_id)
-            batch_position_ids.append(position_id)
-            batch_attention_mask.append(attention_mask)
-
-        batch_input_ids = torch.vstack(batch_input_ids)
-        batch_position_ids = torch.vstack(batch_position_ids)
-        batch_attention_mask = torch.vstack(batch_attention_mask)
-
-        if past is None:
-            past = past_key_values
-
-        if past is not None or past_key_values is not None:
-            self.kernel.set_context_mode(False)
+        # 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:
-            self.kernel.set_context_mode(self.config.use_cache)
-
-        return {
-            "input_ids": batch_input_ids,
-            "past_key_values": past_key_values,
-            "position_ids": batch_position_ids,
-            "attention_mask": batch_attention_mask
-        }
+            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/config.json

@@ -1,25 +0,0 @@
-{
-  "_name_or_path": "THUDM/chatglm-6b",
-  "architectures": [
-    "ChatGLMModel"
-  ],
-  "auto_map": {
-    "AutoConfig": "configuration_chatglm.ChatGLMConfig",
-    "AutoModel": "modeling_chatglm.ChatGLMForConditionalGeneration",
-    "AutoModelForSeq2SeqLM": "modeling_chatglm.ChatGLMForConditionalGeneration"
-  },
-  "bos_token_id": 150004,
-  "eos_token_id": 150005,
-  "hidden_size": 4096,
-  "inner_hidden_size": 16384,
-  "layernorm_epsilon": 1e-05,
-  "max_sequence_length": 2048,
-  "model_type": "chatglm",
-  "num_attention_heads": 32,
-  "num_layers": 28,
-  "position_encoding_2d": true,
-  "torch_dtype": "float16",
-  "transformers_version": "4.23.1",
-  "use_cache": true,
-  "vocab_size": 150528
-}

models/configuration_chatglm.py

@@ -1,92 +0,0 @@
-""" ChatGLM model configuration """
-
-from transformers.configuration_utils import PretrainedConfig
-from transformers.utils import logging
-
-logger = logging.get_logger(__name__)
-
-
-class ChatGLMConfig(PretrainedConfig):
-    r"""
-    This is the configuration class to store the configuration of a [`~ChatGLMModel`].
-    It is used to instantiate an ChatGLM model according to the specified arguments, defining the model
-    architecture. Instantiating a configuration with the defaults will yield a similar configuration to that of
-    the ChatGLM-6B [THUDM/ChatGLM-6B](https://huggingface.co/THUDM/chatglm-6b) architecture.
-
-    Configuration objects inherit from  [`PretrainedConfig`] and can be used
-    to control the model outputs. Read the documentation from  [`PretrainedConfig`]
-    for more information.
-
-
-    Args:
-        vocab_size (`int`, *optional*, defaults to 150528):
-            Vocabulary size of the ChatGLM-6B model. Defines the number of different tokens that can be represented by the
-            `inputs_ids` passed when calling [`~ChatGLMModel`] or
-            [`~TFChatGLMModel`].
-        hidden_size (`int`, *optional*, defaults to 4096):
-            Dimension of the encoder layers and the pooler layer.
-        num_hidden_layers (`int`, *optional*, defaults to 28):
-            Number of hidden layers in the Transformer encoder.
-        num_attention_heads (`int`, *optional*, defaults to 32):
-            Number of attention heads for each attention layer in the Transformer encoder.
-        inner_hidden_size (`int`, *optional*, defaults to 16384):
-            Dimension of the "intermediate" (i.e., feed-forward) layer in the Transformer encoder.
-        max_sequence_length (`int`, *optional*, defaults to 512):
-            The maximum sequence length that this model might ever be used with.
-            Typically set this to something large just in case (e.g., 512 or 1024 or 2048).
-        layernorm_epsilon (`float`, *optional*, defaults to 1e-5):
-            The epsilon used by the layer normalization layers.
-        use_cache (`bool`, *optional*, defaults to `True`):
-            Whether the model should return the last key/values attentions (not used by all models).
-        Example:
-
-    ```python
-    >>> from configuration_chatglm import ChatGLMConfig
-    >>> from modeling_chatglm import ChatGLMModel
-
-    >>> # Initializing a ChatGLM-6B THUDM/ChatGLM-6B style configuration
-    >>> configuration = ChatGLMConfig()
-
-    >>> # Initializing a model from the THUDM/ChatGLM-6B style configuration
-    >>> model = ChatGLMModel(configuration)
-
-    >>> # Accessing the model configuration
-    >>> configuration = model.config
-    ```
-"""
-    model_type = "chatglm"
-
-    def __init__(
-            self,
-            vocab_size=150528,
-            hidden_size=4096,
-            num_layers=28,
-            num_attention_heads=32,
-            layernorm_epsilon=1e-5,
-            use_cache=False,
-            bos_token_id=150004,
-            eos_token_id=150005,
-            pad_token_id=0,
-            max_sequence_length=2048,
-            inner_hidden_size=16384,
-            position_encoding_2d=True,
-            **kwargs
-    ):
-        self.num_layers = num_layers
-        self.vocab_size = vocab_size
-        self.hidden_size = hidden_size
-        self.num_attention_heads = num_attention_heads
-        self.max_sequence_length = max_sequence_length
-        self.layernorm_epsilon = layernorm_epsilon
-        self.inner_hidden_size = inner_hidden_size
-        self.use_cache = use_cache
-        self.bos_token_id = bos_token_id
-        self.eos_token_id = eos_token_id
-        self.pad_token_id = pad_token_id
-        self.position_encoding_2d = position_encoding_2d
-        super().__init__(
-            pad_token_id=pad_token_id,
-            bos_token_id=bos_token_id,
-            eos_token_id=eos_token_id,
-            **kwargs
-        )

models/ice_text.model

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

models/tokenization_chatglm.py

@@ -1,17 +1,13 @@
 """Tokenization classes for ChatGLM."""
-import sys
-import unicodedata
 from typing import List, Optional, Union
-from functools import lru_cache
 import os
-import collections
-import re
 
 from transformers.tokenization_utils import PreTrainedTokenizer
-from icetk.text_tokenizer import TextTokenizer
-from icetk.utils import auto_create
-import icetk.sentencepiece_model_pb2 as sp_model
-from transformers.utils import logging
+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__)
 
@@ -20,61 +16,55 @@ PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES = {
 }
 
 
+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,
-        max_blank_length=80,
-        byte_fallback=True,
+            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 = self._build_text_tokenizer(encode_special_tokens=False)
-        self.special_text_tokenizer = self._build_text_tokenizer(encode_special_tokens=True)
-
-    @staticmethod
-    def _configure_tokenizer(
-        text_tokenizer: TextTokenizer,
-        special_tokens: List[str],
-        max_blank_length: int,
-        byte_fallback: bool,
-        encode_special_tokens=False,
-    ):
-        # special token
-        special_token_type = 4 if encode_special_tokens else 3  # 3 - CONTROL, 4 - USER_DEFINE
-        for token in special_tokens:
-            text_tokenizer.proto.pieces.append(
-                sp_model.ModelProto.SentencePiece(piece=token, score=0.0, type=special_token_type)
-            )
-        # whitespaces
-        for token in [SPTokenizer.get_tab_token()] + [
-            SPTokenizer.get_blank_token(i) for i in range(2, max_blank_length + 1)
-        ]:
-            text_tokenizer.proto.pieces.append(sp_model.ModelProto.SentencePiece(piece=token, score=0.0, type=4))
-        # byte fallback
-        if byte_fallback:
-            text_tokenizer.proto.trainer_spec.byte_fallback = True
-            for i in range(256):
-                text_tokenizer.proto.pieces.append(
-                    sp_model.ModelProto.SentencePiece(piece="<0x{:02X}>".format(i), score=0.0, type=6)
-                )
-        text_tokenizer.refresh()
-
-    def _build_text_tokenizer(self, encode_special_tokens=False):
-        tokenizer = TextTokenizer(self.vocab_file)
-        self._configure_tokenizer(
-            tokenizer, self.special_tokens, self.max_blank_length, self.byte_fallback, encode_special_tokens
-        )
-        return tokenizer
+        self.text_tokenizer = TextTokenizer(vocab_file)
 
-    def _get_text_tokenizer(self, encode_special_tokens=False):
-        if encode_special_tokens:
-            return self.special_text_tokenizer
-        else:
-            return self.text_tokenizer
+    def _get_text_tokenizer(self):
+        return self.text_tokenizer
 
     @staticmethod
     def get_blank_token(length: int):
@@ -85,10 +75,6 @@ class SPTokenizer:
     def get_tab_token():
         return f"<|tab|>"
 
-    @property
-    def num_image_tokens(self):
-        return 20000
-
     @property
     def num_text_tokens(self):
         return self.text_tokenizer.num_tokens
@@ -112,7 +98,7 @@ class SPTokenizer:
         return text
 
     def encode(
-        self, text: str, linebreak=True, whitespaces=True, special_tokens=False, add_dummy_prefix=True
+            self, text: str, linebreak=True, whitespaces=True, add_dummy_prefix=True
     ) -> List[int]:
         """
         @param text: Text to encode.
@@ -124,22 +110,31 @@ class SPTokenizer:
         text = self._preprocess(text, linebreak, whitespaces)
         if not add_dummy_prefix:
             text = "<n>" + text
-        tmp = self._get_text_tokenizer(encode_special_tokens=special_tokens).encode(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 decode(self, text_ids: List[int], special_tokens=False) -> 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(encode_special_tokens=special_tokens).decode(ids)
+    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, special_tokens=False, add_dummy_prefix=True
+            self, text: str, linebreak=True, whitespaces=True, add_dummy_prefix=True
     ) -> List[str]:
         """
         @param text: Text to encode.
@@ -151,7 +146,7 @@ class SPTokenizer:
         text = self._preprocess(text, linebreak, whitespaces)
         if not add_dummy_prefix:
             text = "<n>" + text
-        tokens = self._get_text_tokenizer(encode_special_tokens=special_tokens).tokenize(text)
+        tokens = self._get_text_tokenizer().tokenize(text)
         return tokens if add_dummy_prefix else tokens[2:]
 
     def __getitem__(self, x: Union[int, str]):
@@ -180,25 +175,36 @@ class ChatGLMTokenizer(PreTrainedTokenizer):
 
     vocab_files_names = {"vocab_file": "ice_text.model"}
     max_model_input_sizes = PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
-    model_input_names = ["input_ids"]
+    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='eos',
-            eop_token='eop',
+            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
         )
 
@@ -208,23 +214,29 @@ class ChatGLMTokenizer(PreTrainedTokenizer):
 
         self.bos_token = bos_token
         self.eos_token = eos_token
-        self.eop_token = eop_token
+        self.end_token = end_token
         self.mask_token = mask_token
-        self.gMASK_token = gmask_token
+        self.gmask_token = gmask_token
 
-        self.sp_tokenizer = SPTokenizer(vocab_file)
+        self.sp_tokenizer = SPTokenizer(vocab_file, num_image_tokens=num_image_tokens)
 
         """ Initialisation """
 
     @property
-    def eop_token_id(self) -> Optional[int]:
+    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 sentence token in the vocabulary. Returns `None` if the token has not been
+        `Optional[int]`: Id of the end of context token in the vocabulary. Returns `None` if the token has not been
         set.
         """
-        if self.eop_token is None:
+        if self.end_token is None:
             return None
-        return self.convert_tokens_to_ids(self.eop_token)
+        return self.convert_tokens_to_ids(self.end_token)
 
     @property
     def vocab_size(self):
@@ -256,25 +268,21 @@ class ChatGLMTokenizer(PreTrainedTokenizer):
 
         return seq
 
-    def decode(
+    def convert_tokens_to_string(self, tokens: List[str]) -> str:
+        return self.sp_tokenizer.decode_tokens(tokens)
+
+    def _decode(
             self,
-            token_ids: Union[List[int], List[List[int]]],
-            skip_special_tokens: bool = False,
-            clean_up_tokenization_spaces: bool = True,
-            spaces_between_special_tokens: bool = True,
+            token_ids: Union[int, List[int]],
             **kwargs
     ) -> str:
-        if isinstance(token_ids[0], list):
-            tokens = []
-            for single_token_ids in token_ids:
-                if self.pad_token_id in single_token_ids:  # remove pad
-                    single_token_ids = list(filter((self.pad_token_id).__ne__, single_token_ids))
-                tokens.append(self.sp_tokenizer.decode(single_token_ids))
-            return (tokens)
-        else:
-            if self.pad_token_id in token_ids:  # remove pad
-                token_ids = list(filter((self.pad_token_id).__ne__, token_ids))
-            return self.sp_tokenizer.decode(token_ids)
+        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. """
@@ -299,7 +307,7 @@ class ChatGLMTokenizer(PreTrainedTokenizer):
         """
         if os.path.isdir(save_directory):
             vocab_file = os.path.join(
-                save_directory, VOCAB_FILES_NAMES["vocab_file"]
+                save_directory, self.vocab_files_names["vocab_file"]
             )
         else:
             vocab_file = save_directory
@@ -331,16 +339,105 @@ class ChatGLMTokenizer(PreTrainedTokenizer):
         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_1
-        mask_ids = self.sp_tokenizer[self.mask_token]
-        gmask_ids = self.sp_tokenizer[self.gMASK_token]
-        if mask_ids not in token_ids_0 and gmask_ids not in token_ids_0:
-            token_ids_0 += [gmask_ids]
-
-        if token_ids_0[-1] != mask_ids and token_ids_0[-1] != gmask_ids:
-            token_ids_0 += [self.sp_tokenizer[self.eos_token]]
+            token_ids_0 = token_ids_0 + token_ids_1 + [eos_id]
+        return token_ids_0
 
-        token_ids_0 += [self.sp_tokenizer[self.bos_token]]
+    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)
 
-        return token_ids_0
+        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

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

requirements.txt

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