commit 93ebdd82f7dfab742f6ef872ea195bf75611ffbf Author: ModelHub XC Date: Tue Jul 21 16:28:07 2026 +0800 初始化项目，由ModelHub XC社区提供模型 Model: katuni4ka/tiny-random-baichuan2-13b Source: Original Platform diff --git a/.gitattributes b/.gitattributes new file mode 100644 index 0000000..53d7257 --- /dev/null +++ b/.gitattributes @@ -0,0 +1,47 @@ +*.7z filter=lfs diff=lfs merge=lfs -text +*.arrow filter=lfs diff=lfs merge=lfs -text +*.bin filter=lfs diff=lfs merge=lfs -text +*.bin.* filter=lfs diff=lfs merge=lfs -text +*.bz2 filter=lfs diff=lfs merge=lfs -text +*.ftz filter=lfs diff=lfs merge=lfs -text +*.gz filter=lfs diff=lfs merge=lfs -text +*.h5 filter=lfs diff=lfs merge=lfs -text +*.joblib filter=lfs diff=lfs merge=lfs -text +*.lfs.* filter=lfs diff=lfs merge=lfs -text +*.model filter=lfs diff=lfs merge=lfs -text +*.msgpack filter=lfs diff=lfs merge=lfs -text +*.onnx filter=lfs diff=lfs merge=lfs -text +*.ot filter=lfs diff=lfs merge=lfs -text +*.parquet filter=lfs diff=lfs merge=lfs -text +*.pb filter=lfs diff=lfs merge=lfs -text +*.pt filter=lfs diff=lfs merge=lfs -text +*.pth filter=lfs diff=lfs merge=lfs -text +*.rar filter=lfs diff=lfs merge=lfs -text +saved_model/**/* filter=lfs diff=lfs merge=lfs -text +*.tar.* filter=lfs diff=lfs merge=lfs -text +*.tflite filter=lfs diff=lfs merge=lfs -text +*.tgz filter=lfs diff=lfs merge=lfs -text +*.xz filter=lfs diff=lfs merge=lfs -text +*.zip filter=lfs diff=lfs merge=lfs -text +*.zstandard filter=lfs diff=lfs merge=lfs -text +*.tfevents* filter=lfs diff=lfs merge=lfs -text +*.db* filter=lfs diff=lfs merge=lfs -text +*.ark* filter=lfs diff=lfs merge=lfs -text +**/*ckpt*data* filter=lfs diff=lfs merge=lfs -text +**/*ckpt*.meta filter=lfs diff=lfs merge=lfs -text +**/*ckpt*.index filter=lfs diff=lfs merge=lfs -text +*.safetensors filter=lfs diff=lfs merge=lfs -text +*.ckpt filter=lfs diff=lfs merge=lfs -text +*.gguf* filter=lfs diff=lfs merge=lfs -text +*.ggml filter=lfs diff=lfs merge=lfs -text +*.llamafile* filter=lfs diff=lfs merge=lfs -text +*.pt2 filter=lfs diff=lfs merge=lfs -text +*.mlmodel filter=lfs diff=lfs merge=lfs -text +*.npy filter=lfs diff=lfs merge=lfs -text +*.npz filter=lfs diff=lfs merge=lfs -text +*.pickle filter=lfs diff=lfs merge=lfs -text +*.pkl filter=lfs diff=lfs merge=lfs -text +*.tar filter=lfs diff=lfs merge=lfs -text +*.wasm filter=lfs diff=lfs merge=lfs -text +*.zst filter=lfs diff=lfs merge=lfs -text +*tfevents* filter=lfs diff=lfs merge=lfs -text \ No newline at end of file diff --git a/README.md b/README.md new file mode 100644 index 0000000..5e78be2 --- /dev/null +++ b/README.md @@ -0,0 +1,47 @@ +--- +license: Apache License 2.0 + +#model-type: +##如 gpt、phi、llama、chatglm、baichuan 等 +#- gpt + +#domain: +##如 nlp、cv、audio、multi-modal +#- nlp + +#language: +##语言代码列表 https://help.aliyun.com/document_detail/215387.html?spm=a2c4g.11186623.0.0.9f8d7467kni6Aa +#- cn + +#metrics: +##如 CIDEr、Blue、ROUGE 等 +#- CIDEr + +#tags: +##各种自定义，包括 pretrained、fine-tuned、instruction-tuned、RL-tuned 等训练方法和其他 +#- pretrained + +#tools: +##如 vllm、fastchat、llamacpp、AdaSeq 等 +#- vllm +--- +### 当前模型的贡献者未提供更加详细的模型介绍。模型文件和权重，可浏览“模型文件”页面获取。 +#### 您可以通过如下git clone命令，或者ModelScope SDK来下载模型 + +SDK下载 +```bash +#安装ModelScope +pip install modelscope +``` +```python +#SDK模型下载 +from modelscope import snapshot_download +model_dir = snapshot_download('katuni4ka/tiny-random-baichuan2-13b') +``` +Git下载 +``` +#Git模型下载 +git clone https://www.modelscope.cn/katuni4ka/tiny-random-baichuan2-13b.git +``` + +

如果您是本模型的贡献者，我们邀请您根据模型贡献文档，及时完善模型卡片内容。

\ No newline at end of file diff --git a/config.json b/config.json new file mode 100644 index 0000000..ed015b1 --- /dev/null +++ b/config.json @@ -0,0 +1,30 @@ +{ + "_from_model_config": true, + "_name_or_path": "/home/ea/work/my_optimum_intel/optimum-intel/tiny-random-baichuan2-13b", + "architectures": [ + "BaichuanForCausalLM" + ], + "auto_map": { + "AutoConfig": "configuration_baichuan.BaichuanConfig", + "AutoModelForCausalLM": "modeling_baichuan.BaichuanForCausalLM" + }, + "bos_token_id": 1, + "eos_token_id": 2, + "hidden_act": "silu", + "hidden_size": 64, + "initializer_range": 0.02, + "intermediate_size": 214, + "model_max_length": 4096, + "model_type": "baichuan", + "num_attention_heads": 2, + "num_hidden_layers": 2, + "pad_token_id": 0, + "rms_norm_eps": 1e-06, + "tie_word_embeddings": false, + "tokenizer_class": "BaichuanTokenizer", + "torch_dtype": "float32", + "transformers_version": "4.39.3", + "use_cache": true, + "vocab_size": 125696, + "z_loss_weight": 0 +} diff --git a/configuration.json b/configuration.json new file mode 100644 index 0000000..bbeeda1 --- /dev/null +++ b/configuration.json @@ -0,0 +1 @@ +{"framework": "pytorch", "task": "text-generation", "allow_remote": true} \ No newline at end of file diff --git a/configuration_baichuan.py b/configuration_baichuan.py new file mode 100644 index 0000000..f781580 --- /dev/null +++ b/configuration_baichuan.py @@ -0,0 +1,48 @@ +# Copyright (c) 2023, Baichuan Intelligent Technology. All rights reserved. + +from transformers.configuration_utils import PretrainedConfig + + +class BaichuanConfig(PretrainedConfig): + model_type = "baichuan" + keys_to_ignore_at_inference = ["past_key_values"] + + def __init__( + self, + vocab_size=125696, + hidden_size=5120, + intermediate_size=13696, + num_hidden_layers=40, + num_attention_heads=40, + hidden_act="silu", + model_max_length=4096, + initializer_range=0.02, + rms_norm_eps=1e-6, + use_cache=True, + pad_token_id=0, + bos_token_id=1, + eos_token_id=2, + tie_word_embeddings=False, + gradient_checkpointing=False, + z_loss_weight=0, + **kwargs, + ): + self.vocab_size = vocab_size + self.model_max_length = model_max_length + self.hidden_size = hidden_size + self.intermediate_size = intermediate_size + self.num_hidden_layers = num_hidden_layers + self.num_attention_heads = num_attention_heads + self.hidden_act = hidden_act + self.initializer_range = initializer_range + self.rms_norm_eps = rms_norm_eps + self.use_cache = use_cache + self.z_loss_weight = z_loss_weight + self.gradient_checkpointing = (gradient_checkpointing,) + super().__init__( + pad_token_id=pad_token_id, + bos_token_id=bos_token_id, + eos_token_id=eos_token_id, + tie_word_embeddings=tie_word_embeddings, + **kwargs, + ) diff --git a/generation_config.json b/generation_config.json new file mode 100644 index 0000000..ad1f8dc --- /dev/null +++ b/generation_config.json @@ -0,0 +1,7 @@ +{ + "_from_model_config": true, + "bos_token_id": 1, + "eos_token_id": 2, + "pad_token_id": 0, + "transformers_version": "4.39.3" +} diff --git a/generation_utils.py b/generation_utils.py new file mode 100644 index 0000000..5771699 --- /dev/null +++ b/generation_utils.py @@ -0,0 +1,83 @@ +from typing import List +from queue import Queue + +import torch + + +def build_chat_input(model, tokenizer, messages: List[dict], max_new_tokens: int=0): + def _parse_messages(messages, split_role="user"): + system, rounds = "", [] + round = [] + for i, message in enumerate(messages): + if message["role"] == "system": + assert i == 0 + system = message["content"] + continue + if message["role"] == split_role and round: + rounds.append(round) + round = [] + round.append(message) + if round: + rounds.append(round) + return system, rounds + + max_new_tokens = max_new_tokens or model.generation_config.max_new_tokens + max_input_tokens = model.config.model_max_length - max_new_tokens + system, rounds = _parse_messages(messages, split_role="user") + system_tokens = tokenizer.encode(system) + max_history_tokens = max_input_tokens - len(system_tokens) + + history_tokens = [] + for round in rounds[::-1]: + round_tokens = [] + for message in round: + if message["role"] == "user": + round_tokens.append(model.generation_config.user_token_id) + else: + round_tokens.append(model.generation_config.assistant_token_id) + round_tokens.extend(tokenizer.encode(message["content"])) + if len(history_tokens) == 0 or len(history_tokens) + len(round_tokens) <= max_history_tokens: + history_tokens = round_tokens + history_tokens # concat left + if len(history_tokens) < max_history_tokens: + continue + break + + input_tokens = system_tokens + history_tokens + if messages[-1]["role"] != "assistant": + input_tokens.append(model.generation_config.assistant_token_id) + input_tokens = input_tokens[-max_input_tokens:] # truncate left + return torch.LongTensor([input_tokens]).to(model.device) + + +class TextIterStreamer: + def __init__(self, tokenizer, skip_prompt=False, skip_special_tokens=False): + self.tokenizer = tokenizer + self.skip_prompt = skip_prompt + self.skip_special_tokens = skip_special_tokens + self.tokens = [] + self.text_queue = Queue() + self.next_tokens_are_prompt = True + + def put(self, value): + if self.skip_prompt and self.next_tokens_are_prompt: + self.next_tokens_are_prompt = False + else: + if len(value.shape) > 1: + value = value[0] + self.tokens.extend(value.tolist()) + self.text_queue.put( + self.tokenizer.decode(self.tokens, skip_special_tokens=self.skip_special_tokens)) + + def end(self): + self.text_queue.put(None) + + def __iter__(self): + return self + + def __next__(self): + value = self.text_queue.get() + if value is None: + raise StopIteration() + else: + return value + diff --git a/handler.py b/handler.py new file mode 100644 index 0000000..9974200 --- /dev/null +++ b/handler.py @@ -0,0 +1,23 @@ +import torch +from typing import Dict, List, Any +from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline +from transformers.generation.utils import GenerationConfig + +# get dtype +dtype = torch.bfloat16 if torch.cuda.get_device_capability()[0] == 8 else torch.float16 + +class EndpointHandler: + def __init__(self, path=""): + # load the model + self.model = AutoModelForCausalLM.from_pretrained(path, device_map="auto", torch_dtype=dtype, trust_remote_code=True) + self.model.generation_config = GenerationConfig.from_pretrained(path) + self.tokenizer = AutoTokenizer.from_pretrained(path, use_fast=False, trust_remote_code=True) + + def __call__(self, data: Any) -> List[List[Dict[str, float]]]: + inputs = data.pop("inputs", data) + # ignoring parameters! Default to configs in generation_config.json. + messages = [{"role": "user", "content": inputs}] + response = self.model.chat(self.tokenizer, messages) + if torch.backends.mps.is_available(): + torch.mps.empty_cache() + return [{'generated_text': response}] diff --git a/model.safetensors b/model.safetensors new file mode 100644 index 0000000..b83c2df --- /dev/null +++ b/model.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:3f4599a2f6cb6992c94e1aa62193985b5745879672f934debea6a8e39d96b8d6 +size 64819216 diff --git a/modeling_baichuan.py b/modeling_baichuan.py new file mode 100644 index 0000000..6c80b44 --- /dev/null +++ b/modeling_baichuan.py @@ -0,0 +1,827 @@ +# Copyright (c) 2023, Baichuan Intelligent Technology. All rights reserved. + +from .configuration_baichuan import BaichuanConfig +from .generation_utils import build_chat_input, TextIterStreamer + +import math +from threading import Thread +from typing import List, Optional, Tuple, Union + +import torch +from torch import nn +from torch.nn import CrossEntropyLoss +from torch.nn import functional as F +from transformers import PreTrainedModel, PretrainedConfig +from transformers.activations import ACT2FN +from transformers.generation.utils import GenerationConfig +from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast +from transformers.utils import logging, ContextManagers + +import os +from contextlib import contextmanager +from accelerate import init_empty_weights + +logger = logging.get_logger(__name__) + +try: + from xformers import ops as xops +except ImportError: + xops = None + logger.warning( + "Xformers is not installed correctly. If you want to use memory_efficient_attention to accelerate training use the following command to install Xformers\npip install xformers." + ) + + +def _get_interleave(n): + def _get_interleave_power_of_2(n): + start = 2 ** (-(2 ** -(math.log2(n) - 3))) + ratio = start + return [start * ratio**i for i in range(n)] + + if math.log2(n).is_integer(): + return _get_interleave_power_of_2(n) + else: + closest_power_of_2 = 2 ** math.floor(math.log2(n)) + return ( + _get_interleave_power_of_2(closest_power_of_2) + + _get_interleave(2 * closest_power_of_2)[0::2][: n - closest_power_of_2] + ) + + +def _fill_with_neg_inf(t): + """FP16-compatible function that fills a tensor with -inf.""" + return t.float().fill_(float("-inf")).type_as(t) + + +def _buffered_future_mask(tensor, maxpos, alibi, attn_heads): + _future_mask = torch.triu(_fill_with_neg_inf(torch.zeros([maxpos, maxpos])), 1) + _future_mask = _future_mask.unsqueeze(0) + alibi + new_future_mask = _future_mask.to(tensor) + return new_future_mask[: tensor.shape[0] * attn_heads, :maxpos, :maxpos] + + +def _gen_alibi_mask(tensor, n_head, max_pos): + slopes = torch.Tensor(_get_interleave(n_head)) + position_point = torch.arange(max_pos) - max_pos + 1 + position_point = position_point.unsqueeze(0).unsqueeze(0).expand(n_head, -1, -1) + diag = torch.diag(position_point[0]) + position_point = position_point - diag.unsqueeze(0).unsqueeze(0).transpose(-1, -2) + alibi = slopes.unsqueeze(1).unsqueeze(1) * position_point + alibi = alibi.view(n_head, 1, max_pos) + alibi_mask = torch.triu(_fill_with_neg_inf(torch.zeros([max_pos, max_pos])), 1) + alibi_mask = alibi_mask.unsqueeze(0) + alibi + return alibi_mask + + +class RMSNorm(torch.nn.Module): + def __init__(self, hidden_size, epsilon=1e-6): + super().__init__() + self.weight = torch.nn.Parameter(torch.empty(hidden_size)) + self.epsilon = epsilon + + def forward(self, hidden_states): + variance = hidden_states.to(torch.float32).pow(2).mean(-1, keepdim=True) + hidden_states = hidden_states * torch.rsqrt(variance + self.epsilon) + + # convert into half-precision + if self.weight.dtype in [torch.float16, torch.bfloat16]: + hidden_states = hidden_states.to(self.weight.dtype) + + return self.weight * hidden_states + + +class MLP(torch.nn.Module): + def __init__( + self, + hidden_size: int, + intermediate_size: int, + hidden_act: str, + ): + super().__init__() + self.gate_proj = torch.nn.Linear(hidden_size, intermediate_size, bias=False) + self.down_proj = torch.nn.Linear(intermediate_size, hidden_size, bias=False) + self.up_proj = torch.nn.Linear(hidden_size, intermediate_size, bias=False) + self.act_fn = ACT2FN[hidden_act] + + def forward(self, x): + return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)) + + +class BaichuanAttention(torch.nn.Module): + def __init__(self, config: BaichuanConfig): + super().__init__() + self.config = config + self.hidden_size = config.hidden_size + self.num_heads = config.num_attention_heads + self.head_dim = self.hidden_size // self.num_heads + self.max_position_embeddings = config.model_max_length + + if (self.head_dim * self.num_heads) != self.hidden_size: + raise ValueError( + f"hidden_size {self.hidden_size} is not divisible by num_heads {self.num_heads}" + ) + self.W_pack = torch.nn.Linear( + self.hidden_size, 3 * self.hidden_size, bias=False + ) + self.o_proj = torch.nn.Linear( + self.num_heads * self.head_dim, self.hidden_size, bias=False + ) + + def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int): + return ( + tensor.view(bsz, seq_len, self.num_heads, self.head_dim) + .transpose(1, 2) + .contiguous() + ) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + past_key_value: Optional[Tuple[torch.Tensor]] = None, + output_attentions: bool = False, + use_cache: bool = False, + ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: + bsz, q_len, _ = hidden_states.size() + + proj = self.W_pack(hidden_states) + proj = ( + proj.unflatten(-1, (3, self.hidden_size)) + .unsqueeze(0) + .transpose(0, -2) + .squeeze(-2) + ) + query_states = ( + proj[0].view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2) + ) + key_states = ( + proj[1].view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2) + ) + value_states = ( + proj[2].view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2) + ) + + kv_seq_len = key_states.shape[-2] + if past_key_value is not None: + kv_seq_len += past_key_value[0].shape[-2] + + if past_key_value is not None: + # reuse k, v, self_attention + key_states = torch.cat([past_key_value[0], key_states], dim=2) + value_states = torch.cat([past_key_value[1], value_states], dim=2) + + past_key_value = (key_states, value_states) if use_cache else None + if xops is not None and self.training: + attn_weights = None + # query_states = query_states.transpose(1, 2) + # key_states = key_states.transpose(1, 2) + # value_states = value_states.transpose(1, 2) + # attn_output = xops.memory_efficient_attention( + # query_states, key_states, value_states, attn_bias=attention_mask + # ) + with torch.backends.cuda.sdp_kernel(enable_flash=True, enable_math=True, enable_mem_efficient=True): + attn_output = F.scaled_dot_product_attention(query_states, key_states, value_states, attn_mask = attention_mask) + attn_output = attn_output.transpose(1, 2) + else: + attn_weights = torch.matmul( + query_states, key_states.transpose(2, 3) + ) / math.sqrt(self.head_dim) + + if attention_mask is not None: + if q_len == 1: # inference with cache + if len(attention_mask.size()) == 4: + attention_mask = attention_mask[:, :, -1:, :] + else: + attention_mask = attention_mask[:, -1:, :] + attn_weights = attn_weights + attention_mask + attn_weights = torch.max( + attn_weights, torch.tensor(torch.finfo(attn_weights.dtype).min) + ) + + attn_weights = torch.nn.functional.softmax(attn_weights, dim=-1) + attn_output = torch.matmul(attn_weights, value_states) + + attn_output = attn_output.transpose(1, 2) + attn_output = attn_output.reshape(bsz, q_len, self.hidden_size) + attn_output = self.o_proj(attn_output) + + if not output_attentions: + attn_weights = None + + return attn_output, attn_weights, past_key_value + + +class BaichuanLayer(torch.nn.Module): + def __init__(self, config: BaichuanConfig): + super().__init__() + self.hidden_size = config.hidden_size + self.self_attn = BaichuanAttention(config=config) + self.mlp = MLP( + hidden_size=self.hidden_size, + intermediate_size=config.intermediate_size, + hidden_act=config.hidden_act, + ) + self.input_layernorm = RMSNorm(config.hidden_size, epsilon=config.rms_norm_eps) + self.post_attention_layernorm = RMSNorm( + config.hidden_size, epsilon=config.rms_norm_eps + ) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + past_key_value: Optional[Tuple[torch.Tensor]] = None, + output_attentions: Optional[bool] = False, + use_cache: Optional[bool] = False, + ) -> Tuple[ + torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]] + ]: + residual = hidden_states + + hidden_states = self.input_layernorm(hidden_states) + + # Self Attention + hidden_states, self_attn_weights, present_key_value = self.self_attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + past_key_value=past_key_value, + output_attentions=output_attentions, + use_cache=use_cache, + ) + hidden_states = residual + hidden_states + + # Fully Connected + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = self.mlp(hidden_states) + hidden_states = residual + hidden_states + + outputs = (hidden_states,) + + if use_cache: + outputs += (present_key_value,) + + return outputs + + +class BaichuanPreTrainedModel(PreTrainedModel): + config_class = BaichuanConfig + base_model_prefix = "model" + supports_gradient_checkpointing = True + _no_split_modules = ["BaichuanLayer"] + _keys_to_ignore_on_load_unexpected = [r"decoder\.version"] + + def _init_weights(self, module): + std = self.config.initializer_range + if isinstance(module, torch.nn.Linear): + module.weight.data.normal_(mean=0.0, std=std) + if module.bias is not None: + module.bias.data.zero_() + elif isinstance(module, torch.nn.Embedding): + module.weight.data.normal_(mean=0.0, std=std) + if module.padding_idx is not None: + module.weight.data[module.padding_idx].zero_() + + def _set_gradient_checkpointing(self, module, value=False): + if isinstance(module, BaichuanModel): + module.gradient_checkpointing = value + + +class BaichuanModel(BaichuanPreTrainedModel): + def __init__(self, config: BaichuanConfig): + super().__init__(config) + self.padding_idx = config.pad_token_id + self.vocab_size = config.vocab_size + self.n_head = config.num_attention_heads + self.embed_tokens = torch.nn.Embedding( + config.vocab_size, config.hidden_size, self.padding_idx + ) + self.layers = torch.nn.ModuleList( + [BaichuanLayer(config) for _ in range(config.num_hidden_layers)] + ) + self.norm = RMSNorm(config.hidden_size, epsilon=config.rms_norm_eps) + + self.gradient_checkpointing = config.gradient_checkpointing + self.post_init() + self.max_cache_pos = config.model_max_length + self.first_run = True + self.alibi_mask = None + + def get_input_embeddings(self): + return self.embed_tokens + + def set_input_embeddings(self, value): + self.embed_tokens = value + + def get_alibi_mask(self, tensor, seq_length_with_past): + if self.training: + slopes = torch.Tensor(_get_interleave(self.n_head)) + position_point = ( + torch.arange(seq_length_with_past) - seq_length_with_past + 1 + ) + position_point = ( + position_point.unsqueeze(0) + .unsqueeze(0) + .expand(self.n_head, seq_length_with_past, -1) + ) + diag = torch.diag(position_point[0]) + position_point = position_point - diag.unsqueeze(0).unsqueeze(0).transpose( + -1, -2 + ) + alibi = slopes.unsqueeze(1).unsqueeze(1) * position_point + mask = _buffered_future_mask( + tensor, seq_length_with_past, alibi, self.n_head + ) + else: + if self.first_run: + self.first_run = False + self.register_buffer( + "future_mask", + _gen_alibi_mask(tensor, self.n_head, self.max_cache_pos).to( + tensor + ), + persistent=False, + ) + if seq_length_with_past > self.max_cache_pos: + self.max_cache_pos = seq_length_with_past + self.register_buffer( + "future_mask", + _gen_alibi_mask(tensor, self.n_head, self.max_cache_pos).to( + tensor + ), + persistent=False, + ) + mask = self.future_mask[ + : self.n_head, :seq_length_with_past, :seq_length_with_past + ] + return mask + + def forward( + self, + input_ids: torch.LongTensor = None, + attention_mask: Optional[torch.Tensor] = None, + past_key_values: Optional[List[torch.FloatTensor]] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + use_cache: Optional[bool] = False, + output_attentions: Optional[bool] = False, + output_hidden_states: Optional[bool] = False, + return_dict: Optional[bool] = True, + ) -> Union[Tuple, BaseModelOutputWithPast]: + if input_ids is not None and inputs_embeds is not None: + raise ValueError( + "You cannot provide both input_ids and inputs_embeds simultaneously" + ) + elif 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 + else: + raise ValueError("You need to provide input_ids or inputs_embeds") + + return_dict = ( + return_dict if return_dict is not None else self.config.use_return_dict + ) + + seq_length_with_past = seq_length + + if past_key_values 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 inputs_embeds is None: + inputs_embeds = self.embed_tokens(input_ids) + + if self.training: + if ( + self.alibi_mask is None + or self.alibi_mask.shape[-1] != seq_length_with_past + ): + self.alibi_mask = self.get_alibi_mask( + inputs_embeds, seq_length_with_past + ) + alibi_mask = self.alibi_mask + else: + alibi_mask = self.get_alibi_mask(inputs_embeds, seq_length_with_past) + + if attention_mask is not None: + if len(attention_mask.shape) == 2: + expanded_mask = attention_mask.to(torch.float32) + expanded_mask = torch.tril( + torch.gt(expanded_mask[:, :, None] * expanded_mask[:, None, :], 0).to(torch.float32) + ) * torch.eq(expanded_mask[:, :, None] - expanded_mask[:, None, :], 0).to(torch.float32) + else: + expanded_mask = attention_mask + bsz = inputs_embeds.size(0) + src_len, tgt_len = alibi_mask.size()[-2:] + expanded_mask = ( + expanded_mask.unsqueeze(1) + .expand(bsz, 1, src_len, tgt_len) + .to(alibi_mask.dtype) + ) + inverted_mask = 1.0 - expanded_mask + inverted_mask = inverted_mask.masked_fill( + inverted_mask.to(torch.bool), torch.finfo(alibi_mask.dtype).min + ) + attention_mask = inverted_mask + alibi_mask.unsqueeze(0) + else: + attention_mask = alibi_mask + + hidden_states = inputs_embeds + + if self.gradient_checkpointing and self.training: + if use_cache: + logger.warning_once( + "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..." + ) + use_cache = False + + # decoder layers + all_hidden_states = () if output_hidden_states else None + all_self_attns = () if output_attentions else None + next_decoder_cache = () if use_cache else None + + for idx, decoder_layer in enumerate(self.layers): + if output_hidden_states: + all_hidden_states += (hidden_states,) + + past_key_value = ( + past_key_values[idx] if past_key_values is not None else None + ) + + if self.gradient_checkpointing and self.training: + + def create_custom_forward(module): + def custom_forward(*inputs): + # None for past_key_value + return module(*inputs, output_attentions, None) + + return custom_forward + + layer_outputs = torch.utils.checkpoint.checkpoint( + create_custom_forward(decoder_layer), + hidden_states, + attention_mask, + None, + ) + else: + layer_outputs = decoder_layer( + hidden_states, + attention_mask=attention_mask, + past_key_value=past_key_value, + output_attentions=output_attentions, + use_cache=use_cache, + ) + + hidden_states = layer_outputs[0] + + if use_cache: + next_decoder_cache += (layer_outputs[2 if output_attentions else 1],) + + if output_attentions: + all_self_attns += (layer_outputs[1],) + + hidden_states = self.norm(hidden_states) + + # add hidden states from the last decoder layer + if output_hidden_states: + all_hidden_states += (hidden_states,) + + next_cache = next_decoder_cache if use_cache else None + if not return_dict: + return tuple( + v + for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] + if v is not None + ) + return BaseModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=next_cache, + hidden_states=all_hidden_states, + attentions=all_self_attns, + ) + + +class NormHead(nn.Module): + def __init__(self, hidden_size, vocab_size, bias=False): + super().__init__() + self.weight = nn.Parameter(torch.empty((vocab_size, hidden_size))) + nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5)) + self.first_flag = True + + def forward(self, hidden_states): + if self.training: + norm_weight = nn.functional.normalize(self.weight) + self.first_flag = True + elif self.first_flag: + self.first_flag = False + self.weight.data = nn.functional.normalize(self.weight) + norm_weight = self.weight + else: + norm_weight = self.weight + return nn.functional.linear(hidden_states, norm_weight) + +_init_weights = True +@contextmanager +def no_init_weights(_enable=True): + global _init_weights + old_init_weights = _init_weights + if _enable: + _init_weights = False + try: + yield + finally: + _init_weights = old_init_weights + + +class BaichuanForCausalLM(BaichuanPreTrainedModel): + def __init__(self, config, *model_args, **model_kwargs): + super().__init__(config, *model_args, **model_kwargs) + self.model = BaichuanModel(config) + self.lm_head = NormHead(config.hidden_size, config.vocab_size, bias=False) + #if hasattr(config, "quantization_config") and config.quantization_config['load_in_4bit']: + if hasattr(config, "quantization_config") and isinstance(config.quantization_config, dict) and config.quantization_config.get('load_in_4bit', False): + try: + from .quantizer import quantize_offline, init_model_weight_int4 + except ImportError: + raise ImportError(f"Needs quantize_offline to run quantize.") + quantize_offline(self, 4) + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.model.embed_tokens + + def set_input_embeddings(self, value): + self.model.embed_tokens = value + + def get_output_embeddings(self): + return self.lm_head + + def set_output_embeddings(self, new_embeddings): + self.lm_head = new_embeddings + + def set_decoder(self, decoder): + self.model = decoder + + def get_decoder(self): + return self.model + + @classmethod + def from_pretrained( + cls, + pretrained_model_name_or_path: Optional[Union[str, os.PathLike]], + *model_args, + config: Optional[Union[PretrainedConfig, str, os.PathLike]] = None, + cache_dir: Optional[Union[str, os.PathLike]] = None, + ignore_mismatched_sizes: bool = False, + force_download: bool = False, + local_files_only: bool = False, + token: Optional[Union[str, bool]] = None, + revision: str = "main", + use_safetensors: bool = None, + **kwargs, + ): + + # Load config if we don't provide a configuration + if not isinstance(config, PretrainedConfig): + config_path = config if config is not None else pretrained_model_name_or_path + config, model_kwargs = cls.config_class.from_pretrained( + config_path, + cache_dir=cache_dir, + return_unused_kwargs=True, + force_download=force_download, + resume_download=False, + proxies=None, + local_files_only=local_files_only, + token=token, + revision=revision, + subfolder="", + _from_auto=False, + _from_pipeline=None, + **kwargs, + ) + else: + model_kwargs = kwargs + + if hasattr(config, "quantization_config") and config.quantization_config['load_in_4bit']: + try: + from .quantizer import init_model_weight_int4 + from accelerate import init_empty_weights, dispatch_model, infer_auto_device_map + from accelerate.utils import CustomDtype + from accelerate.utils import get_balanced_memory + except ImportError: + raise ImportError(f"Needs import model weight init func to run quantize.") + # Instantiate model. + init_contexts = [no_init_weights(_enable=True)] + init_contexts.append(init_empty_weights()) + with ContextManagers(init_contexts): + model = cls(config) + + model_file = os.path.join(pretrained_model_name_or_path, 'pytorch_model.bin') + state_dict = torch.load(model_file, map_location="cpu") + model.is_quantized = True + + device_map = kwargs.pop("device_map", None) + torch_dtype = kwargs.pop("torch_dtype", None) + if device_map is not None: + kwargs = {"no_split_module_classes": model._no_split_modules} + target_dtype = CustomDtype.INT4 + max_memory = get_balanced_memory( + model, + dtype=target_dtype, + low_zero=(device_map == "balanced_low_0"), + max_memory=None, + **kwargs, + ) + kwargs["max_memory"] = max_memory + device_map = infer_auto_device_map(model, dtype=target_dtype, **kwargs) + model = init_model_weight_int4(config, model, state_dict) + + # Set model in evaluation mode to deactivate DropOut modules by default + model.eval() + # If it is a model with generation capabilities, attempt to load the generation config + if model.can_generate(): + try: + model.generation_config = GenerationConfig.from_pretrained( + pretrained_model_name_or_path, + cache_dir=cache_dir, + force_download=force_download, + resume_download=False, + proxies=None, + local_files_only=local_files_only, + token=token, + revision=revision, + subfolder="", + _from_auto=False, + _from_pipeline=None, + **kwargs, + ) + except (OSError, TypeError): + logger.info( + "Generation config file not found, using a generation config created from the model config." + ) + pass + + if device_map is not None: + dispatch_model(model, device_map=device_map) + + return model + + return super(BaichuanForCausalLM, cls).from_pretrained(pretrained_model_name_or_path, *model_args, + config=config, cache_dir=cache_dir, ignore_mismatched_sizes=ignore_mismatched_sizes, + force_download=force_download, local_files_only=local_files_only, token=token, revision=revision, + use_safetensors=use_safetensors, **kwargs) + + def forward( + self, + input_ids: torch.LongTensor = None, + attention_mask: Optional[torch.Tensor] = None, + past_key_values: Optional[List[torch.FloatTensor]] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + labels: Optional[torch.LongTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = False, + output_hidden_states: Optional[bool] = False, + return_dict: Optional[bool] = True, + **kwargs, + ) -> Union[Tuple, CausalLMOutputWithPast]: + return_dict = ( + return_dict if return_dict is not None else self.config.use_return_dict + ) + + # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn) + outputs = self.model( + input_ids=input_ids, + attention_mask=attention_mask, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + hidden_states = outputs[0] + logits = self.lm_head(hidden_states) + loss = None + if labels is not None: + # Shift so that tokens < n predict n + shift_logits = logits[..., :-1, :].contiguous() + shift_labels = labels[..., 1:].contiguous() + # Flatten the tokens + loss_fct = CrossEntropyLoss() + shift_logits = shift_logits.view(-1, self.config.vocab_size) + shift_labels = shift_labels.view(-1) + softmax_normalizer = shift_logits.max(-1).values ** 2 + z_loss = self.config.z_loss_weight * softmax_normalizer.mean() + # Enable model parallelism + shift_labels = shift_labels.to(shift_logits.device) + loss = loss_fct(shift_logits, shift_labels) + z_loss + + if not return_dict: + output = (logits,) + outputs[1:] + return (loss,) + output if loss is not None else output + + return CausalLMOutputWithPast( + loss=loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + + def quantize(self, bits: int): + try: + from .quantizer import quantize_online + except ImportError: + raise ImportError(f"Needs QLinear to run quantize.") + return quantize_online(self, bits) + + 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, + ): + if past_key_values: + input_ids = input_ids[:, -1:] + + # if `inputs_embeds` are passed, we only want to use them in the 1st generation step + 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 + + @staticmethod + def _reorder_cache(past_key_values, beam_idx): + return tuple( + tuple(past_state.index_select(0, beam_idx) for past_state in layer_past) + for layer_past in past_key_values + ) + + def _build_chat_input( + self, tokenizer, messages: List[dict], max_new_tokens: int = 0 + ): + max_new_tokens = max_new_tokens or self.generation_config.max_new_tokens + max_input_tokens = self.config.model_max_length - max_new_tokens + max_input_tokens = max(self.config.model_max_length // 2, max_input_tokens) + total_input, round_input = [], [] + for i, message in enumerate(messages[::-1]): + content_tokens = tokenizer.encode(message["content"]) + if message["role"] == "user": + round_input = ( + [self.generation_config.user_token_id] + + content_tokens + + round_input + ) + if ( + total_input + and len(total_input) + len(round_input) > max_input_tokens + ): + break + else: + total_input = round_input + total_input + if len(total_input) >= max_input_tokens: + break + else: + round_input = [] + elif message["role"] == "assistant": + round_input = ( + [self.generation_config.assistant_token_id] + + content_tokens + + [self.generation_config.eos_token_id] + + round_input + ) + else: + raise ValueError(f"message role not supported yet: {message['role']}") + total_input = total_input[-max_input_tokens:] # truncate left + total_input.append(self.generation_config.assistant_token_id) + total_input = torch.LongTensor([total_input]).to(self.device) + return total_input + + def chat(self, tokenizer, messages: List[dict], stream=False, + generation_config: Optional[GenerationConfig]=None): + generation_config = generation_config or self.generation_config + input_ids = build_chat_input(self, tokenizer, messages, generation_config.max_new_tokens) + if stream: + streamer = TextIterStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True) + Thread(target=self.generate, kwargs=dict( + inputs=input_ids, streamer=streamer, + generation_config=generation_config, + )).start() + return streamer + else: + outputs = self.generate(input_ids, generation_config=generation_config) + response = tokenizer.decode(outputs[0][len(input_ids[0]):], skip_special_tokens=True) + return response diff --git a/quantizer.py b/quantizer.py new file mode 100644 index 0000000..cdd0409 --- /dev/null +++ b/quantizer.py @@ -0,0 +1,211 @@ +#import bitsandbytes as bnb +from accelerate import init_empty_weights +#from bitsandbytes.nn.modules import Params4bit, Int8Params +import torch + +def Params4bitCuda(self, device): + self.data = self.data.cuda(device) + self.quant_state[0] = self.quant_state[0].cuda(device) + self.quant_state[4][0] = self.quant_state[4][0].cuda(device) + self.quant_state[4][1][0] = self.quant_state[4][1][0].cuda(device) + self.quant_state[4][1][1] = self.quant_state[4][1][1].cuda(device) + + self.quant_state[6] = self.quant_state[6].cuda(device) + return self + +class Linear4bitOnline(torch.nn.Module): + def __init__(self, weight, bias, quant_type): + super().__init__() + self.weight = Params4bit( + weight.data, requires_grad=False, compress_statistics=True, quant_type=quant_type + ) + self.compute_dtype = None + #self.weight.cuda(weight.device) + self.bias = bias + + def forward(self, x: torch.Tensor): + # weights are cast automatically as Int8Params, but the bias has to be cast manually + if self.bias is not None and self.bias.dtype != x.dtype: + self.bias.data = self.bias.data.to(x.dtype) + + if getattr(self.weight, "quant_state", None) is None: + print( + "FP4 quantization state not initialized. Please call .cuda() or .to(device) on the LinearFP4 layer first." + ) + inp_dtype = x.dtype + if self.compute_dtype is not None: + x = x.to(self.compute_dtype) + + bias = None if self.bias is None else self.bias.to(self.compute_dtype) + out = bnb.matmul_4bit( + x, self.weight.t(), bias=bias, quant_state=self.weight.quant_state + ) + + out = out.to(inp_dtype) + + return out + +class Linear8bitLtOnline(torch.nn.Module): + def __init__( + self, + weight, + bias, + has_fp16_weights=True, + memory_efficient_backward=False, + threshold=0.0, + index=None, + ): + super().__init__() + assert ( + not memory_efficient_backward + ), "memory_efficient_backward is no longer required and the argument is deprecated in 0.37.0 and will be removed in 0.39.0" + self.state = bnb.MatmulLtState() + self.index = index + + # Necessary for stacked layers + self.state.threshold = threshold + self.state.has_fp16_weights = has_fp16_weights + self.state.memory_efficient_backward = memory_efficient_backward + if threshold > 0.0 and not has_fp16_weights: + self.state.use_pool = True + + self.weight = Int8Params( + weight.data, + has_fp16_weights=has_fp16_weights, + requires_grad=has_fp16_weights, + ) + self.bias = bias + + def init_8bit_state(self): + self.state.CB = self.weight.CB + self.state.SCB = self.weight.SCB + self.weight.CB = None + self.weight.SCB = None + + def forward(self, x: torch.Tensor): + self.state.is_training = self.training + if self.weight.CB is not None: + self.init_8bit_state() + + # weights are cast automatically as Int8Params, but the bias has to be cast manually + if self.bias is not None and self.bias.dtype != x.dtype: + self.bias.data = self.bias.data.to(x.dtype) + + out = bnb.matmul(x, self.weight, bias=self.bias, state=self.state) + + if not self.state.has_fp16_weights: + if self.state.CB is not None and self.state.CxB is not None: + # we converted 8-bit row major to turing/ampere format in the first inference pass + # we no longer need the row-major weight + del self.state.CB + self.weight.data = self.state.CxB + return out + +def quantize_offline(model, bits: int): + assert (bits == 4), f'bits: {bits} is not supported' + + for i, layer in enumerate(model.model.layers): + layer.self_attn.W_pack = bnb.nn.Linear4bit( + layer.self_attn.W_pack.weight.shape[1], + layer.self_attn.W_pack.weight.shape[0], + False, + torch.float16, + compress_statistics=True, + quant_type="nf4", + ) + layer.self_attn.o_proj = bnb.nn.Linear4bit( + layer.self_attn.o_proj.weight.shape[1], + layer.self_attn.o_proj.weight.shape[0], + False, + torch.float16, + compress_statistics=True, + quant_type="nf4", + ) + + layer.mlp.gate_proj = bnb.nn.Linear4bit( + layer.mlp.gate_proj.weight.shape[1], + layer.mlp.gate_proj.weight.shape[0], + False, + torch.float16, + compress_statistics=True, + quant_type="nf4", + ) + layer.mlp.down_proj = bnb.nn.Linear4bit( + layer.mlp.down_proj.weight.shape[1], + layer.mlp.down_proj.weight.shape[0], + False, + torch.float16, + compress_statistics=True, + quant_type="nf4", + ) + layer.mlp.up_proj = bnb.nn.Linear4bit( + layer.mlp.up_proj.weight.shape[1], + layer.mlp.up_proj.weight.shape[0], + False, + torch.float16, + compress_statistics=True, + quant_type="nf4", + ) + return model + +def quantize_online(model, bits: int): + def quant(weight, bias=None): + if bits == 8: + linear = Linear8bitLtOnline( + weight, + bias, + has_fp16_weights=False, + threshold=6.0, + ) + if bias is not None: + linear.bias = torch.nn.Parameter(bias) + elif bits == 4: + linear = Linear4bitOnline( + weight, + bias, + quant_type="nf4", #fp4/nf4 + ) + else: + raise ValueError("quantize only support 4/8 bit") + return linear + + for i, layer in enumerate(model.model.layers): + layer.self_attn.W_pack = quant(layer.self_attn.W_pack.weight) + layer.self_attn.o_proj = quant(layer.self_attn.o_proj.weight) + layer.mlp.gate_proj = quant(layer.mlp.gate_proj.weight) + layer.mlp.down_proj = quant(layer.mlp.down_proj.weight) + layer.mlp.up_proj = quant(layer.mlp.up_proj.weight) + return model + +def init_model_weight_int4(config, model, state_dict): + #replace Params4bit.cuda with Params4bitCuda + Params4bit.cuda = Params4bitCuda + + for i in range(config.num_hidden_layers): + weight_data = state_dict[f'model.layers.{i}.self_attn.W_pack.weight.data'] + weight_quant_state = state_dict[f'model.layers.{i}.self_attn.W_pack.weight.quant_state'] + model.model.layers[i].self_attn.W_pack.weight = Params4bit(weight_data, requires_grad=False, quant_state=weight_quant_state) + + weight_data = state_dict[f'model.layers.{i}.self_attn.o_proj.weight.data'] + weight_quant_state = state_dict[f'model.layers.{i}.self_attn.o_proj.weight.quant_state'] + model.model.layers[i].self_attn.o_proj.weight = Params4bit(weight_data, requires_grad=False, quant_state=weight_quant_state) + + weight_data = state_dict[f'model.layers.{i}.mlp.gate_proj.weight.data'] + weight_quant_state = state_dict[f'model.layers.{i}.mlp.gate_proj.weight.quant_state'] + model.model.layers[i].mlp.gate_proj.weight = Params4bit(weight_data, requires_grad=False, quant_state=weight_quant_state) + + weight_data = state_dict[f'model.layers.{i}.mlp.up_proj.weight.data'] + weight_quant_state = state_dict[f'model.layers.{i}.mlp.up_proj.weight.quant_state'] + model.model.layers[i].mlp.up_proj.weight = Params4bit(weight_data, requires_grad=False, quant_state=weight_quant_state) + + weight_data = state_dict[f'model.layers.{i}.mlp.down_proj.weight.data'] + weight_quant_state = state_dict[f'model.layers.{i}.mlp.down_proj.weight.quant_state'] + model.model.layers[i].mlp.down_proj.weight = Params4bit(weight_data, requires_grad=False, quant_state=weight_quant_state) + + model.model.layers[i].input_layernorm.weight = state_dict[f'model.layers.{i}.input_layernorm.weight'] + model.model.layers[i].post_attention_layernorm.weight = state_dict[f'model.layers.{i}.post_attention_layernorm.weight'] + + model.model.embed_tokens.weight = state_dict['model.embed_tokens.weight'] + model.model.norm.weight = state_dict['model.norm.weight'] + model.lm_head.weight = state_dict['lm_head.weight'] + return model diff --git a/special_tokens_map.json b/special_tokens_map.json new file mode 100644 index 0000000..28d1171 --- /dev/null +++ b/special_tokens_map.json @@ -0,0 +1,30 @@ +{ + "bos_token": { + "content": "", + "lstrip": false, + "normalized": true, + "rstrip": false, + "single_word": true + }, + "eos_token": { + "content": "", + "lstrip": false, + "normalized": true, + "rstrip": false, + "single_word": true + }, + "unk_token": { + "content": "", + "lstrip": false, + "normalized": true, + "rstrip": false, + "single_word": true + }, + "pad_token": { + "content": "", + "lstrip": false, + "normalized": true, + "rstrip": false, + "single_word": true + } +} diff --git a/tokenization_baichuan.py b/tokenization_baichuan.py new file mode 100644 index 0000000..ffc923f --- /dev/null +++ b/tokenization_baichuan.py @@ -0,0 +1,260 @@ +# Copyright (c) 2023, Baichuan Intelligent Technology. All rights reserved. + +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"} + +PRETRAINED_VOCAB_FILES_MAP = { + "vocab_file": {}, + "tokenizer_file": {}, +} +PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES = {} + + +class BaichuanTokenizer(PreTrainedTokenizer): + """ + Construct a Baichuan tokenizer. Based on byte-level Byte-Pair-Encoding. + + Args: + vocab_file (`str`): + Path to the vocabulary file. + """ + + vocab_files_names = VOCAB_FILES_NAMES + pretrained_vocab_files_map = PRETRAINED_VOCAB_FILES_MAP + max_model_input_sizes = PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES + model_input_names = ["input_ids", "attention_mask"] + + def __init__( + self, + vocab_file, + unk_token="", + bos_token="", + eos_token="", + pad_token=None, + 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.vocab_file = vocab_file + self.add_bos_token = add_bos_token + self.add_eos_token = add_eos_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, + ) + + 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 diff --git a/tokenizer.model b/tokenizer.model new file mode 100644 index 0000000..b3902c4 --- /dev/null +++ b/tokenizer.model @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:79452955be6b419a65984273a9f08af86042e1c2a75ee3ba989cbf620a133cc2 +size 2001107 diff --git a/tokenizer_config.json b/tokenizer_config.json new file mode 100644 index 0000000..e441aae --- /dev/null +++ b/tokenizer_config.json @@ -0,0 +1,46 @@ +{ + "add_bos_token": false, + "add_eos_token": false, + "auto_map": { + "AutoTokenizer": [ + "tokenization_baichuan.BaichuanTokenizer", + null + ] + }, + "bos_token": { + "__type": "AddedToken", + "content": "", + "lstrip": false, + "normalized": true, + "rstrip": false, + "single_word": true + }, + "clean_up_tokenization_spaces": false, + "eos_token": { + "__type": "AddedToken", + "content": "", + "lstrip": false, + "normalized": true, + "rstrip": false, + "single_word": true + }, + "model_max_length": 4096, + "pad_token": { + "__type": "AddedToken", + "content": "", + "lstrip": false, + "normalized": true, + "rstrip": false, + "single_word": true + }, + "sp_model_kwargs": {}, + "tokenizer_class": "BaichuanTokenizer", + "unk_token": { + "__type": "AddedToken", + "content": "", + "lstrip": false, + "normalized": true, + "rstrip": false, + "single_word": true + } +}