Skip to content

GitLab

Commit f7db21eb authored by lvzhen's avatar lvzhen
Browse files

first

parents
Pipeline #1580 canceled with stages
Hide whitespace changes
Inline Side-by-side
Showing with 2131 additions and 0 deletions
models/12B/specail_tokens_map.json 0 → 100644
View file @ f7db21eb
{
"bos_token": {
"content": "<_start>",
"lstrip": false,
"normalized": true,
"rstrip": false,
"single_word": false
},
"eos_token": {
"content": "<_end>",
"lstrip": false,
"normalized": true,
"rstrip": false,
"single_word": false
},
"pad_token": {
"content": "<_pad>",
"lstrip": false,
"normalized": true,
"rstrip": false,
"single_word": false
},
"unk_token": {
"content": "<_unk>",
"lstrip": false,
"normalized": true,
"rstrip": false,
"single_word": false
}
}
\ No newline at end of file
models/12B/tokenization_telechat.py 0 → 100644
View file @ f7db21eb
import os
from shutil import copyfile
from typing import Any, Dict, List, Optional, Tuple
import sentencepiece as spm
from transformers.tokenization_utils import AddedToken, PreTrainedTokenizer
from transformers.utils import logging
logger = logging.get_logger(__name__)
VOCAB_FILES_NAMES = {"vocab_file": "tokenizer.model"}
# TODO: when we get download url from huggingface, refresh the map
PRETRAINED_VOCAB_FILES_MAP = {
"vocab_file": {},
"tokenizer_file": {},
}
class TelechatTokenizer(PreTrainedTokenizer):
vocab_files_names = VOCAB_FILES_NAMES
pretrained_vocab_files_map = PRETRAINED_VOCAB_FILES_MAP
model_input_names = ["input_ids", "attention_mask"]
def __init__(
self,
vocab_file,
unk_token="<unk>",
bos_token="<_start>",
eos_token="<_end>",
pad_token="<_pad>",
sp_model_kwargs: Optional[Dict[str, Any]] = None,
add_bos_token=True,
add_eos_token=False,
clean_up_tokenization_spaces=False,
**kwargs,
):
self.sp_model_kwargs = {} if sp_model_kwargs is None else sp_model_kwargs
bos_token = AddedToken(bos_token, lstrip=False, rstrip=False) if isinstance(bos_token, str) else bos_token
eos_token = AddedToken(eos_token, lstrip=False, rstrip=False) if isinstance(eos_token, str) else eos_token
unk_token = AddedToken(unk_token, lstrip=False, rstrip=False) if isinstance(unk_token, str) else unk_token
pad_token = AddedToken(pad_token, lstrip=False, rstrip=False) if isinstance(pad_token, str) else pad_token
self.sp_model = spm.SentencePieceProcessor(**self.sp_model_kwargs)
self.sp_model.Load(vocab_file)
super().__init__(
bos_token=bos_token,
eos_token=eos_token,
unk_token=unk_token,
pad_token=pad_token,
add_bos_token=add_bos_token,
add_eos_token=add_eos_token,
sp_model_kwargs=self.sp_model_kwargs,
clean_up_tokenization_spaces=clean_up_tokenization_spaces,
**kwargs,
)
self.vocab_file = vocab_file
self.add_bos_token = add_bos_token
self.add_eos_token = add_eos_token
def __getstate__(self):
state = self.__dict__.copy()
state["sp_model"] = None
return state
def __setstate__(self, d):
self.__dict__ = d
self.sp_model = spm.SentencePieceProcessor(**self.sp_model_kwargs)
self.sp_model.Load(self.vocab_file)
@property
def vocab_size(self):
"""Returns vocab size"""
return self.sp_model.get_piece_size()
def get_vocab(self):
"""Returns vocab as a dict"""
vocab = {self.convert_ids_to_tokens(i): i for i in range(self.vocab_size)}
vocab.update(self.added_tokens_encoder)
return vocab
def _tokenize(self, text):
"""Returns a tokenized string."""
return self.sp_model.encode(text, out_type=str)
def _convert_token_to_id(self, token):
"""Converts a token (str) in an id using the vocab."""
return self.sp_model.piece_to_id(token)
def _convert_id_to_token(self, index):
"""Converts an index (integer) in a token (str) using the vocab."""
token = self.sp_model.IdToPiece(index)
return token
def convert_tokens_to_string(self, tokens):
"""Converts a sequence of tokens (string) in a single string."""
current_sub_tokens = []
out_string = ""
prev_is_special = False
for i, token in enumerate(tokens):
# make sure that special tokens are not decoded using sentencepiece model
if token in self.all_special_tokens:
if not prev_is_special and i != 0:
out_string += " "
out_string += self.sp_model.decode(current_sub_tokens) + token
prev_is_special = True
current_sub_tokens = []
else:
current_sub_tokens.append(token)
prev_is_special = False
out_string += self.sp_model.decode(current_sub_tokens)
return out_string
def save_vocabulary(self, save_directory, filename_prefix: Optional[str] = None) -> Tuple[str]:
"""
Save the vocabulary and special tokens file to a directory.
Args:
save_directory (`str`):
The directory in which to save the vocabulary.
Returns:
`Tuple(str)`: Paths to the files saved.
"""
if not os.path.isdir(save_directory):
logger.error(f"Vocabulary path ({save_directory}) should be a directory")
return
out_vocab_file = os.path.join(
save_directory, (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["vocab_file"]
)
if os.path.abspath(self.vocab_file) != os.path.abspath(out_vocab_file) and os.path.isfile(self.vocab_file):
copyfile(self.vocab_file, out_vocab_file)
elif not os.path.isfile(self.vocab_file):
with open(out_vocab_file, "wb") as fi:
content_spiece_model = self.sp_model.serialized_model_proto()
fi.write(content_spiece_model)
return (out_vocab_file,)
def build_inputs_with_special_tokens(self, token_ids_0, token_ids_1=None):
bos_token_id = [self.bos_token_id] if self.add_bos_token else []
eos_token_id = [self.eos_token_id] if self.add_eos_token else []
output = bos_token_id + token_ids_0 + eos_token_id
if token_ids_1 is not None:
output = output + bos_token_id + token_ids_1 + eos_token_id
return output
def get_special_tokens_mask(
self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None, already_has_special_tokens: bool = False
) -> List[int]:
"""
Retrieve sequence ids from a token list that has no special tokens added. This method is called when adding
special tokens using the tokenizer `prepare_for_model` method.
Args:
token_ids_0 (`List[int]`):
List of IDs.
token_ids_1 (`List[int]`, *optional*):
Optional second list of IDs for sequence pairs.
already_has_special_tokens (`bool`, *optional*, defaults to `False`):
Whether or not the token list is already formatted with special tokens for the model.
Returns:
`List[int]`: A list of integers in the range [0, 1]: 1 for a special token, 0 for a sequence token.
"""
if already_has_special_tokens:
return super().get_special_tokens_mask(
token_ids_0=token_ids_0, token_ids_1=token_ids_1, already_has_special_tokens=True
)
bos_token_id = [1] if self.add_bos_token else []
eos_token_id = [1] if self.add_eos_token else []
if token_ids_1 is None:
return bos_token_id + ([0] * len(token_ids_0)) + eos_token_id
return (
bos_token_id
+ ([0] * len(token_ids_0))
+ eos_token_id
+ bos_token_id
+ ([0] * len(token_ids_1))
+ eos_token_id
)
def create_token_type_ids_from_sequences(
self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None
) -> List[int]:
"""
Creates a mask from the two sequences passed to be used in a sequence-pair classification task. An ALBERT
sequence pair mask has the following format:
```
0 0 0 0 0 0 0 0 0 0 0 1 1 1 1 1 1 1 1 1
| first sequence | second sequence |
```
if token_ids_1 is None, only returns the first portion of the mask (0s).
Args:
token_ids_0 (`List[int]`):
List of ids.
token_ids_1 (`List[int]`, *optional*):
Optional second list of IDs for sequence pairs.
Returns:
`List[int]`: List of [token type IDs](../glossary#token-type-ids) according to the given sequence(s).
"""
bos_token_id = [self.bos_token_id] if self.add_bos_token else []
eos_token_id = [self.eos_token_id] if self.add_eos_token else []
output = [0] * len(bos_token_id + token_ids_0 + eos_token_id)
if token_ids_1 is not None:
output += [1] * len(bos_token_id + token_ids_1 + eos_token_id)
return output
models/12B/tokenizer_config.json 0 → 100644
View file @ f7db21eb
{
"name_or_path": "ChinaTelecom/telechat-12b",
"tokenizer_class": "TelechatTokenizer",
"auto_map": {
"AutoTokenizer": [
"tokenization_telechat.TelechatTokenizer",
null
]
},
"add_bos_token": false,
"add_eos_token": false,
"use_fast": false,
"clean_up_tokenization_spaces": false,
"eos_token": {
"__type": "AddedToken",
"content": "<_end>",
"lstrip": false,
"normalized": true,
"rstrip": false,
"single_word": true
},
"model_max_length": 100000000,
"sp_model_kwargs": {},
"pad_token": {
"__type": "AddedToken",
"content": "<_pad>",
"lstrip": false,
"normalized": true,
"rstrip": false,
"single_word": true
},
"unk_token": {
"__type": "AddedToken",
"content": "<_end>",
"lstrip": false,
"normalized": true,
"rstrip": false,
"single_word": true
}
}
models/12B_4bit/config.json 0 → 100644
View file @ f7db21eb
{
"apply_residual_connection_post_layernorm": false,
"architectures": [
"TelechatForCausalLM"
],
"auto_map": {
"AutoConfig": "configuration_telechat.TelechatConfig",
"AutoModelForCausalLM": "modeling_telechat.TelechatForCausalLM"
},
"attention_dropout": 0.0,
"attention_softmax_in_fp32": true,
"bias_dropout_fusion": true,
"bos_token_id": 1,
"eos_token_id": 2,
"hidden_dropout": 0.0,
"hidden_size": 5120,
"initializer_range": 0.02,
"layer_norm_epsilon": 1e-05,
"masked_softmax_fusion": true,
"model_type": "telechat",
"n_head": 32,
"n_inner": null,
"n_layer": 38,
"offset_alibi": 100,
"pad_token_id": 3,
"pretraining_tp": 2,
"seq_length": 8192,
"skip_bias_add": true,
"skip_bias_add_qkv": false,
"slow_but_exact": false,
"transformers_version": "4.24.0",
"unk_token_id": 0,
"use_cache": true,
"vocab_size": 120000,
"ffn_hidden_size": 12288,
"flash_attn":true,
"tie_word_embeddings":false,
"training_seqlen":8192,
"logn":false,
"semi_causal":false,
"embed_layernorm":false
}
models/12B_4bit/configuration_telechat.py 0 → 100644
View file @ f7db21eb
# coding=utf-8
# Copyright 2022 the Big Science Workshop and HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
""" Telechat configuration"""
from packaging import version
from collections import OrderedDict
from transformers.utils import is_torch_available, logging
from transformers.configuration_utils import PretrainedConfig
from typing import TYPE_CHECKING, Any, List, Mapping, Optional
logger = logging.get_logger(__name__)
class TelechatConfig(PretrainedConfig):
"""
Args:
vocab_size (`int`, *optional*, defaults to 160256): Vocabulary size of the Telechat model.
hidden_size (`int`, *optional*, defaults to 4096): Dimensionality of the embeddings and hidden states.
ffn_hidden_size (`int`, *optional*, defaults to 12288): Dimensionality of the feed-forward hidden states.
n_layer (`int`, *optional*, defaults to 30): Number of hidden layers in the Transformer
n_head (`int`, *optional*, defaults to 32): Number of attention heads for each attention layer.
layer_norm_epsilon (`float`, *optional*, defaults to 1e-5): The epsilon to use in the layer normalization layers.
initializer_range (`float`, *optional*, defaults to 0.02): The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
apply_residual_connection_post_layernorm (`bool`, *optional*, defaults to `False`): If enabled, use the layer norm of the hidden states as the residual in the transformer blocks
hidden_dropout (`float`, *optional*, defaults to 0.0): Dropout rate of the dropout function on the bias dropout.
attention_dropout (`float`, *optional*, defaults to 0.0): Dropout rate applied to the attention probs
use_cache (`bool`, *optional*, defaults to `True`): Whether or not the model should return the last key/values attentions.
training_seqlen (`int`, *optional*, defaults to 8192): Sequence length during last finetuning.
logn (`bool`, *optional*, defaults to `True`): Whether or not to use logN during extrapolation.
embed_layernorm (`bool`, *optional*, defaults to `True`): Whether or not to use embedding layernorm.
"""
model_type = "telechat"
keys_to_ignore_at_inference = ["past_key_values"]
attribute_map = {
"num_hidden_layers": "n_layer",
"num_attention_heads": "n_head",
}
def __init__(
self,
vocab_size=160256,
hidden_size=4096,
n_layer=30,
n_head=32,
layer_norm_epsilon=1e-5,
initializer_range=0.02,
use_cache=True,
bos_token_id=1,
eos_token_id=2,
apply_residual_connection_post_layernorm=False,
hidden_dropout=0.0,
attention_dropout=0.0,
ffn_hidden_size=12288,
training_seqlen = 8192,
logn = True,
embed_layernorm = False,
**kwargs,
):
self.vocab_size = vocab_size
n_embed = kwargs.pop("n_embed", None)
self.hidden_size = hidden_size if n_embed is None else n_embed
self.n_layer = n_layer
self.n_head = n_head
self.layer_norm_epsilon = layer_norm_epsilon
self.initializer_range = initializer_range
self.use_cache = use_cache
self.apply_residual_connection_post_layernorm = apply_residual_connection_post_layernorm
self.hidden_dropout = hidden_dropout
self.attention_dropout = attention_dropout
self.bos_token_id = bos_token_id
self.eos_token_id = eos_token_id
self.logn = logn
self.ffn_hidden_size = ffn_hidden_size
self.training_seqlen = training_seqlen
self.embed_layernorm = embed_layernorm
super().__init__(bos_token_id=bos_token_id, eos_token_id=eos_token_id, **kwargs)
models/12B_4bit/generation_config.json 0 → 100644
View file @ f7db21eb
{
"max_length": 8192,
"do_sample": false,
"use_cache": true,
"temperature": 0.3,
"top_k": 5,
"top_p": 0.85,
"repetition_penalty": 1.01,
"pad_token_id": 3,
"bos_token_id": 1,
"eos_token_id": 2,
"user_token_id": 20,
"bot_token_id": 21
}
models/12B_4bit/generation_utils.py 0 → 100644
View file @ f7db21eb
from typing import Optional
from collections import deque
from queue import Queue
import copy
class History:
def __init__(self, tokenizer, history):
'''
init from a list of dict
'''
# use deque to meet some special situation
self.input_history = deque()
self.tokenizer = tokenizer
if history:
self._transfer_from_list(history)
def _transfer_from_list(self, history):
for message in history:
content = message.get("content")
# the token result may not be equal to the result model gen
message.update(self.tokenizer(content))
self.input_history.append(message)
def append(self, message):
content = message.get("content")
if "input_ids" not in message or "attention_mask" not in message:
message.update(self.tokenizer(content))
self.input_history.append(message)
def append_left(self, message):
content = message.get("content")
if "input_ids" not in message or "attention_mask" not in message:
message.update(self.tokenizer(content))
self.input_history.appendleft(message)
def pop(self):
x = self.input_history.pop()
return x
def pop_left(self):
x = self.pop_left()
return x
def update(self, message):
self.input_history.pop()
self.append(message)
def __len__(self):
return self.input_history.__len__()
def __str__(self):
return self.input_history.__str__()
def __copy__(self):
new_instance = type(self)(self.tokenizer, [])
new_instance.input_history = copy.copy(self.input_history)
return new_instance
def __deepcopy__(self, memodict={}):
new_instance = type(self)(self.tokenizer, [])
new_instance.input_history = copy.deepcopy(self.input_history)
return new_instance
class TelechatIterTextStreamer:
"""
With reference to the TextIterStreamers in transformers, we have rewritten this class
"""
def __init__(
self, tokenizer, history: History = None, skip_prompt: bool = False, timeout: Optional[float] = None,
**decode_kwargs
):
self.tokenizer = tokenizer
self.history = history
self.skip_prompt = skip_prompt
self.timeout = timeout
self.decode_kwargs = decode_kwargs
self.text_queue = Queue()
self.cache_time = 0
self.text_until = ""
self.token_until = []
self.stop_signal = None
self.next_tokens_are_prompt = True
self.history.append({"role": "bot", "content": self.text_until})
def put(self, value):
"""
put printable text into queue
"""
if len(value.shape) > 1 and value.shape[0] > 1:
raise ValueError("TextStreamer only supports batch size 1")
elif len(value.shape) > 1:
value = value[0]
if self.skip_prompt and self.next_tokens_are_prompt:
self.next_tokens_are_prompt = False
return
if value[-1] == self.tokenizer.eos_token_id:
return
# there may be some smart way to decode.
self.token_until.extend(value.tolist())
text = self.tokenizer.decode(self.token_until, **self.decode_kwargs)
if self._is_printable(text) or self.cache_time >= 6:
output_text = text[len(self.text_until):]
self.text_until = text
else:
self.cache_time+=1
return
self.on_finalized_text(output_text)
def end(self):
"""Flushes any remaining cache and prints a newline to stdout."""
# Flush the cache, if it exists
text = self.tokenizer.decode(self.token_until, **self.decode_kwargs)
output_text = text[len(self.text_until):]
self.text_until = text
self.on_finalized_text(output_text, stream_end=True)
self.clear_cache()
def clear_cache(self):
self.cache_time = 0
self.token_until = []
self.text_until = ""
self.history = None
self.next_tokens_are_prompt = True
def on_finalized_text(self, text: str, stream_end: bool = False):
"""Put the text tuple in the queue."""
self.history.update({"role": "bot", "content": self.text_until, "input_ids": self.token_until,
"attention_mask": [1] * len(self.token_until)})
self.text_queue.put((text, self.history), timeout=self.timeout)
if stream_end:
self.text_queue.put((self.stop_signal, self.history), timeout=self.timeout)
@staticmethod
def _is_printable(cp):
"""Checks whether tokens can be decoded or not"""
if "�" in cp:
return False
return True
def __iter__(self):
return self
def __next__(self):
value_now, history_until = self.text_queue.get(timeout=self.timeout)
if value_now == self.stop_signal:
raise StopIteration()
else:
return value_now, history_until
models/12B_4bit/modeling_telechat.py 0 → 100644
View file @ f7db21eb
# coding=utf-8
# Copyright 2022 HuggingFace Inc. team and BigScience workshop.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.
# Copyright (c) 2021 EleutherAI
# This file is based on code by the authors denoted below and has been modified from its original version.
#
# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""PyTorch TELECHAT model."""
import warnings
from typing import Optional, Tuple, Union, List, Dict
from threading import Thread
import torch
import math
import copy
from torch import nn
import torch.utils.checkpoint
from torch.nn import functional as F
from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, LayerNorm, MSELoss
from transformers.modeling_outputs import (
BaseModelOutputWithPastAndCrossAttentions,
CausalLMOutputWithCrossAttentions
)
from transformers.modeling_utils import PreTrainedModel
from transformers.utils import logging
from transformers import GenerationConfig
from .configuration_telechat import TelechatConfig
from .generation_utils import History, TelechatIterTextStreamer
logger = logging.get_logger(__name__)
_CHECKPOINT_FOR_DOC = "telechat"
_CONFIG_FOR_DOC = "TelechatConfig"
TELECHAT_PRETRAINED_MODEL_ARCHIVE_LIST = []
try:
from einops import rearrange
except ImportError:
rearrange = None
use_flash_attn = True
try:
from flash_attn.flash_attn_interface import flash_attn_unpadded_func
except ImportError:
try:
from flash_attn.flash_attn_interface import flash_attn_varlen_func as flash_attn_unpadded_func
except ImportError:
flash_attn_unpadded_func = None
class RotaryEmbedding(torch.nn.Module):
# Extracted from: https://github.com/EleutherAI/gpt-neox
def __init__(self, dim, config, base=10000,precision=torch.half):
super().__init__()
self.config = config
self.dim = dim
self.base = base
self.max_seq_len_cached = None
self.cos_cached = None
self.sin_cached = None
self.precision = precision
def get_mscale(self, scale=1):
if scale <= 1:
return 1.0
return 0.1 * math.log(scale) + 1.0
def get_ntk_alpha(self, true_seq_len):
context_value = math.log(true_seq_len / 8192, 2) + 1
# ntk_alpha = 2 ** context_value - 1
ntk_alpha = 2 ** math.ceil(context_value) - 1
ntk_alpha = max(ntk_alpha, 1)
return ntk_alpha
def forward(self, x, seq_dim=0, seq_len=None):
if seq_len is None:
seq_len = x.shape[seq_dim]
seq_len = max(seq_len, self.config.training_seqlen)
ntk_alpha = self.get_ntk_alpha(seq_len)
mscale = float(self.get_mscale(seq_len / self.config.training_seqlen))
base = self.base * ntk_alpha ** (self.dim / (self.dim - 2))
inv_freq = 1.0 / (base ** (torch.arange(0, self.dim, 2, device=x.device).float( )/ self.dim ))
max_seq_len_cached = seq_len
t = torch.arange(max_seq_len_cached, device=x.device, dtype=inv_freq.dtype)
freqs = torch.einsum('i,j->ij', t, inv_freq)
# Different from paper, but it uses a different permutation in order to obtain the same calculation
emb = torch.cat((freqs, freqs), dim=-1).to(x.device)
if self.precision == torch.bfloat16:
emb = emb.float()
# [sx, 1 (b * np), hn]
cos_cached = mscale *emb.cos()[:, None, :].half()
sin_cached = mscale *emb.sin()[:, None, :].half()
if self.precision == torch.bfloat16:
cos_cached = cos_cached.bfloat16()
sin_cached = sin_cached.bfloat16()
return cos_cached[:seq_len, ...], sin_cached[:seq_len, ...]
# rotary pos emb helpers:
def rotate_half(x):
x1, x2 = x[..., :x.shape[-1] // 2], x[..., x.shape[-1] // 2:]
return torch.cat((-x2, x1), dim=x1.ndim - 1) # dim=-1 triggers a bug in earlier torch versions
def apply_rotary_pos_emb_torch(q, k, cos, sin, offset: int = 0): # jitting fails with bf16
cos, sin = cos[offset:q.shape[0] + offset, ...], sin[offset:q.shape[0] + offset, ...]
return (q * cos) + (rotate_half(q) * sin), (k * cos) + (rotate_half(k) * sin)
class MixedFusedRMSNorm(nn.Module):
# Extracted from https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/modeling_llama.py
def __init__(self, hidden_size, eps=1e-6):
super().__init__()
self.weight = nn.Parameter(torch.ones(hidden_size))
self.variance_epsilon = eps
def forward(self, hidden_states):
input_dtype = hidden_states.dtype
hidden_states = hidden_states.to(torch.float32)
variance = hidden_states.pow(2).mean(-1, keepdim=True)
hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
return self.weight * hidden_states.to(input_dtype)
class FlashSelfAttention(torch.nn.Module):
# Extracted from https://github.com/microsoft/Megatron-DeepSpeed/blob/main/megatron/model/transformer.py
"""Implement the scaled dot product attention with softmax.
Arguments
---------
softmax_scale: The temperature to use for the softmax attention.
(default: 1/sqrt(d_keys) where d_keys is computed at
runtime)
attention_dropout: The dropout rate to apply to the attention
(default: 0.0)
"""
def __init__(self, causal=False, softmax_scale=None, attention_dropout=0.0,
device=None, dtype=None):
super().__init__()
assert flash_attn_unpadded_func is not None, ('Please install FlashAttention first, '
'e.g., with pip install flash-attn')
assert rearrange is not None, 'Please install einops first, e.g., with pip install einops'
self.causal = causal
self.softmax_scale = softmax_scale
self.dropout_p = attention_dropout
def forward(self, q, k, v):
"""Implements the multihead softmax attention.
Arguments
---------
q, k, v: The tensor containing the query, key, and value. (B, S, H, D)
"""
assert all((i.dtype in [torch.float16, torch.bfloat16] for i in (q, k, v)))
assert all((i.is_cuda for i in (q, k, v)))
batch_size, seqlen_q = q.shape[0], q.shape[1]
seqlen_k = k.shape[1]
q, k, v = [rearrange(x, 'b s ... -> (b s) ...') for x in [q, k, v]]
cu_seqlens_q = torch.arange(0, (batch_size + 1) * seqlen_q, step=seqlen_q, dtype=torch.int32,
device=q.device)
if self.training:
# during training q,k,v always have same seqlen
assert seqlen_k == seqlen_q
is_causal = self.causal
cu_seqlens_k = cu_seqlens_q
dropout_p = self.dropout_p
else:
# turn off FA causal mask after first inference autoregressive iteration
# only on first autoregressive step q,k,v have same seqlen
is_causal = seqlen_q == seqlen_k
cu_seqlens_k = torch.arange(0, (batch_size + 1) * seqlen_k, step=seqlen_k, dtype=torch.int32,
device=q.device)
dropout_p = 0
output = flash_attn_unpadded_func(
q, k, v, cu_seqlens_q, cu_seqlens_k, seqlen_q, seqlen_k,
dropout_p=dropout_p,
softmax_scale=self.softmax_scale, causal=is_causal
)
output = rearrange(output, '(b s) ... -> b s ...', b=batch_size)
return output
def _make_causal_mask(
input_ids_shape: torch.Size, device: torch.device, past_key_values_length: int
) -> torch.BoolTensor:
"""
Make causal mask used for self-attention.
"""
batch_size, target_length = input_ids_shape
mask = torch.empty((target_length, target_length + past_key_values_length), dtype=torch.bool, device=device)
# ONNX doesn't support `torch.Tensor.triu` properly, thus we use this workaround
seq_ids = torch.arange(target_length, device=device)
mask[:, past_key_values_length:] = seq_ids[:, None] < seq_ids[None, :]
if past_key_values_length > 0:
mask[:, :past_key_values_length] = False
expanded_mask = mask[None, None, :, :].expand(batch_size, 1, target_length, target_length + past_key_values_length)
return expanded_mask
def _expand_mask(mask: torch.Tensor, tgt_length: int) -> torch.BoolTensor:
"""
Expands attention_mask from `[batch_size, src_length]` to `[batch_size, 1, tgt_length, src_length]`.
"""
batch_size, src_length = mask.shape
tgt_length = tgt_length if tgt_length is not None else src_length
expanded_mask = ~(mask[:, None, None, :].to(torch.bool))
return expanded_mask.expand(batch_size, 1, tgt_length, src_length)
def dropout_add(x: torch.Tensor, residual: torch.Tensor, prob: float, training: bool) -> torch.Tensor:
"""
Dropout add function
Args:
x (`torch.tensor`, *required*):
input tensor
residual (`torch.tensor`, *required*):
residual tensor
prob (`float`, *required*):
dropout probability
training (`bool`, *required*):
training mode
"""
out = F.dropout(x, p=prob, training=training)
out = residual + out
return out
def telechat_gelu_forward(x: torch.Tensor) -> torch.Tensor:
"""
Custom bias GELU function. Adapted from Megatron-DeepSpeed code. Here we use a simple implementation (inference) to
make the model jitable.
Args:
x (`torch.tensor`, *required*):
input hidden states
"""
return x * 0.5 * (1.0 + torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)))
def telechat_gelu_back(g: torch.Tensor, x: torch.Tensor) -> torch.Tensor:
"""
gradient of tanh approximation of gelu gradient of actual gelu is: 0.5 * (1. + torch.erf(x * 0.70710678)) +
0.3989423 * x * torch.exp(-0.5 * x * x)
Args:
g (`torch.tensor`, *required*):
gradient output tensor
x (`torch.tensor`, *required*):
input tensor
"""
x = x[0] # x is a tuple of 1 element, needs to unpack it first
tanh_out = torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x))
# sqrt(2/pi) * 3 * 0.044715 -> 0.1070322243
ff = 0.5 * x * ((1 - tanh_out * tanh_out) * (0.79788456 + 0.1070322243 * x * x)) + 0.5 * (1 + tanh_out)
return ff * g
class GeLUFunction(torch.autograd.Function):
@staticmethod
def forward(ctx, input: torch.Tensor) -> torch.Tensor:
ctx.save_for_backward(input)
return telechat_gelu_forward(input)
@staticmethod
def backward(ctx, grad_output: torch.Tensor) -> torch.Tensor:
input = ctx.saved_tensors
tmp = telechat_gelu_back(grad_output, input)
return tmp
class TelechatGelu(nn.Module):
"""
TelechatBiasGelu wrapper function that make use of the simple function on inference mode to make the model
torchscriptable and use the autograd function in training mode to get the accurate results of the gradients Partly
copied from Megatron-DeepSpeed code and adapted for our needs
See here why autograd functions are not torchscriptable: https://github.com/pytorch/pytorch/issues/22329
"""
def __init__(self):
super().__init__()
def forward(self, x: torch.Tensor) -> torch.Tensor:
if self.training:
return GeLUFunction.apply(x)
else:
return telechat_gelu_forward(x)
class TelechatAttention(nn.Module):
def __init__(self, config: TelechatConfig, layer_idx):
super().__init__()
self.kv_cache = None
self.layer_idx = layer_idx
self.hidden_size = config.hidden_size
self.num_heads = config.n_head
self.head_dim = self.hidden_size // self.num_heads
self.split_size = self.hidden_size
self.hidden_dropout = config.hidden_dropout
self.config = config
if self.head_dim * self.num_heads != self.hidden_size:
raise ValueError(
f"`hidden_size` must be divisible by num_heads (got `hidden_size`: {self.hidden_size} and `num_heads`:"
f" {self.num_heads})."
)
# Layer-wise attention scaling
self.inv_norm_factor = 1.0 / math.sqrt(self.head_dim)
self.beta = 1.0
self.num_key_value_heads = self.num_heads
kv_projection_size = self.head_dim * self.num_key_value_heads
self.num_key_value_groups = self.num_heads // self.num_key_value_heads
self.query = nn.Linear(self.hidden_size, self.hidden_size, bias=False)
self.key_value = nn.Linear(self.hidden_size, kv_projection_size * 2, bias=False)
self.dense = nn.Linear(self.hidden_size, self.hidden_size)
self.attention_dropout = nn.Dropout(config.attention_dropout)
self.rotary_emb = RotaryEmbedding(self.head_dim, config=config)
self.core_attention_flash = FlashSelfAttention(
causal=True, attention_dropout=config.attention_dropout
)
self.last_key_layer = None
def repeat_kv(self, hidden_states, n_rep):
slen, batch, num_key_value_heads_per_partition, head_dim = hidden_states.shape
if n_rep == 1:
return hidden_states
hidden_states = hidden_states[:, :, :, None, :].expand(slen, batch, num_key_value_heads_per_partition, n_rep,
head_dim)
return hidden_states.reshape(slen, batch, num_key_value_heads_per_partition * n_rep, head_dim)
def split_tensor_along_last_dim(self,
tensor: torch.Tensor,
num_partitions: int,
contiguous_split_chunks: bool = False,
):
# Get the size and dimension.
last_dim = tensor.dim() - 1
last_dim_size = tensor.size()[last_dim] // num_partitions
# Split.
tensor_list = torch.split(tensor, last_dim_size, dim=last_dim)
# Note: torch.split does not create contiguous tensors by default.
if contiguous_split_chunks:
return tuple(chunk.contiguous() for chunk in tensor_list)
return tensor_list
def _merge_heads(self, x: torch.Tensor) -> torch.Tensor:
batch_size_and_num_heads, seq_length, _ = x.shape
batch_size = batch_size_and_num_heads // self.num_heads
x = x.view(batch_size, self.num_heads, seq_length, self.head_dim)
x = x.permute(0, 2, 1, 3)
return x.reshape(batch_size, seq_length, self.num_heads * self.head_dim)
def forward(
self,
hidden_states: torch.Tensor,
residual: torch.Tensor,
attention_mask: torch.Tensor,
layer_past: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
use_cache: bool = False,
output_attentions: bool = False,
):
hidden_states = hidden_states.transpose(1, 0)
query_layer = self.query(hidden_states)
new_tensor_shape = query_layer.size()[:-1] + \
(self.num_heads,
self.head_dim)
query_layer = query_layer.view(*new_tensor_shape)
mixed_kv_layer = self.key_value(hidden_states)
new_tensor_shape = mixed_kv_layer.size()[:-1] + \
(self.num_key_value_heads,
2 * self.head_dim)
mixed_kv_layer = mixed_kv_layer.view(*new_tensor_shape)
(key_layer, value_layer) = self.split_tensor_along_last_dim(mixed_kv_layer, 2)
output_size = (query_layer.size(1),
query_layer.size(2),
query_layer.size(0),
key_layer.size(0))
query_layer = query_layer.view(output_size[2], output_size[0] * output_size[1], -1)
key_layer = key_layer.view(output_size[3], output_size[0] * output_size[1], -1)
apply_rotary_fn = apply_rotary_pos_emb_torch
seq_len = key_layer.shape[0]
offset = 0
if use_cache and layer_past != None:
past_key, past_value = layer_past
offset = past_key.shape[0]
seq_len += offset
cos, sin = self.rotary_emb(value_layer)
query_layer, key_layer = apply_rotary_fn(query_layer, key_layer, cos, sin, offset=offset)
if use_cache:
if layer_past != None:
past_key, past_value = layer_past
key_layer = torch.cat((past_key, key_layer[-1, ...].unsqueeze(0)), dim=0)
value_layer = torch.cat((past_value, value_layer[-1, ...].unsqueeze(0)), dim=0)
layer_past = key_layer, value_layer
s, bz, head, dim = value_layer.shape
s_key = key_layer.shape[0]
s_query = query_layer.shape[0]
query_layer = query_layer.reshape((s_query, bz, head, dim))
key_layer = key_layer.reshape((s_key, bz, head, dim))
if self.config.flash_attn:
q, k, v = [rearrange(x, 's b ... -> b s ...').contiguous() for x in
(query_layer, key_layer, value_layer)]
context_layer = self.core_attention_flash(q, k, v)
context_layer = rearrange(context_layer, 'b s h d -> b s (h d)').contiguous()
else:
##[sq, b, np, hn] -> [sq, b * np, hn]
query_layer = query_layer.reshape(s_query, bz * self.num_heads, dim)
# [sk, b, np, hn] -> [sk, b * np, hn]
key_layer = key_layer.reshape(s_key, bz * self.num_heads, dim)
matmul_result = self.inv_norm_factor * torch.einsum('bik,bkj->bij', query_layer.transpose(0, 1),
key_layer.transpose(0, 1).transpose(1, 2))
attention_scores = matmul_result.view(bz, self.num_heads, s_query, s_key)
input_dtype = attention_scores.dtype
if input_dtype == torch.float16 or input_dtype == torch.bfloat16:
attention_scores = attention_scores.to(torch.float)
attn_weights = torch.masked_fill(attention_scores, attention_mask, torch.finfo(attention_scores.dtype).min)
attention_probs = F.softmax(attn_weights, dim=-1).to(input_dtype) ##dtype = torch.float32
attention_probs = self.attention_dropout(attention_probs)
attention_probs_reshaped = attention_probs.view(bz * self.num_heads, s_query, s_key)
value_layer = value_layer.reshape(s_key, bz * self.num_heads, dim)
context_layer = torch.bmm(attention_probs_reshaped, value_layer.transpose(0, 1))
context_layer = self._merge_heads(context_layer)
output_tensor = self.dense(context_layer)
output_tensor = dropout_add(output_tensor, residual, self.hidden_dropout, self.training)
present = None
outputs = (output_tensor, present)
if output_attentions:
outputs += (attention_probs,)
return output_tensor, layer_past
class TelechatMLP(nn.Module):
def __init__(self, config: TelechatConfig):
super().__init__()
hidden_size = config.hidden_size
self.gate_proj = nn.Linear(hidden_size, config.ffn_hidden_size, bias=False)
self.up_proj = nn.Linear(hidden_size, config.ffn_hidden_size, bias=False)
self.down_proj = nn.Linear(config.ffn_hidden_size, hidden_size, bias=True)
self.hidden_dropout = config.hidden_dropout
def forward(self, hidden_states: torch.Tensor, residual: torch.Tensor) -> torch.Tensor:
intermediate_output = self.down_proj(F.silu(self.gate_proj(hidden_states)) * self.up_proj(hidden_states))
output = dropout_add(intermediate_output, residual, self.hidden_dropout, self.training)
return output
class TelechatBlock(nn.Module):
def __init__(self, config: TelechatConfig, layer_idx):
super().__init__()
hidden_size = config.hidden_size
self.input_layernorm = MixedFusedRMSNorm(hidden_size, eps=config.layer_norm_epsilon)
self.num_heads = config.n_head
self.layer_idx = layer_idx
self.self_attention = TelechatAttention(config, layer_idx)
self.post_attention_layernorm = MixedFusedRMSNorm(hidden_size, eps=config.layer_norm_epsilon)
self.mlp = TelechatMLP(config)
self.apply_residual_connection_post_layernorm = config.apply_residual_connection_post_layernorm
self.hidden_dropout = config.hidden_dropout
def forward(
self,
hidden_states: torch.Tensor,
attention_mask: torch.Tensor,
layer_past: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
use_cache: bool = False,
output_attentions: bool = False,
):
layernorm_output = self.input_layernorm(hidden_states)
if self.apply_residual_connection_post_layernorm:
residual = layernorm_output
else:
residual = hidden_states
attn_outputs = self.self_attention(
layernorm_output,
residual,
layer_past=layer_past,
attention_mask=attention_mask,
use_cache=use_cache,
output_attentions=output_attentions,
)
attention_output = attn_outputs[0]
outputs = attn_outputs[1:]
layernorm_output = self.post_attention_layernorm(attention_output)
if self.apply_residual_connection_post_layernorm:
residual = layernorm_output
else:
residual = attention_output
output = self.mlp(layernorm_output, residual)
if use_cache:
outputs = (output,) + outputs
else:
outputs = (output,) + outputs[1:]
return outputs
class TelechatPreTrainedModel(PreTrainedModel):
config_class = TelechatConfig
base_model_prefix = "transformer"
supports_gradient_checkpointing = True
_no_split_modules = ["TelechatBlock"]
_skip_keys_device_placement = "past_key_values"
def __init__(self, *inputs, **kwargs):
super().__init__(*inputs, **kwargs)
def _init_weights(self, module: nn.Module):
"""Initialize the weights."""
if isinstance(module, nn.Linear):
module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
if module.bias is not None:
module.bias.data.zero_()
elif isinstance(module, nn.Embedding):
module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
if module.padding_idx is not None:
module.weight.data[module.padding_idx].zero_()
elif isinstance(module, LayerNorm):
module.bias.data.zero_()
module.weight.data.fill_(1.0)
def _set_gradient_checkpointing(self, module: nn.Module, value: bool = False):
if isinstance(module, TelechatModel):
module.gradient_checkpointing = value
class TelechatModel(TelechatPreTrainedModel):
def __init__(self, config: TelechatConfig):
super().__init__(config)
self.embed_dim = config.hidden_size
self.num_heads = config.n_head
self.config = config
self.word_embeddings = nn.Embedding(config.vocab_size, self.embed_dim)
if self.config.embed_layernorm:
self.word_embeddings_layernorm = MixedFusedRMSNorm(self.embed_dim, eps=config.layer_norm_epsilon)
self.h = nn.ModuleList([TelechatBlock(config, _) for _ in range(config.num_hidden_layers)])
self.ln_f = MixedFusedRMSNorm(self.embed_dim, eps=config.layer_norm_epsilon)
self.gradient_checkpointing = False
self.post_init()
def get_input_embeddings(self):
return self.word_embeddings
def _prepare_attn_mask(
self, attention_mask: torch.Tensor, input_shape: Tuple[int, int], past_key_values_length: int
) -> torch.BoolTensor:
combined_attention_mask = None
device = attention_mask.device
_, src_length = input_shape
if src_length > 1:
combined_attention_mask = _make_causal_mask(
input_shape, device=device, past_key_values_length=past_key_values_length
)
expanded_attn_mask = _expand_mask(attention_mask, tgt_length=src_length)
combined_attention_mask = (
expanded_attn_mask if combined_attention_mask is None else expanded_attn_mask | combined_attention_mask
)
return combined_attention_mask
def set_input_embeddings(self, new_embeddings: torch.Tensor):
self.word_embeddings = new_embeddings
def forward(
self,
input_ids: Optional[torch.LongTensor] = None,
past_key_values: Optional[Tuple[Tuple[torch.Tensor, torch.Tensor], ...]] = None,
attention_mask: Optional[torch.Tensor] = None,
inputs_embeds: Optional[torch.LongTensor] = None,
use_cache: Optional[bool] = None,
output_attentions: Optional[bool] = None,
output_hidden_states: Optional[bool] = None,
return_dict: Optional[bool] = None,
**deprecated_arguments,
) -> Union[Tuple[torch.Tensor, ...], BaseModelOutputWithPastAndCrossAttentions]:
output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
output_hidden_states = (
output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
)
use_cache = use_cache if use_cache is not None else self.config.use_cache
return_dict = return_dict if return_dict is not None else self.config.use_return_dict
if input_ids is not None:
batch_size, seq_length = input_ids.shape
elif inputs_embeds is not None:
batch_size, seq_length, _ = inputs_embeds.shape
if past_key_values is None:
past_key_values = tuple([None] * len(self.h))
if inputs_embeds is None:
inputs_embeds = self.word_embeddings(input_ids)
hidden_states = inputs_embeds
if self.config.embed_layernorm:
hidden_states = self.word_embeddings_layernorm(inputs_embeds)
presents = () if use_cache else None
all_self_attentions = () if output_attentions else None
all_hidden_states = () if output_hidden_states else None
if self.gradient_checkpointing and self.training:
if use_cache:
use_cache = False
seq_length_with_past = seq_length
past_key_values_length = 0
if past_key_values[0] is not None:
past_key_values_length = past_key_values[0][0].shape[2]
seq_length_with_past = seq_length_with_past + past_key_values_length
if attention_mask is None:
attention_mask = torch.ones((batch_size, seq_length_with_past), device=hidden_states.device)
else:
attention_mask = attention_mask.to(hidden_states.device)
causal_mask = self._prepare_attn_mask(
attention_mask,
input_shape=(batch_size, seq_length),
past_key_values_length=past_key_values_length,
)
for i, (block, layer_past) in enumerate(zip(self.h, past_key_values)):
if output_hidden_states:
all_hidden_states = all_hidden_states + (hidden_states,)
if self.gradient_checkpointing and self.training:
def create_custom_forward(module):
def custom_forward(*inputs):
# None for past_key_value
return module(*inputs, use_cache=use_cache, output_attentions=output_attentions)
return custom_forward
outputs = torch.utils.checkpoint.checkpoint(
create_custom_forward(block),
hidden_states,
causal_mask,
layer_past,
)
else:
outputs = block(
hidden_states,
layer_past=layer_past,
attention_mask=causal_mask,
use_cache=use_cache,
output_attentions=output_attentions,
)
hidden_states = outputs[0]
if use_cache is True:
presents = presents + (outputs[1],)
if output_attentions:
all_self_attentions = all_self_attentions + (outputs[2 if use_cache else 1],)
hidden_states = self.ln_f(hidden_states)
if output_hidden_states:
all_hidden_states = all_hidden_states + (hidden_states,)
if not return_dict:
return tuple(v for v in [hidden_states, presents, all_hidden_states, all_self_attentions] if v is not None)
return BaseModelOutputWithPastAndCrossAttentions(
last_hidden_state=hidden_states,
past_key_values=presents,
hidden_states=all_hidden_states,
attentions=all_self_attentions,
)
class TelechatForCausalLM(TelechatPreTrainedModel):
# _tied_weights_keys = ["lm_head.weight"]
_keys_to_ignore_on_load_missing = [r"lm_head.weight"]
def __init__(self, config: TelechatConfig):
super().__init__(config)
self.transformer = TelechatModel(config)
self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
self.post_init()
def get_output_embeddings(self):
return self.lm_head
def set_output_embeddings(self, new_embeddings: torch.Tensor):
self.lm_head = new_embeddings
def prepare_inputs_for_generation(
self,
input_ids: torch.LongTensor,
past_key_values: Optional[torch.Tensor] = None,
attention_mask: Optional[torch.Tensor] = None,
inputs_embeds: Optional[torch.Tensor] = None,
**kwargs,
) -> dict:
if past_key_values:
input_ids = input_ids[:, -1].unsqueeze(-1)
if inputs_embeds is not None and past_key_values is None:
model_inputs = {"inputs_embeds": inputs_embeds}
else:
model_inputs = {"input_ids": input_ids}
model_inputs.update(
{
"past_key_values": past_key_values,
"use_cache": kwargs.get("use_cache"),
"attention_mask": attention_mask,
}
)
return model_inputs
def forward(
self,
input_ids: Optional[torch.LongTensor] = None,
past_key_values: Optional[Tuple[Tuple[torch.Tensor, torch.Tensor], ...]] = None,
attention_mask: Optional[torch.Tensor] = None,
inputs_embeds: Optional[torch.Tensor] = None,
labels: Optional[torch.Tensor] = None,
use_cache: Optional[bool] = None,
output_attentions: Optional[bool] = None,
output_hidden_states: Optional[bool] = None,
return_dict: Optional[bool] = None,
**deprecated_arguments,
) -> Union[Tuple[torch.Tensor], CausalLMOutputWithCrossAttentions]:
return_dict = return_dict if return_dict is not None else self.config.use_return_dict
transformer_outputs = self.transformer(
input_ids,
past_key_values=past_key_values,
attention_mask=attention_mask,
inputs_embeds=inputs_embeds,
use_cache=use_cache,
output_attentions=output_attentions,
output_hidden_states=output_hidden_states,
return_dict=return_dict,
)
hidden_states = transformer_outputs[0]
lm_logits = self.lm_head(hidden_states)
loss = None
if labels is not None:
labels = labels.to(lm_logits.device)
shift_logits = lm_logits[..., :-1, :].contiguous()
shift_labels = labels[..., 1:].contiguous()
batch_size, seq_length, vocab_size = shift_logits.shape
loss_fct = CrossEntropyLoss()
loss = loss_fct(
shift_logits.view(batch_size * seq_length, vocab_size), shift_labels.view(batch_size * seq_length)
)
if not return_dict:
output = (lm_logits,) + transformer_outputs[1:]
return ((loss,) + output) if loss is not None else output
return CausalLMOutputWithCrossAttentions(
loss=loss,
logits=lm_logits,
past_key_values=transformer_outputs.past_key_values,
hidden_states=transformer_outputs.hidden_states,
attentions=transformer_outputs.attentions,
)
def chat(self, tokenizer, question: str = '', history: Union[List[Dict], History] = None, stream: bool = False,
generation_config: Optional[GenerationConfig] = None, **kwargs):
"""
Args:
tokenizer: the tokenizer of telechat
question: question which the model reply in this turn
history: history which will format the input for telechat
stream: if return the full text at last or yield the text in token
generation_config: configuration for generation
**kwargs: args which will update the generation config or pass to model forward
"""
generation_config = generation_config or self.generation_config
if not generation_config:
logger.error("generation_config is None")
raise ValueError("generation_config must not be None")
if not question:
logger.error("question is empty")
raise ValueError("question must not be empty")
if history is None:
history = []
# we update and check generate_config here for building inputs.
generation_config = copy.deepcopy(generation_config)
user_id = generation_config.user_token_id
bot_id = generation_config.bot_token_id
model_kwargs = generation_config.update(**kwargs)
generation_config.validate()
# transfer to History
if not isinstance(history, History):
history = History(tokenizer, history)
inputs = self.build_inputs_for_chat(tokenizer, question, history, generation_config, user_id, bot_id)
history.append({"role": "user", "content": question})
if stream:
streamer = TelechatIterTextStreamer(tokenizer, history,skip_prompt=True)
Thread(target=self.generate, kwargs=dict(
inputs=inputs.to(self.device), streamer=streamer,
generation_config=generation_config, **model_kwargs
)).start()
return streamer
else:
outputs = self.generate(inputs.to(self.device), generation_config=generation_config, **model_kwargs)
response = tokenizer.decode(outputs[0][len(inputs[0]):-1])
history.append({"role": "bot", "content": response})
return response, history
def build_inputs_for_chat(self, tokenizer, question, history, generation_config, usr_id, bot_id):
"""
check history and build inputs here
"""
# first tokenize question
q_token = tokenizer(question)
qa_history = copy.deepcopy(history)
# get the max length we should build our inputs in
model_max_length = self.config.seq_length
build_max_length = max(0, model_max_length - generation_config.max_new_tokens) \
if generation_config.max_new_tokens else max(0, generation_config.max_length)
if build_max_length < 3:
logger.warning("the model can not meet the requirements of input length,Please check config")
raise ValueError("")
# trunc left
input_tokens = [usr_id] + q_token["input_ids"][-build_max_length + 1:] + [bot_id]
length = len(input_tokens)
while len(qa_history) != 0:
message = qa_history.pop()
if message["role"] == "user":
tokens = [usr_id] + message["input_ids"]
elif message["role"] == "bot":
tokens = [bot_id] + message["input_ids"] + [generation_config.eos_token_id]
else:
tokens = []
if len(tokens) + length >= build_max_length:
break
else:
input_tokens = tokens + input_tokens
return torch.tensor([input_tokens], dtype=torch.int64)
models/12B_4bit/quantize_config.json 0 → 100644
View file @ f7db21eb
{
"bits": 4,
"group_size": 128,
"damp_percent": 0.01,
"desc_act": false,
"sym": true,
"true_sequential": true,
"model_name_or_path": null,
"model_file_base_name": null
}
\ No newline at end of file
models/12B_4bit/specail_tokens_map.json 0 → 100644
View file @ f7db21eb
{
"bos_token": {
"content": "<_start>",
"lstrip": false,
"normalized": true,
"rstrip": false,
"single_word": false
},
"eos_token": {
"content": "<_end>",
"lstrip": false,
"normalized": true,
"rstrip": false,
"single_word": false
},
"pad_token": {
"content": "<_pad>",
"lstrip": false,
"normalized": true,
"rstrip": false,
"single_word": false
},
"unk_token": {
"content": "<_unk>",
"lstrip": false,
"normalized": true,
"rstrip": false,
"single_word": false
}
}
\ No newline at end of file
models/12B_4bit/tokenization_telechat.py 0 → 100644
View file @ f7db21eb
import os
from shutil import copyfile
from typing import Any, Dict, List, Optional, Tuple
import sentencepiece as spm
from transformers.tokenization_utils import AddedToken, PreTrainedTokenizer
from transformers.utils import logging
logger = logging.get_logger(__name__)
VOCAB_FILES_NAMES = {"vocab_file": "tokenizer.model"}
# TODO: when we get download url from huggingface, refresh the map
PRETRAINED_VOCAB_FILES_MAP = {
"vocab_file": {},
"tokenizer_file": {},
}
class TelechatTokenizer(PreTrainedTokenizer):
vocab_files_names = VOCAB_FILES_NAMES
pretrained_vocab_files_map = PRETRAINED_VOCAB_FILES_MAP
model_input_names = ["input_ids", "attention_mask"]
def __init__(
self,
vocab_file,
unk_token="<unk>",
bos_token="<_start>",
eos_token="<_end>",
pad_token="<_pad>",
sp_model_kwargs: Optional[Dict[str, Any]] = None,
add_bos_token=True,
add_eos_token=False,
clean_up_tokenization_spaces=False,
**kwargs,
):
self.sp_model_kwargs = {} if sp_model_kwargs is None else sp_model_kwargs
bos_token = AddedToken(bos_token, lstrip=False, rstrip=False) if isinstance(bos_token, str) else bos_token
eos_token = AddedToken(eos_token, lstrip=False, rstrip=False) if isinstance(eos_token, str) else eos_token
unk_token = AddedToken(unk_token, lstrip=False, rstrip=False) if isinstance(unk_token, str) else unk_token
pad_token = AddedToken(pad_token, lstrip=False, rstrip=False) if isinstance(pad_token, str) else pad_token
self.sp_model = spm.SentencePieceProcessor(**self.sp_model_kwargs)
self.sp_model.Load(vocab_file)
super().__init__(
bos_token=bos_token,
eos_token=eos_token,
unk_token=unk_token,
pad_token=pad_token,
add_bos_token=add_bos_token,
add_eos_token=add_eos_token,
sp_model_kwargs=self.sp_model_kwargs,
clean_up_tokenization_spaces=clean_up_tokenization_spaces,
**kwargs,
)
self.vocab_file = vocab_file
self.add_bos_token = add_bos_token
self.add_eos_token = add_eos_token
def __getstate__(self):
state = self.__dict__.copy()
state["sp_model"] = None
return state
def __setstate__(self, d):
self.__dict__ = d
self.sp_model = spm.SentencePieceProcessor(**self.sp_model_kwargs)
self.sp_model.Load(self.vocab_file)
@property
def vocab_size(self):
"""Returns vocab size"""
return self.sp_model.get_piece_size()
def get_vocab(self):
"""Returns vocab as a dict"""
vocab = {self.convert_ids_to_tokens(i): i for i in range(self.vocab_size)}
vocab.update(self.added_tokens_encoder)
return vocab
def _tokenize(self, text):
"""Returns a tokenized string."""
return self.sp_model.encode(text, out_type=str)
def _convert_token_to_id(self, token):
"""Converts a token (str) in an id using the vocab."""
return self.sp_model.piece_to_id(token)
def _convert_id_to_token(self, index):
"""Converts an index (integer) in a token (str) using the vocab."""
token = self.sp_model.IdToPiece(index)
return token
def convert_tokens_to_string(self, tokens):
"""Converts a sequence of tokens (string) in a single string."""
current_sub_tokens = []
out_string = ""
prev_is_special = False
for i, token in enumerate(tokens):
# make sure that special tokens are not decoded using sentencepiece model
if token in self.all_special_tokens:
if not prev_is_special and i != 0:
out_string += " "
out_string += self.sp_model.decode(current_sub_tokens) + token
prev_is_special = True
current_sub_tokens = []
else:
current_sub_tokens.append(token)
prev_is_special = False
out_string += self.sp_model.decode(current_sub_tokens)
return out_string
def save_vocabulary(self, save_directory, filename_prefix: Optional[str] = None) -> Tuple[str]:
"""
Save the vocabulary and special tokens file to a directory.
Args:
save_directory (`str`):
The directory in which to save the vocabulary.
Returns:
`Tuple(str)`: Paths to the files saved.
"""
if not os.path.isdir(save_directory):
logger.error(f"Vocabulary path ({save_directory}) should be a directory")
return
out_vocab_file = os.path.join(
save_directory, (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["vocab_file"]
)
if os.path.abspath(self.vocab_file) != os.path.abspath(out_vocab_file) and os.path.isfile(self.vocab_file):
copyfile(self.vocab_file, out_vocab_file)
elif not os.path.isfile(self.vocab_file):
with open(out_vocab_file, "wb") as fi:
content_spiece_model = self.sp_model.serialized_model_proto()
fi.write(content_spiece_model)
return (out_vocab_file,)
def build_inputs_with_special_tokens(self, token_ids_0, token_ids_1=None):
bos_token_id = [self.bos_token_id] if self.add_bos_token else []
eos_token_id = [self.eos_token_id] if self.add_eos_token else []
output = bos_token_id + token_ids_0 + eos_token_id
if token_ids_1 is not None:
output = output + bos_token_id + token_ids_1 + eos_token_id
return output
def get_special_tokens_mask(
self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None, already_has_special_tokens: bool = False
) -> List[int]:
"""
Retrieve sequence ids from a token list that has no special tokens added. This method is called when adding
special tokens using the tokenizer `prepare_for_model` method.
Args:
token_ids_0 (`List[int]`):
List of IDs.
token_ids_1 (`List[int]`, *optional*):
Optional second list of IDs for sequence pairs.
already_has_special_tokens (`bool`, *optional*, defaults to `False`):
Whether or not the token list is already formatted with special tokens for the model.
Returns:
`List[int]`: A list of integers in the range [0, 1]: 1 for a special token, 0 for a sequence token.
"""
if already_has_special_tokens:
return super().get_special_tokens_mask(
token_ids_0=token_ids_0, token_ids_1=token_ids_1, already_has_special_tokens=True
)
bos_token_id = [1] if self.add_bos_token else []
eos_token_id = [1] if self.add_eos_token else []
if token_ids_1 is None:
return bos_token_id + ([0] * len(token_ids_0)) + eos_token_id
return (
bos_token_id
+ ([0] * len(token_ids_0))
+ eos_token_id
+ bos_token_id
+ ([0] * len(token_ids_1))
+ eos_token_id
)
def create_token_type_ids_from_sequences(
self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None
) -> List[int]:
"""
Creates a mask from the two sequences passed to be used in a sequence-pair classification task. An ALBERT
sequence pair mask has the following format:
```
0 0 0 0 0 0 0 0 0 0 0 1 1 1 1 1 1 1 1 1
| first sequence | second sequence |
```
if token_ids_1 is None, only returns the first portion of the mask (0s).
Args:
token_ids_0 (`List[int]`):
List of ids.
token_ids_1 (`List[int]`, *optional*):
Optional second list of IDs for sequence pairs.
Returns:
`List[int]`: List of [token type IDs](../glossary#token-type-ids) according to the given sequence(s).
"""
bos_token_id = [self.bos_token_id] if self.add_bos_token else []
eos_token_id = [self.eos_token_id] if self.add_eos_token else []
output = [0] * len(bos_token_id + token_ids_0 + eos_token_id)
if token_ids_1 is not None:
output += [1] * len(bos_token_id + token_ids_1 + eos_token_id)
return output
models/12B_4bit/tokenizer_config.json 0 → 100644
View file @ f7db21eb
{
"name_or_path": "ChinaTelecom/telechat-12b",
"tokenizer_class": "TelechatTokenizer",
"auto_map": {
"AutoTokenizer": [
"tokenization_telechat.TelechatTokenizer",
null
]
},
"add_bos_token": false,
"add_eos_token": false,
"use_fast": false,
"clean_up_tokenization_spaces": false,
"eos_token": {
"__type": "AddedToken",
"content": "<_end>",
"lstrip": false,
"normalized": true,
"rstrip": false,
"single_word": true
},
"model_max_length": 100000000,
"sp_model_kwargs": {},
"pad_token": {
"__type": "AddedToken",
"content": "<_pad>",
"lstrip": false,
"normalized": true,
"rstrip": false,
"single_word": true
},
"unk_token": {
"__type": "AddedToken",
"content": "<_end>",
"lstrip": false,
"normalized": true,
"rstrip": false,
"single_word": true
}
}
models/12B_8bit/config.json 0 → 100644
View file @ f7db21eb
{
"apply_residual_connection_post_layernorm": false,
"architectures": [
"TelechatForCausalLM"
],
"auto_map": {
"AutoConfig": "configuration_telechat.TelechatConfig",
"AutoModelForCausalLM": "modeling_telechat.TelechatForCausalLM"
},
"attention_dropout": 0.0,
"attention_softmax_in_fp32": true,
"bias_dropout_fusion": true,
"bos_token_id": 1,
"eos_token_id": 2,
"hidden_dropout": 0.0,
"hidden_size": 5120,
"initializer_range": 0.02,
"layer_norm_epsilon": 1e-05,
"masked_softmax_fusion": true,
"model_type": "telechat",
"n_head": 32,
"n_inner": null,
"n_layer": 38,
"offset_alibi": 100,
"pad_token_id": 3,
"pretraining_tp": 2,
"seq_length": 8192,
"skip_bias_add": true,
"skip_bias_add_qkv": false,
"slow_but_exact": false,
"transformers_version": "4.24.0",
"unk_token_id": 0,
"use_cache": true,
"vocab_size": 120000,
"ffn_hidden_size": 12288,
"flash_attn":true,
"tie_word_embeddings":false,
"training_seqlen":8192,
"logn":false,
"semi_causal":false,
"embed_layernorm":false
}
models/12B_8bit/configuration_telechat.py 0 → 100644
View file @ f7db21eb
# coding=utf-8
# Copyright 2022 the Big Science Workshop and HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
""" Telechat configuration"""
from packaging import version
from collections import OrderedDict
from transformers.utils import is_torch_available, logging
from transformers.configuration_utils import PretrainedConfig
from typing import TYPE_CHECKING, Any, List, Mapping, Optional
logger = logging.get_logger(__name__)
class TelechatConfig(PretrainedConfig):
"""
Args:
vocab_size (`int`, *optional*, defaults to 160256): Vocabulary size of the Telechat model.
hidden_size (`int`, *optional*, defaults to 4096): Dimensionality of the embeddings and hidden states.
ffn_hidden_size (`int`, *optional*, defaults to 12288): Dimensionality of the feed-forward hidden states.
n_layer (`int`, *optional*, defaults to 30): Number of hidden layers in the Transformer
n_head (`int`, *optional*, defaults to 32): Number of attention heads for each attention layer.
layer_norm_epsilon (`float`, *optional*, defaults to 1e-5): The epsilon to use in the layer normalization layers.
initializer_range (`float`, *optional*, defaults to 0.02): The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
apply_residual_connection_post_layernorm (`bool`, *optional*, defaults to `False`): If enabled, use the layer norm of the hidden states as the residual in the transformer blocks
hidden_dropout (`float`, *optional*, defaults to 0.0): Dropout rate of the dropout function on the bias dropout.
attention_dropout (`float`, *optional*, defaults to 0.0): Dropout rate applied to the attention probs
use_cache (`bool`, *optional*, defaults to `True`): Whether or not the model should return the last key/values attentions.
training_seqlen (`int`, *optional*, defaults to 8192): Sequence length during last finetuning.
logn (`bool`, *optional*, defaults to `True`): Whether or not to use logN during extrapolation.
embed_layernorm (`bool`, *optional*, defaults to `True`): Whether or not to use embedding layernorm.
"""
model_type = "telechat"
keys_to_ignore_at_inference = ["past_key_values"]
attribute_map = {
"num_hidden_layers": "n_layer",
"num_attention_heads": "n_head",
}
def __init__(
self,
vocab_size=160256,
hidden_size=4096,
n_layer=30,
n_head=32,
layer_norm_epsilon=1e-5,
initializer_range=0.02,
use_cache=True,
bos_token_id=1,
eos_token_id=2,
apply_residual_connection_post_layernorm=False,
hidden_dropout=0.0,
attention_dropout=0.0,
ffn_hidden_size=12288,
training_seqlen = 8192,
logn = True,
embed_layernorm = False,
**kwargs,
):
self.vocab_size = vocab_size
n_embed = kwargs.pop("n_embed", None)
self.hidden_size = hidden_size if n_embed is None else n_embed
self.n_layer = n_layer
self.n_head = n_head
self.layer_norm_epsilon = layer_norm_epsilon
self.initializer_range = initializer_range
self.use_cache = use_cache
self.apply_residual_connection_post_layernorm = apply_residual_connection_post_layernorm
self.hidden_dropout = hidden_dropout
self.attention_dropout = attention_dropout
self.bos_token_id = bos_token_id
self.eos_token_id = eos_token_id
self.logn = logn
self.ffn_hidden_size = ffn_hidden_size
self.training_seqlen = training_seqlen
self.embed_layernorm = embed_layernorm
super().__init__(bos_token_id=bos_token_id, eos_token_id=eos_token_id, **kwargs)
models/12B_8bit/generation_config.json 0 → 100644
View file @ f7db21eb
{
"max_length": 8192,
"do_sample": false,
"use_cache": true,
"temperature": 0.3,
"top_k": 5,
"top_p": 0.85,
"repetition_penalty": 1.01,
"pad_token_id": 3,
"bos_token_id": 1,
"eos_token_id": 2,
"user_token_id": 20,
"bot_token_id": 21
}
models/12B_8bit/generation_utils.py 0 → 100644
View file @ f7db21eb
from typing import Optional
from collections import deque
from queue import Queue
import copy
class History:
def __init__(self, tokenizer, history):
'''
init from a list of dict
'''
# use deque to meet some special situation
self.input_history = deque()
self.tokenizer = tokenizer
if history:
self._transfer_from_list(history)
def _transfer_from_list(self, history):
for message in history:
content = message.get("content")
# the token result may not be equal to the result model gen
message.update(self.tokenizer(content))
self.input_history.append(message)
def append(self, message):
content = message.get("content")
if "input_ids" not in message or "attention_mask" not in message:
message.update(self.tokenizer(content))
self.input_history.append(message)
def append_left(self, message):
content = message.get("content")
if "input_ids" not in message or "attention_mask" not in message:
message.update(self.tokenizer(content))
self.input_history.appendleft(message)
def pop(self):
x = self.input_history.pop()
return x
def pop_left(self):
x = self.pop_left()
return x
def update(self, message):
self.input_history.pop()
self.append(message)
def __len__(self):
return self.input_history.__len__()
def __str__(self):
return self.input_history.__str__()
def __copy__(self):
new_instance = type(self)(self.tokenizer, [])
new_instance.input_history = copy.copy(self.input_history)
return new_instance
def __deepcopy__(self, memodict={}):
new_instance = type(self)(self.tokenizer, [])
new_instance.input_history = copy.deepcopy(self.input_history)
return new_instance
class TelechatIterTextStreamer:
"""
With reference to the TextIterStreamers in transformers, we have rewritten this class
"""
def __init__(
self, tokenizer, history: History = None, skip_prompt: bool = False, timeout: Optional[float] = None,
**decode_kwargs
):
self.tokenizer = tokenizer
self.history = history
self.skip_prompt = skip_prompt
self.timeout = timeout
self.decode_kwargs = decode_kwargs
self.text_queue = Queue()
self.cache_time = 0
self.text_until = ""
self.token_until = []
self.stop_signal = None
self.next_tokens_are_prompt = True
self.history.append({"role": "bot", "content": self.text_until})
def put(self, value):
"""
put printable text into queue
"""
if len(value.shape) > 1 and value.shape[0] > 1:
raise ValueError("TextStreamer only supports batch size 1")
elif len(value.shape) > 1:
value = value[0]
if self.skip_prompt and self.next_tokens_are_prompt:
self.next_tokens_are_prompt = False
return
if value[-1] == self.tokenizer.eos_token_id:
return
# there may be some smart way to decode.
self.token_until.extend(value.tolist())
text = self.tokenizer.decode(self.token_until, **self.decode_kwargs)
if self._is_printable(text) or self.cache_time >= 6:
output_text = text[len(self.text_until):]
self.text_until = text
else:
self.cache_time+=1
return
self.on_finalized_text(output_text)
def end(self):
"""Flushes any remaining cache and prints a newline to stdout."""
# Flush the cache, if it exists
text = self.tokenizer.decode(self.token_until, **self.decode_kwargs)
output_text = text[len(self.text_until):]
self.text_until = text
self.on_finalized_text(output_text, stream_end=True)
self.clear_cache()
def clear_cache(self):
self.cache_time = 0
self.token_until = []
self.text_until = ""
self.history = None
self.next_tokens_are_prompt = True
def on_finalized_text(self, text: str, stream_end: bool = False):
"""Put the text tuple in the queue."""
self.history.update({"role": "bot", "content": self.text_until, "input_ids": self.token_until,
"attention_mask": [1] * len(self.token_until)})
self.text_queue.put((text, self.history), timeout=self.timeout)
if stream_end:
self.text_queue.put((self.stop_signal, self.history), timeout=self.timeout)
@staticmethod
def _is_printable(cp):
"""Checks whether tokens can be decoded or not"""
if "�" in cp:
return False
return True
def __iter__(self):
return self
def __next__(self):
value_now, history_until = self.text_queue.get(timeout=self.timeout)
if value_now == self.stop_signal:
raise StopIteration()
else:
return value_now, history_until
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment