diff --git a/Baichuan2 模型社区许可协议.pdf b/Baichuan2 模型社区许可协议.pdf
new file mode 100644
index 0000000..411d36a
Binary files /dev/null and b/Baichuan2 模型社区许可协议.pdf differ
diff --git a/Community License for Baichuan2 Model.pdf b/Community License for Baichuan2 Model.pdf
new file mode 100644
index 0000000..8f1aef4
Binary files /dev/null and b/Community License for Baichuan2 Model.pdf differ
diff --git a/README.md b/README.md
index 73e1fcc..661afcd 100644
--- a/README.md
+++ b/README.md
@@ -1,3 +1,152 @@
-# Baichuan2-7B-Chat_a13444794114109440680734
+---
+language:
+ - en
+ - zh
+license: other
+tasks:
+ - text-generation
+---
-百川2-7B-对话模型
\ No newline at end of file
+
+
+
+
+ Baichuan 2
+
+
+
+
+
+
+# 目录
+
+- [📖 模型介绍](#模型介绍)
+- [⚙️ 快速开始](#快速开始)
+- [📊 Benchmark评估](#评估)
+- [📜 声明与协议](#声明与协议)
+
+# 模型介绍
+
+- Baichuan 2 是[百川智能]推出的**新一代开源大语言模型**，采用 **2.6 万亿** Tokens 的高质量语料训练。
+- Baichuan 2 在多个权威的中文、英文和多语言的通用、领域 benchmark 上取得同尺寸**最佳**的效果。
+- 本次发布包含有 **7B**、**13B** 的 **Base** 和 **Chat** 版本，并提供了 Chat 版本的 **4bits 量化**。
+- 所有版本对学术研究完全开放。同时，开发者通过邮件申请并获得官方商用许可后，即可**免费商用**，请参考[协议](#协议)章节。
+- 欢迎阅读我们的技术报告 [Baichuan 2: Open Large-scale Language Models] 获取更多信息。
+
+本次发布版本和下载链接见下表：
+
+| | 基座模型 | 对齐模型 | 对齐模型 4bits 量化 |
+|:---:|:--------------------:|:--------------------:|:--------------------------:|
+| 7B | [Baichuan2-7B-Base] | [Baichuan2-7B-Chat] | [Baichuan2-7B-Chat-4bits] |
+| 13B | [Baichuan2-13B-Base] | [Baichuan2-13B-Chat] | [Baichuan2-13B-Chat-4bits] |
+
+# 快速开始
+
+```python
+import torch
+from modelscope import snapshot_download, AutoModelForCausalLM, AutoTokenizer,GenerationConfig
+model_dir = snapshot_download("baichuan-inc/Baichuan2-7B-Chat", revision='v1.0.4')
+tokenizer = AutoTokenizer.from_pretrained(model_dir, device_map="auto",
+ trust_remote_code=True, torch_dtype=torch.float16)
+model = AutoModelForCausalLM.from_pretrained(model_dir, device_map="auto",
+ trust_remote_code=True, torch_dtype=torch.float16)
+model.generation_config = GenerationConfig.from_pretrained(model_dir)
+messages = []
+messages.append({"role": "user", "content": "讲解一下“温故而知新”"})
+response = model.chat(tokenizer, messages)
+print(response)
+messages.append({'role': 'assistant', 'content': response})
+messages.append({"role": "user", "content": "背诵一下将进酒"})
+response = model.chat(tokenizer, messages)
+print(response)
+```
+
+# Benchmark 结果
+
+我们在[通用]、[法律]、[医疗]、[数学]、[代码]和[多语言翻译]六个领域的中英文权威数据集上对模型进行了广泛测试，更多详细测评结果可查看[GitHub]。
+
+### 7B 模型结果
+
+| | **C-Eval** | **MMLU** | **CMMLU** | **Gaokao** | **AGIEval** | **BBH** |
+|:-----------------------:|:----------:|:--------:|:---------:|:----------:|:-----------:|:-------:|
+| | 5-shot | 5-shot | 5-shot | 5-shot | 5-shot | 3-shot |
+| **GPT-4** | 68.40 | 83.93 | 70.33 | 66.15 | 63.27 | 75.12 |
+| **GPT-3.5 Turbo** | 51.10 | 68.54 | 54.06 | 47.07 | 46.13 | 61.59 |
+| **LLaMA-7B** | 27.10 | 35.10 | 26.75 | 27.81 | 28.17 | 32.38 |
+| **LLaMA2-7B** | 28.90 | 45.73 | 31.38 | 25.97 | 26.53 | 39.16 |
+| **MPT-7B** | 27.15 | 27.93 | 26.00 | 26.54 | 24.83 | 35.20 |
+| **Falcon-7B** | 24.23 | 26.03 | 25.66 | 24.24 | 24.10 | 28.77 |
+| **ChatGLM2-6B** | 50.20 | 45.90 | 49.00 | 49.44 | 45.28 | 31.65 |
+| **[Baichuan-7B]** | 42.80 | 42.30 | 44.02 | 36.34 | 34.44 | 32.48 |
+| **[Baichuan2-7B-Base]** | 54.00 | 54.16 | 57.07 | 47.47 | 42.73 | 41.56 |
+
+### 13B 模型结果
+
+| | **C-Eval** | **MMLU** | **CMMLU** | **Gaokao** | **AGIEval** | **BBH** |
+|:---------------------------:|:----------:|:--------:|:---------:|:----------:|:-----------:|:-------:|
+| | 5-shot | 5-shot | 5-shot | 5-shot | 5-shot | 3-shot |
+| **GPT-4** | 68.40 | 83.93 | 70.33 | 66.15 | 63.27 | 75.12 |
+| **GPT-3.5 Turbo** | 51.10 | 68.54 | 54.06 | 47.07 | 46.13 | 61.59 |
+| **LLaMA-13B** | 28.50 | 46.30 | 31.15 | 28.23 | 28.22 | 37.89 |
+| **LLaMA2-13B** | 35.80 | 55.09 | 37.99 | 30.83 | 32.29 | 46.98 |
+| **Vicuna-13B** | 32.80 | 52.00 | 36.28 | 30.11 | 31.55 | 43.04 |
+| **Chinese-Alpaca-Plus-13B** | 38.80 | 43.90 | 33.43 | 34.78 | 35.46 | 28.94 |
+| **XVERSE-13B** | 53.70 | 55.21 | 58.44 | 44.69 | 42.54 | 38.06 |
+| **[Baichuan-13B-Base]** | 52.40 | 51.60 | 55.30 | 49.69 | 43.20 | 43.01 |
+| **[Baichuan2-13B-Base]** | 58.10 | 59.17 | 61.97 | 54.33 | 48.17 | 48.78 |
+
+
+## 训练过程模型
+
+除了训练了 2.6 万亿 Tokens 的 [Baichuan2-7B-Base] 模型，我们还提供了在此之前的另外 11 个中间过程的模型（分别对应训练了约 0.2 ~ 2.4 万亿 Tokens）供社区研究使用（[训练过程checkpoint下载]）。下图给出了这些 checkpoints 在 C-Eval、MMLU、CMMLU 三个 benchmark 上的效果变化：
+
+
+
+# 声明与协议
+
+## 声明
+
+我们在此声明，我们的开发团队并未基于 Baichuan 2 模型开发任何应用，无论是在 iOS、Android、网页或任何其他平台。我们强烈呼吁所有使用者，不要利用
+Baichuan 2 模型进行任何危害国家社会安全或违法的活动。另外，我们也要求使用者不要将 Baichuan 2
+模型用于未经适当安全审查和备案的互联网服务。我们希望所有的使用者都能遵守这个原则，确保科技的发展能在规范和合法的环境下进行。
+
+我们已经尽我们所能，来确保模型训练过程中使用的数据的合规性。然而，尽管我们已经做出了巨大的努力，但由于模型和数据的复杂性，仍有可能存在一些无法预见的问题。因此，如果由于使用
+Baichuan 2 开源模型而导致的任何问题，包括但不限于数据安全问题、公共舆论风险，或模型被误导、滥用、传播或不当利用所带来的任何风险和问题，我们将不承担任何责任。
+
+## 协议
+
+* Baichuan 2 模型的社区使用需遵循[《Baichuan 2 模型社区许可协议》]。
+* Baichuan 2 支持商用，如果将 Baichuan 2 模型或其衍生品用作商业用途，请您按照如下方式联系许可方，以进行登记并向许可方申请书面授权：联系邮箱 [opensource@baichuan-inc.com]。
+
+[GitHub]:https://github.com/baichuan-inc/Baichuan2
+[Baichuan2]:https://github.com/baichuan-inc/Baichuan2
+
+[Baichuan-7B]:https://modelscope.cn/models/baichuan-inc/baichuan-7B/summary
+[Baichuan2-7B-Base]:https://modelscope.cn/models/baichuan-inc/Baichuan2-7B-Base/summary
+[Baichuan2-7B-Chat]:https://modelscope.cn/models/baichuan-inc/Baichuan2-7B-Chat/summary
+[Baichuan2-7B-Chat-4bits]:https://modelscope.cn/models/baichuan-inc/Baichuan2-7B-Chat-4bits/summary
+[Baichuan-13B-Base]:https://modelscope.cn/models/baichuan-inc/Baichuan-13B-Base/summary
+[Baichuan2-13B-Base]:https://modelscope.cn/models/baichuan-inc/Baichuan2-13B-Base/summary
+[Baichuan2-13B-Chat]:https://modelscope.cn/models/baichuan-inc/Baichuan2-13B-Chat/summary
+[Baichuan2-13B-Chat-4bits]:https://modelscope.cn/models/baichuan-inc/Baichuan2-13B-Chat-4bits/summary
+
+[通用]:https://github.com/baichuan-inc/Baichuan2#%E9%80%9A%E7%94%A8%E9%A2%86%E5%9F%9F
+[法律]:https://github.com/baichuan-inc/Baichuan2#%E6%B3%95%E5%BE%8B%E5%8C%BB%E7%96%97
+[医疗]:https://github.com/baichuan-inc/Baichuan2#%E6%B3%95%E5%BE%8B%E5%8C%BB%E7%96%97
+[数学]:https://github.com/baichuan-inc/Baichuan2#%E6%95%B0%E5%AD%A6%E4%BB%A3%E7%A0%81
+[代码]:https://github.com/baichuan-inc/Baichuan2#%E6%95%B0%E5%AD%A6%E4%BB%A3%E7%A0%81
+[多语言翻译]:https://github.com/baichuan-inc/Baichuan2#%E5%A4%9A%E8%AF%AD%E8%A8%80%E7%BF%BB%E8%AF%91
+
+[《Baichuan 2 模型社区许可协议》]:https://modelscope.cn/models/baichuan-inc/Baichuan2-7B-Base/file/view/master/Baichuan%202%E6%A8%A1%E5%9E%8B%E7%A4%BE%E5%8C%BA%E8%AE%B8%E5%8F%AF%E5%8D%8F%E8%AE%AE.pdf
+
+[邮件申请]: mailto:opensource@baichuan-inc.com
+[Email]: mailto:opensource@baichuan-inc.com
+[opensource@baichuan-inc.com]: mailto:opensource@baichuan-inc.com
+[训练过程checkpoint下载]: https://modelscope.cn/models/baichuan-inc/Baichuan2-7B-Intermediate-Checkpoints/summary
+[百川智能]: https://www.baichuan-ai.com
+[Baichuan 2: Open Large-scale Language Models]:https://cdn.baichuan-ai.com/paper/Baichuan2-technical-report.pdf
\ No newline at end of file
diff --git a/config.json b/config.json
new file mode 100644
index 0000000..83027d2
--- /dev/null
+++ b/config.json
@@ -0,0 +1,29 @@
+{
+ "architectures": [
+ "BaichuanForCausalLM"
+ ],
+ "auto_map": {
+ "AutoConfig": "configuration_baichuan.BaichuanConfig",
+ "AutoModelForCausalLM": "modeling_baichuan.BaichuanForCausalLM"
+ },
+ "tokenizer_class": "BaichuanTokenizer",
+ "bos_token_id": 1,
+ "eos_token_id": 2,
+ "hidden_act": "silu",
+ "hidden_size": 4096,
+ "initializer_range": 0.02,
+ "intermediate_size": 11008,
+ "max_position_embeddings": 4096,
+ "model_max_length": 4096,
+ "model_type": "baichuan",
+ "num_attention_heads": 32,
+ "num_hidden_layers": 32,
+ "pad_token_id": 0,
+ "rms_norm_eps": 1e-06,
+ "_from_model_config": true,
+ "tie_word_embeddings": false,
+ "torch_dtype": "bfloat16",
+ "transformers_version": "4.29.2",
+ "use_cache": true,
+ "vocab_size": 125696
+}
diff --git a/configuration.json b/configuration.json
new file mode 100644
index 0000000..dc23a15
--- /dev/null
+++ b/configuration.json
@@ -0,0 +1,11 @@
+{
+ "framework": "pytorch",
+ "task": "text-generation",
+ "model": {
+ "type": "Baichuan2-7B-Chat"
+ },
+ "pipeline": {
+ "type": "Baichuan2-7B-chatbot-pipe"
+ },
+ "allow_remote": true
+}
diff --git a/configuration_baichuan.py b/configuration_baichuan.py
new file mode 100644
index 0000000..a260931
--- /dev/null
+++ b/configuration_baichuan.py
@@ -0,0 +1,69 @@
+# Copyright 2023 Baichuan Inc. All Rights Reserved.
+
+# Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved.
+#
+# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
+# and OPT implementations in this library. It has been modified from its
+# original forms to accommodate minor architectural differences compared
+# to GPT-NeoX and OPT used by the Meta AI team that trained the model.
+#
+# 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.
+
+from transformers.configuration_utils import PretrainedConfig
+from transformers.utils import logging
+
+
+logger = logging.get_logger(__name__)
+
+
+class BaichuanConfig(PretrainedConfig):
+ model_type = "baichuan"
+ keys_to_ignore_at_inference = ["past_key_values"]
+
+ def __init__(
+ self,
+ vocab_size=125696,
+ hidden_size=4096,
+ intermediate_size=11008,
+ num_hidden_layers=32,
+ num_attention_heads=32,
+ hidden_act="silu",
+ max_position_embeddings=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,
+ z_loss_weight=0,
+ **kwargs,
+ ):
+ self.vocab_size = vocab_size
+ self.max_position_embeddings = max_position_embeddings
+ 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
+ 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..b1cd1e5
--- /dev/null
+++ b/generation_config.json
@@ -0,0 +1,14 @@
+{
+ "pad_token_id": 0,
+ "bos_token_id": 1,
+ "eos_token_id": 2,
+ "user_token_id": 195,
+ "assistant_token_id": 196,
+ "max_new_tokens": 2048,
+ "temperature": 0.3,
+ "top_k": 5,
+ "top_p": 0.85,
+ "repetition_penalty": 1.05,
+ "do_sample": true,
+ "transformers_version": "4.29.2"
+}
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/modeling_baichuan.py b/modeling_baichuan.py
new file mode 100644
index 0000000..a202cb8
--- /dev/null
+++ b/modeling_baichuan.py
@@ -0,0 +1,785 @@
+# Copyright 2023 Baichuan Inc. All Rights Reserved.
+
+# Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved.
+#
+# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
+# and OPT implementations in this library. It has been modified from its
+# original forms to accommodate minor architectural differences compared
+# to GPT-NeoX and OPT used by the Meta AI team that trained the model.
+#
+# 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.
+
+
+from .configuration_baichuan import BaichuanConfig
+from .generation_utils import build_chat_input, TextIterStreamer
+
+import math
+from typing import List, Optional, Tuple, Union
+from threading import Thread
+
+import torch
+import torch.utils.checkpoint
+from torch import nn
+from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
+from torch.nn import functional as F
+from transformers import PreTrainedModel, PretrainedConfig
+from transformers.activations import ACT2FN
+from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
+from transformers.generation.utils import GenerationConfig
+from transformers.utils import logging, ContextManagers
+
+import os
+from contextlib import contextmanager
+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."
+ )
+
+
+# Copied from transformers.models.bart.modeling_bart._make_causal_mask
+def _make_causal_mask(
+ input_ids_shape: torch.Size, dtype: torch.dtype, device: torch.device, past_key_values_length: int = 0
+):
+ """
+ Make causal mask used for bi-directional self-attention.
+ """
+ bsz, tgt_len = input_ids_shape
+ mask = torch.full((tgt_len, tgt_len), torch.tensor(torch.finfo(dtype).min, device=device), device=device)
+ mask_cond = torch.arange(mask.size(-1), device=device)
+ mask.masked_fill_(mask_cond < (mask_cond + 1).view(mask.size(-1), 1), 0)
+ mask = mask.to(dtype)
+
+ if past_key_values_length > 0:
+ mask = torch.cat([torch.zeros(tgt_len, past_key_values_length, dtype=dtype, device=device), mask], dim=-1)
+ return mask[None, None, :, :].expand(bsz, 1, tgt_len, tgt_len + past_key_values_length)
+
+def _expand_mask(mask: torch.Tensor, dtype: torch.dtype, tgt_len: Optional[int] = None):
+ """
+ Expands attention_mask from `[bsz, seq_len]` to `[bsz, 1, tgt_seq_len, src_seq_len]`.
+ """
+ if len(mask.size()) == 3:
+ bsz, src_len, _ = mask.size()
+ tgt_len = tgt_len if tgt_len is not None else src_len
+ expanded_mask = mask[:,None,:,:].expand(bsz, 1, tgt_len, src_len).to(dtype)
+ else:
+ bsz, src_len = mask.size()
+ tgt_len = tgt_len if tgt_len is not None else src_len
+ expanded_mask = mask[:, None, None, :].expand(bsz, 1, tgt_len, src_len).to(dtype)
+
+ inverted_mask = 1.0 - expanded_mask
+
+ return inverted_mask.masked_fill(inverted_mask.to(torch.bool), torch.finfo(dtype).min)
+
+
+class RMSNorm(nn.Module):
+ def __init__(self, hidden_size, eps=1e-6):
+ """
+ RMSNorm is equivalent to T5LayerNorm
+ """
+ super().__init__()
+ self.weight = nn.Parameter(torch.ones(hidden_size))
+ self.variance_epsilon = eps
+
+ 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.variance_epsilon)
+
+ # convert into half-precision if necessary
+ if self.weight.dtype in [torch.float16, torch.bfloat16]:
+ hidden_states = hidden_states.to(self.weight.dtype)
+
+ return self.weight * hidden_states
+
+
+class RotaryEmbedding(torch.nn.Module):
+ def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None):
+ super().__init__()
+ self.inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float().to(device) / dim))
+ self.max_seq_len_cached = max_position_embeddings
+ t = torch.arange(self.max_seq_len_cached, device=self.inv_freq.device, dtype=torch.float32)
+ freqs = torch.outer(t, self.inv_freq)
+ emb = torch.cat((freqs, freqs), dim=-1)
+ self.cos_cached = emb.cos()[None, None, :, :].to(torch.float32)
+ self.sin_cached = emb.sin()[None, None, :, :].to(torch.float32)
+ def forward(self, x, seq_len=None):
+ # x: [bs, num_attention_heads, seq_len, head_size]
+ # This `if` block is unlikely to be run after we build sin/cos in `__init__`. Keep the logic here just in case.
+ if seq_len > self.max_seq_len_cached:
+ self.max_seq_len_cached = seq_len
+ t = torch.arange(self.max_seq_len_cached, device=self.inv_freq.device, dtype=torch.float32)
+ freqs = torch.outer(t, self.inv_freq)
+ emb = torch.cat((freqs, freqs), dim=-1)
+ self.cos_cached = emb.cos()[None, None, :, :].to(torch.float32).to(x.device)
+ self.sin_cached = emb.sin()[None, None, :, :].to(torch.float32).to(x.device)
+ elif self.cos_cached.device != x.device:
+ self.cos_cached = self.cos_cached.to(x.device)
+ self.sin_cached = self.sin_cached.to(x.device)
+ return (
+ self.cos_cached[:, :, :seq_len, ...],
+ self.sin_cached[:, :, :seq_len, ...],
+ )
+
+
+def rotate_half(x):
+ """Rotates half the hidden dims of the input."""
+ x1 = x[..., : x.shape[-1] // 2]
+ x2 = x[..., x.shape[-1] // 2:]
+ return torch.cat((-x2, x1), dim=-1)
+
+
+def apply_rotary_pos_emb(q, k, cos_, sin_, position_ids):
+ cos = cos_.squeeze(1).squeeze(0) # [seq_len, dim]
+ sin = sin_.squeeze(1).squeeze(0) # [seq_len, dim]
+ cos = cos[position_ids].unsqueeze(1) # [bs, 1, seq_len, dim]
+ sin = sin[position_ids].unsqueeze(1) # [bs, 1, seq_len, dim]
+ q_embed = (q.float() * cos) + (rotate_half(q.float()) * sin)
+ k_embed = (k.float() * cos) + (rotate_half(k.float()) * sin)
+ return q_embed.to(q.dtype), k_embed.to(k.dtype)
+
+
+class MLP(nn.Module):
+ def __init__(
+ self,
+ hidden_size: int,
+ intermediate_size: int,
+ hidden_act: str,
+ ):
+ super().__init__()
+ self.gate_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
+ self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=False)
+ self.up_proj = 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 Attention(nn.Module):
+ """Multi-headed attention from 'Attention Is All You Need' paper"""
+ 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.max_position_embeddings
+
+ 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}"
+ f" and `num_heads`: {self.num_heads})."
+ )
+ self.W_pack = nn.Linear(self.hidden_size, 3 * self.hidden_size, bias=False)
+ self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False)
+ self.rotary_emb = RotaryEmbedding(self.head_dim, max_position_embeddings=self.max_position_embeddings)
+
+ 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,
+ position_ids: Optional[torch.LongTensor] = 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]
+ cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
+ query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)
+ # [bsz, nh, t, hd]
+
+ 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=xops.LowerTriangularMask()
+ )
+ else:
+ 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)
+ 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 DecoderLayer(nn.Module):
+ def __init__(self, config: BaichuanConfig):
+ super().__init__()
+ self.hidden_size = config.hidden_size
+ self.self_attn = Attention(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, eps=config.rms_norm_eps)
+ self.post_attention_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ attention_mask: Optional[torch.Tensor] = None,
+ position_ids: Optional[torch.LongTensor] = 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,
+ position_ids=position_ids,
+ 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 output_attentions:
+ outputs += (self_attn_weights,)
+
+ 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 = ["DecoderLayer"]
+ _keys_to_ignore_on_load_unexpected = [r"decoder\.version"]
+
+ def _init_weights(self, module):
+ std = self.config.initializer_range
+ if isinstance(module, nn.Linear):
+ module.weight.data.normal_(mean=0.0, std=std)
+ if module.bias is not None:
+ module.bias.data.zero_()
+ elif isinstance(module, 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.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
+ self.layers = nn.ModuleList([DecoderLayer(config) for _ in range(config.num_hidden_layers)])
+ self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+ self.gradient_checkpointing = False
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ def get_input_embeddings(self):
+ return self.embed_tokens
+
+ def set_input_embeddings(self, value):
+ self.embed_tokens = value
+
+ # Copied from transformers.models.bart.modeling_bart.BartDecoder._prepare_decoder_attention_mask
+ def _prepare_decoder_attention_mask(self, attention_mask, input_shape, inputs_embeds, past_key_values_length):
+ # create causal mask
+ # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
+ combined_attention_mask = None
+ if input_shape[-1] > 1:
+ combined_attention_mask = _make_causal_mask(
+ input_shape,
+ inputs_embeds.dtype,
+ device=inputs_embeds.device,
+ past_key_values_length=past_key_values_length,
+ )
+
+ if attention_mask is not None:
+ # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
+ expanded_attn_mask = _expand_mask(attention_mask, inputs_embeds.dtype, tgt_len=input_shape[-1]).to(
+ inputs_embeds.device
+ )
+ combined_attention_mask = (
+ expanded_attn_mask if combined_attention_mask is None else expanded_attn_mask + combined_attention_mask
+ )
+
+ return combined_attention_mask
+
+ def forward(
+ self,
+ input_ids: torch.LongTensor = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ position_ids: Optional[torch.LongTensor] = None,
+ past_key_values: Optional[List[torch.FloatTensor]] = None,
+ inputs_embeds: Optional[torch.FloatTensor] = None,
+ use_cache: Optional[bool] = None,
+ output_attentions: Optional[bool] = None,
+ output_hidden_states: Optional[bool] = None,
+ return_dict: Optional[bool] = None,
+ ) -> Union[Tuple, BaseModelOutputWithPast]:
+ 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
+
+ # retrieve input_ids and inputs_embeds
+ if input_ids is not None and inputs_embeds is not None:
+ raise ValueError("You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time")
+ 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 have to specify either decoder_input_ids or decoder_inputs_embeds")
+
+ seq_length_with_past = seq_length
+ past_key_values_length = 0
+
+ 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 position_ids is None:
+ device = input_ids.device if input_ids is not None else inputs_embeds.device
+ position_ids = torch.arange(
+ past_key_values_length, seq_length + past_key_values_length, dtype=torch.long, device=device
+ )
+ position_ids = position_ids.unsqueeze(0).view(-1, seq_length)
+ else:
+ position_ids = position_ids.view(-1, seq_length).long()
+
+ if inputs_embeds is None:
+ inputs_embeds = self.embed_tokens(input_ids)
+ # embed positions
+ if attention_mask is None:
+ attention_mask = torch.ones(
+ (batch_size, seq_length_with_past), dtype=torch.bool, device=inputs_embeds.device
+ )
+ attention_mask = self._prepare_decoder_attention_mask(
+ attention_mask, (batch_size, seq_length), inputs_embeds, past_key_values_length
+ )
+
+ 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,
+ position_ids,
+ None,
+ )
+ else:
+ layer_outputs = decoder_layer(
+ hidden_states,
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ 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 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 QLinear 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,
+ position_ids: Optional[torch.LongTensor] = 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] = None,
+ output_hidden_states: Optional[bool] = None,
+ return_dict: Optional[bool] = None,
+ ) -> Union[Tuple, CausalLMOutputWithPast]:
+
+ 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
+ )
+ 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,
+ position_ids=position_ids,
+ 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 prepare_inputs_for_generation(
+ self, input_ids, past_key_values=None, attention_mask=None, inputs_embeds=None, **kwargs
+ ):
+ if past_key_values:
+ input_ids = input_ids[:, -1:]
+
+ position_ids = kwargs.get("position_ids", None)
+ if attention_mask is not None and position_ids is None:
+ # create position_ids on the fly for batch generation
+ position_ids = attention_mask.long().cumsum(-1) - 1
+ position_ids.masked_fill_(attention_mask == 0, 1)
+ if past_key_values:
+ position_ids = position_ids[:, -1].unsqueeze(-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(
+ {
+ "position_ids": position_ids,
+ "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):
+ reordered_past = ()
+ for layer_past in past_key_values:
+ reordered_past += (tuple(past_state.index_select(0, beam_idx) for past_state in layer_past),)
+ return reordered_past
+
+ 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 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/ms_wrapper.py b/ms_wrapper.py
new file mode 100644
index 0000000..0fa3d36
--- /dev/null
+++ b/ms_wrapper.py
@@ -0,0 +1,74 @@
+import os
+import torch
+from typing import Union, Dict, Any
+from modelscope.pipelines.builder import PIPELINES
+from modelscope.models.builder import MODELS
+from modelscope.utils.constant import Tasks
+from modelscope.pipelines.base import Pipeline
+from modelscope.outputs import OutputKeys
+from modelscope.pipelines.nlp.text_generation_pipeline import TextGenerationPipeline
+from modelscope.models.base import Model, TorchModel
+from modelscope.utils.logger import get_logger
+from transformers import AutoModelForCausalLM, AutoTokenizer, GenerationConfig
+from transformers.generation.utils import GenerationConfig
+
+@PIPELINES.register_module(Tasks.text_generation, module_name='Baichuan2-7B-chatbot-pipe')
+class Baichuan7BChatTextGenerationPipeline(TextGenerationPipeline):
+ def __init__(
+ self,
+ model: Union[Model, str],
+ *args,
+ **kwargs):
+ self.model = Baichuan7BChatTextGeneration(model) if isinstance(model, str) else model
+ super().__init__(model=model, **kwargs)
+
+ def preprocess(self, inputs, **preprocess_params) -> Dict[str, Any]:
+ return inputs
+
+ def _sanitize_parameters(self, **pipeline_parameters):
+ return {},pipeline_parameters,{}
+
+ # define the forward pass
+ def forward(self, inputs: Dict, **forward_params) -> Dict[str, Any]:
+ output = {}
+ device = self.model.model.device
+ input_ids = self.model.tokenizer(inputs, return_tensors="pt").input_ids.to(device)
+ pred = self.model.model.generate(input_ids,**forward_params)
+ out = self.model.tokenizer.decode(pred.cpu()[0], skip_special_tokens=True)
+ output['text'] = out
+ return output
+
+ # format the outputs from pipeline
+ def postprocess(self, input, **kwargs) -> Dict[str, Any]:
+ return input
+
+
+@MODELS.register_module(Tasks.text_generation, module_name='Baichuan2-7B-Chat')
+class Baichuan7BChatTextGeneration(TorchModel):
+ def __init__(self, model_dir=None, *args, **kwargs):
+ super().__init__(model_dir, *args, **kwargs)
+ self.logger = get_logger()
+ # loading tokenizer
+ self.tokenizer = AutoTokenizer.from_pretrained(model_dir, trust_remote_code=True)
+ self.model = AutoModelForCausalLM.from_pretrained(model_dir, device_map="auto", torch_dtype=torch.float16, trust_remote_code=True)
+ # self.model = AutoModelForCausalLM.from_pretrained(model_dir, device_map="auto",trust_remote_code=True)
+ self.model.generation_config = GenerationConfig.from_pretrained(model_dir)
+ self.model = self.model.eval()
+
+ def forward(self,input: Dict, *args, **kwargs) -> Dict[str, Any]:
+ output = {}
+ response = self.model.chat(self.tokenizer, input, *args, **kwargs)
+ history = input.copy()
+ history.append({'role': 'assistant', 'content': response})
+ return {OutputKeys.RESPONSE:response, OutputKeys.HISTORY: history}
+
+ def quantize(self, bits: int):
+ self.model = self.model.quantize(bits)
+ return self
+
+ def infer(self, input, **kwargs):
+ device = self.model.device
+ input_ids = self.tokenizer(input, return_tensors="pt").input_ids.to(device)
+ pred = self.model.generate(input_ids,**kwargs)
+ out = self.tokenizer.decode(pred.cpu()[0], skip_special_tokens=True)
+ return out
diff --git a/pytorch_model.bin b/pytorch_model.bin
new file mode 100644
index 0000000..7a268cc
--- /dev/null
+++ b/pytorch_model.bin
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:594b540d36751fa3b199c609617deafc0329050937c0767fdfb2d38b72d8bec9
+size 15012021529
diff --git a/quantizer.py b/quantizer.py
new file mode 100644
index 0000000..239a2fb
--- /dev/null
+++ b/quantizer.py
@@ -0,0 +1,210 @@
+import bitsandbytes as bnb
+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
\ No newline at end of file
diff --git a/special_tokens_map.json b/special_tokens_map.json
new file mode 100644
index 0000000..5819ea2
--- /dev/null
+++ b/special_tokens_map.json
@@ -0,0 +1,30 @@
+{
+ "bos_token": {
+ "content": "",
+ "lstrip": false,
+ "normalized": true,
+ "rstrip": false,
+ "single_word": false
+ },
+ "eos_token": {
+ "content": "",
+ "lstrip": false,
+ "normalized": true,
+ "rstrip": false,
+ "single_word": false
+ },
+ "pad_token": {
+ "content": "",
+ "lstrip": false,
+ "normalized": true,
+ "rstrip": false,
+ "single_word": false
+ },
+ "unk_token": {
+ "content": "",
+ "lstrip": false,
+ "normalized": true,
+ "rstrip": false,
+ "single_word": false
+ }
+}
diff --git a/tokenization_baichuan.py b/tokenization_baichuan.py
new file mode 100644
index 0000000..f8ed2d1
--- /dev/null
+++ b/tokenization_baichuan.py
@@ -0,0 +1,253 @@
+# Copyright 2023 Baichuan Inc. All Rights Reserved.
+
+# Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved.
+#
+# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
+# and OPT implementations in this library. It has been modified from its
+# original forms to accommodate minor architectural differences compared
+# to GPT-NeoX and OPT used by the Meta AI team that trained the model.
+#
+# 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.
+
+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..b3e0f76
--- /dev/null
+++ b/tokenizer_config.json
@@ -0,0 +1,36 @@
+{
+ "auto_map": {
+ "AutoTokenizer": ["tokenization_baichuan.BaichuanTokenizer", null]
+ },
+ "add_bos_token": false,
+ "add_eos_token": false,
+ "use_fast": false,
+ "clean_up_tokenization_spaces": false,
+ "eos_token": {
+ "__type": "AddedToken",
+ "content": "",
+ "lstrip": false,
+ "normalized": true,
+ "rstrip": false,
+ "single_word": true
+ },
+ "model_max_length": 4096,
+ "sp_model_kwargs": {},
+ "tokenizer_class": "BaichuanTokenizer",
+ "pad_token": {
+ "__type": "AddedToken",
+ "content": "",
+ "lstrip": false,
+ "normalized": true,
+ "rstrip": false,
+ "single_word": true
+ },
+ "unk_token": {
+ "__type": "AddedToken",
+ "content": "",
+ "lstrip": false,
+ "normalized": true,
+ "rstrip": false,
+ "single_word": true
+ }
+}