kanhatakeyama OsakanaTeishoku commited on
Commit
a6dd058
0 Parent(s):

Duplicate from team-hatakeyama-phase2/Tanuki-8x8B-vision-v0-stage2-ckpt900

Browse files

Co-authored-by: OsakanaTeishoku <[email protected]>

.gitattributes ADDED
@@ -0,0 +1,35 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ *.7z filter=lfs diff=lfs merge=lfs -text
2
+ *.arrow filter=lfs diff=lfs merge=lfs -text
3
+ *.bin filter=lfs diff=lfs merge=lfs -text
4
+ *.bz2 filter=lfs diff=lfs merge=lfs -text
5
+ *.ckpt filter=lfs diff=lfs merge=lfs -text
6
+ *.ftz filter=lfs diff=lfs merge=lfs -text
7
+ *.gz filter=lfs diff=lfs merge=lfs -text
8
+ *.h5 filter=lfs diff=lfs merge=lfs -text
9
+ *.joblib filter=lfs diff=lfs merge=lfs -text
10
+ *.lfs.* filter=lfs diff=lfs merge=lfs -text
11
+ *.mlmodel filter=lfs diff=lfs merge=lfs -text
12
+ *.model filter=lfs diff=lfs merge=lfs -text
13
+ *.msgpack filter=lfs diff=lfs merge=lfs -text
14
+ *.npy filter=lfs diff=lfs merge=lfs -text
15
+ *.npz filter=lfs diff=lfs merge=lfs -text
16
+ *.onnx filter=lfs diff=lfs merge=lfs -text
17
+ *.ot filter=lfs diff=lfs merge=lfs -text
18
+ *.parquet filter=lfs diff=lfs merge=lfs -text
19
+ *.pb filter=lfs diff=lfs merge=lfs -text
20
+ *.pickle filter=lfs diff=lfs merge=lfs -text
21
+ *.pkl filter=lfs diff=lfs merge=lfs -text
22
+ *.pt filter=lfs diff=lfs merge=lfs -text
23
+ *.pth filter=lfs diff=lfs merge=lfs -text
24
+ *.rar filter=lfs diff=lfs merge=lfs -text
25
+ *.safetensors filter=lfs diff=lfs merge=lfs -text
26
+ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
27
+ *.tar.* filter=lfs diff=lfs merge=lfs -text
28
+ *.tar filter=lfs diff=lfs merge=lfs -text
29
+ *.tflite filter=lfs diff=lfs merge=lfs -text
30
+ *.tgz filter=lfs diff=lfs merge=lfs -text
31
+ *.wasm filter=lfs diff=lfs merge=lfs -text
32
+ *.xz filter=lfs diff=lfs merge=lfs -text
33
+ *.zip filter=lfs diff=lfs merge=lfs -text
34
+ *.zst filter=lfs diff=lfs merge=lfs -text
35
+ *tfevents* filter=lfs diff=lfs merge=lfs -text
README.md ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ 8x8bの最終
2
+
3
+ tanuki-8x8b-baseの古いcheckpoint( https://huggingface.co/team-hatakeyama-phase2/Tanuki-8x8B-checkpoint-for-vision )をllava化した
config.json ADDED
@@ -0,0 +1,51 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "_name_or_path": "/storage5/shared/Nishijima/Llama-3-8b-MoE/5th_tonyu_iter_20000",
3
+ "architectures": [
4
+ "LlavaTanukiForCausalLM"
5
+ ],
6
+ "attention_bias": false,
7
+ "attention_dropout": 0.0,
8
+ "auto_map": {
9
+ "AutoConfig": "configuration_tanuki.TanukiConfig",
10
+ "AutoModelForCausalLM": "modeling_tanuki.TanukiForCausalLM"
11
+ },
12
+ "bos_token_id": 1,
13
+ "eos_token_id": 2,
14
+ "hidden_act": "silu",
15
+ "hidden_size": 4096,
16
+ "image_aspect_ratio": "square",
17
+ "initializer_range": 0.02,
18
+ "intermediate_size": 14336,
19
+ "max_position_embeddings": 2048,
20
+ "mm_hidden_size": 2304,
21
+ "mm_projector_type": "mlp2x_gelu",
22
+ "mm_vision_select_feature": "patch",
23
+ "mm_vision_select_layer": -2,
24
+ "mm_vision_tower": "google/siglip-so400m-patch14-384",
25
+ "model_type": "llava-jp",
26
+ "num_attention_heads": 32,
27
+ "num_experts_per_tok": 2,
28
+ "num_experts_per_topk": 2,
29
+ "num_hidden_layers": 32,
30
+ "num_key_value_heads": 8,
31
+ "num_local_experts": 8,
32
+ "output_router_logits": false,
33
+ "rms_norm_eps": 1e-05,
34
+ "rope_theta": 10000,
35
+ "router_aux_loss_coef": 0.001,
36
+ "router_jitter_noise": 0.0,
37
+ "scales": [
38
+ 1.0,
39
+ 0.5
40
+ ],
41
+ "sliding_window": null,
42
+ "tie_word_embeddings": false,
43
+ "tokenizer_model_max_length": 8192,
44
+ "tokenizer_padding_side": "right",
45
+ "torch_dtype": "bfloat16",
46
+ "transformers_version": "4.43.3",
47
+ "tune_mm_mlp_adapter": false,
48
+ "use_cache": false,
49
+ "use_mm_proj": true,
50
+ "vocab_size": 65024
51
+ }
configuration_tanuki.py ADDED
@@ -0,0 +1,169 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # coding=utf-8
2
+ # Copyright 2023 Mixtral AI and the HuggingFace Inc. team. All rights reserved.
3
+ #
4
+ # Licensed under the Apache License, Version 2.0 (the "License");
5
+ # you may not use this file except in compliance with the License.
6
+ # You may obtain a copy of the License at
7
+ #
8
+ # http://www.apache.org/licenses/LICENSE-2.0
9
+ #
10
+ # Unless required by applicable law or agreed to in writing, software
11
+ # distributed under the License is distributed on an "AS IS" BASIS,
12
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13
+ # See the License for the specific language governing permissions and
14
+ # limitations under the License.
15
+ """Tanuki model configuration"""
16
+
17
+ from transformers.configuration_utils import PretrainedConfig
18
+ from transformers.utils import logging
19
+
20
+
21
+ logger = logging.get_logger(__name__)
22
+
23
+
24
+ class TanukiConfig(PretrainedConfig):
25
+ r"""
26
+ This is the configuration class to store the configuration of a [`TanukiModel`]. It is used to instantiate an
27
+ Tanuki model according to the specified arguments, defining the model architecture. Instantiating a configuration
28
+ with the defaults will yield a similar configuration to that of the Tanuki-7B-v0.1 or Tanuki-7B-Instruct-v0.1.
29
+
30
+ [mixtralai/Tanuki-8x7B](https://huggingface.co/mixtralai/Tanuki-8x7B)
31
+ [mixtralai/Tanuki-7B-Instruct-v0.1](https://huggingface.co/mixtralai/Tanuki-7B-Instruct-v0.1)
32
+
33
+ Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
34
+ documentation from [`PretrainedConfig`] for more information.
35
+
36
+
37
+ Args:
38
+ vocab_size (`int`, *optional*, defaults to 32000):
39
+ Vocabulary size of the Tanuki model. Defines the number of different tokens that can be represented by the
40
+ `inputs_ids` passed when calling [`TanukiModel`]
41
+ hidden_size (`int`, *optional*, defaults to 4096):
42
+ Dimension of the hidden representations.
43
+ intermediate_size (`int`, *optional*, defaults to 14336):
44
+ Dimension of the MLP representations.
45
+ num_hidden_layers (`int`, *optional*, defaults to 32):
46
+ Number of hidden layers in the Transformer encoder.
47
+ num_attention_heads (`int`, *optional*, defaults to 32):
48
+ Number of attention heads for each attention layer in the Transformer encoder.
49
+ num_key_value_heads (`int`, *optional*, defaults to 8):
50
+ This is the number of key_value heads that should be used to implement Grouped Query Attention. If
51
+ `num_key_value_heads=num_attention_heads`, the model will use Multi Head Attention (MHA), if
52
+ `num_key_value_heads=1` the model will use Multi Query Attention (MQA) otherwise GQA is used. When
53
+ converting a multi-head checkpoint to a GQA checkpoint, each group key and value head should be constructed
54
+ by meanpooling all the original heads within that group. For more details checkout [this
55
+ paper](https://arxiv.org/pdf/2305.13245.pdf). If it is not specified, will default to `8`.
56
+ hidden_act (`str` or `function`, *optional*, defaults to `"silu"`):
57
+ The non-linear activation function (function or string) in the decoder.
58
+ max_position_embeddings (`int`, *optional*, defaults to `4096*32`):
59
+ The maximum sequence length that this model might ever be used with. Tanuki's sliding window attention
60
+ allows sequence of up to 4096*32 tokens.
61
+ initializer_range (`float`, *optional*, defaults to 0.02):
62
+ The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
63
+ rms_norm_eps (`float`, *optional*, defaults to 1e-05):
64
+ The epsilon used by the rms normalization layers.
65
+ use_cache (`bool`, *optional*, defaults to `True`):
66
+ Whether or not the model should return the last key/values attentions (not used by all models). Only
67
+ relevant if `config.is_decoder=True`.
68
+ pad_token_id (`int`, *optional*):
69
+ The id of the padding token.
70
+ bos_token_id (`int`, *optional*, defaults to 1):
71
+ The id of the "beginning-of-sequence" token.
72
+ eos_token_id (`int`, *optional*, defaults to 2):
73
+ The id of the "end-of-sequence" token.
74
+ tie_word_embeddings (`bool`, *optional*, defaults to `False`):
75
+ Whether the model's input and output word embeddings should be tied.
76
+ rope_theta (`float`, *optional*, defaults to 1000000.0):
77
+ The base period of the RoPE embeddings.
78
+ sliding_window (`int`, *optional*):
79
+ Sliding window attention window size. If not specified, will default to `4096`.
80
+ attention_dropout (`float`, *optional*, defaults to 0.0):
81
+ The dropout ratio for the attention probabilities.
82
+ num_experts_per_tok (`int`, *optional*, defaults to 2):
83
+ The number of experts to route per-token, can be also interpreted as the `top-k` routing
84
+ parameter
85
+ num_local_experts (`int`, *optional*, defaults to 8):
86
+ Number of experts per Sparse MLP layer.
87
+ output_router_logits (`bool`, *optional*, defaults to `False`):
88
+ Whether or not the router logits should be returned by the model. Enabeling this will also
89
+ allow the model to output the auxiliary loss. See [here]() for more details
90
+ router_aux_loss_coef (`float`, *optional*, defaults to 0.001):
91
+ The aux loss factor for the total loss.
92
+ router_jitter_noise (`float`, *optional*, defaults to 0.0):
93
+ Amount of noise to add to the router.
94
+
95
+ ```python
96
+ >>> from transformers import TanukiModel, TanukiConfig
97
+
98
+ >>> # Initializing a Tanuki 7B style configuration
99
+ >>> configuration = TanukiConfig()
100
+
101
+ >>> # Initializing a model from the Tanuki 7B style configuration
102
+ >>> model = TanukiModel(configuration)
103
+
104
+ >>> # Accessing the model configuration
105
+ >>> configuration = model.config
106
+ ```"""
107
+
108
+ model_type = "Tanuki"
109
+ keys_to_ignore_at_inference = ["past_key_values"]
110
+
111
+ def __init__(
112
+ self,
113
+ vocab_size=32000,
114
+ hidden_size=4096,
115
+ intermediate_size=14336,
116
+ num_hidden_layers=32,
117
+ num_attention_heads=32,
118
+ num_key_value_heads=8,
119
+ hidden_act="silu",
120
+ max_position_embeddings=4096 * 32,
121
+ initializer_range=0.02,
122
+ rms_norm_eps=1e-5,
123
+ use_cache=True,
124
+ pad_token_id=None,
125
+ bos_token_id=1,
126
+ eos_token_id=2,
127
+ tie_word_embeddings=False,
128
+ rope_theta=1e6,
129
+ sliding_window=None,
130
+ attention_dropout=0.0,
131
+ num_experts_per_tok=2,
132
+ num_local_experts=8,
133
+ output_router_logits=False,
134
+ router_aux_loss_coef=0.001,
135
+ router_jitter_noise=0.0,
136
+ **kwargs,
137
+ ):
138
+ self.vocab_size = vocab_size
139
+ self.max_position_embeddings = max_position_embeddings
140
+ self.hidden_size = hidden_size
141
+ self.intermediate_size = intermediate_size
142
+ self.num_hidden_layers = num_hidden_layers
143
+ self.num_attention_heads = num_attention_heads
144
+ self.sliding_window = sliding_window
145
+
146
+ # for backward compatibility
147
+ if num_key_value_heads is None:
148
+ num_key_value_heads = num_attention_heads
149
+
150
+ self.num_key_value_heads = num_key_value_heads
151
+ self.hidden_act = hidden_act
152
+ self.initializer_range = initializer_range
153
+ self.rms_norm_eps = rms_norm_eps
154
+ self.use_cache = use_cache
155
+ self.rope_theta = rope_theta
156
+ self.attention_dropout = attention_dropout
157
+
158
+ self.num_experts_per_tok = num_experts_per_tok
159
+ self.num_local_experts = num_local_experts
160
+ self.output_router_logits = output_router_logits
161
+ self.router_aux_loss_coef = router_aux_loss_coef
162
+ self.router_jitter_noise = router_jitter_noise
163
+ super().__init__(
164
+ pad_token_id=pad_token_id,
165
+ bos_token_id=bos_token_id,
166
+ eos_token_id=eos_token_id,
167
+ tie_word_embeddings=tie_word_embeddings,
168
+ **kwargs,
169
+ )
generation_config.json ADDED
@@ -0,0 +1,6 @@
 
 
 
 
 
 
 
1
+ {
2
+ "_from_model_config": true,
3
+ "bos_token_id": 1,
4
+ "eos_token_id": 2,
5
+ "transformers_version": "4.43.3"
6
+ }
model-00001-of-00020.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:ceae63ea816cd2743272e69a112ccf1edd36b55837e69b277ff0a2d49214f83d
3
+ size 4928460912
model-00002-of-00020.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:d242df0803ecedd5c0c1ff0e7378cf233e586919ecd81161aff97d2aecba6913
3
+ size 4983004016
model-00003-of-00020.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:5c32f51dcbd8549383e07da0284d0bb1886d7cf8aab8ba68c65b5b696a484acd
3
+ size 4899035200
model-00004-of-00020.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:e275df2bbad319f6e278f714ab9b3134f8ebeb6404829bf1433f2205467ce822
3
+ size 4983004016
model-00005-of-00020.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:54fc2a84f03ccc21b00a2c287d1b8c6dc62466d41db4a5c8f4f1785e726a6fef
3
+ size 4983004016
model-00006-of-00020.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:117f34ebe34dcf7deb401a605a8557d86965aff79a3f8701b666d439a29cab66
3
+ size 4983004016
model-00007-of-00020.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:7bce81972b1c84002bdbf11a21f78c37070bb96273c3abff5ae499b8e0dda7e2
3
+ size 4899035248
model-00008-of-00020.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:eb7e231515d6e4dd768528af0c26a783df4e8eca5510715f762dc19bac0862d5
3
+ size 4983004072
model-00009-of-00020.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:25e4541fdad9025d654b19084b7b1f09fbb9f853ff48e7eb1181232fb59bd87f
3
+ size 4983004072
model-00010-of-00020.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:14eb1d0af2c0e92d3975ccee4f20309f38ca8e73c5cf3c0aa7ae479642f6ce37
3
+ size 4899035248
model-00011-of-00020.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:77870113bd073007c621e36b2d0d8e5669922d0ed7c37542cd9eb71bd1ee7275
3
+ size 4983004072
model-00012-of-00020.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:c66802325a591aa0c92ea919e57a882d5599b32650331bb2a1e0cf70d0a63615
3
+ size 4983004072
model-00013-of-00020.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:3137cb1a2dfe43649adf45ceaab5c81398147bf7443806647bb3241626bef6d0
3
+ size 4983004072
model-00014-of-00020.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:a5a513ee80aa8a7e059664b4abb95f603fd87456f3ab2d68a99ce31a65ae8ec4
3
+ size 4899035248
model-00015-of-00020.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:5407a7e22f41505b178cb0e4597ac82218bbad5ab1593e8d512096ae33ba234c
3
+ size 4983004072
model-00016-of-00020.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:86aed4d02e4e90d5ec33b3e2c561003f3c9166766a49064790c06c6fd7535da0
3
+ size 4983004072
model-00017-of-00020.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:833b0e40f2c3da4434ebb06c989d540e11267f627b1a30106c8da3feab32c48a
3
+ size 4899035248
model-00018-of-00020.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:847f73a475703f6389cbaade47e11a0de0ee9521927939e9df0cdf925e996659
3
+ size 4983004072
model-00019-of-00020.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:64080336642a5e3fda0161255591187950b5422fc7c83d491d8d717bdfb16d41
3
+ size 4997943552
model-00020-of-00020.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:dab213e3544c346f3991be3e928ae2054f88aafdda24ea34c566fe7437720244
3
+ size 638116456
model.safetensors.index.json ADDED
The diff for this file is too large to render. See raw diff
 
modeling_tanuki.py ADDED
@@ -0,0 +1,1757 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # coding=utf-8
2
+ # Copyright 2023 Mistral AI and the HuggingFace Inc. team. All rights reserved.
3
+ #
4
+ # This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
5
+ # and OPT implementations in this library. It has been modified from its
6
+ # original forms to accommodate minor architectural differences compared
7
+ # to GPT-NeoX and OPT used by the Meta AI team that trained the model.
8
+ #
9
+ # Licensed under the Apache License, Version 2.0 (the "License");
10
+ # you may not use this file except in compliance with the License.
11
+ # You may obtain a copy of the License at
12
+ #
13
+ # http://www.apache.org/licenses/LICENSE-2.0
14
+ #
15
+ # Unless required by applicable law or agreed to in writing, software
16
+ # distributed under the License is distributed on an "AS IS" BASIS,
17
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
18
+ # See the License for the specific language governing permissions and
19
+ # limitations under the License.
20
+ """PyTorch Tanuki model."""
21
+
22
+ import inspect
23
+ import math
24
+ from typing import List, Optional, Tuple, Union
25
+
26
+ import torch
27
+ import torch.nn.functional as F
28
+ import torch.utils.checkpoint
29
+ from torch import nn
30
+ from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
31
+
32
+ from transformers.activations import ACT2FN
33
+ from transformers.cache_utils import Cache, DynamicCache, StaticCache
34
+ from transformers.modeling_attn_mask_utils import (
35
+ AttentionMaskConverter,
36
+ _prepare_4d_causal_attention_mask,
37
+ )
38
+ from transformers.modeling_outputs import (
39
+ MoeCausalLMOutputWithPast,
40
+ MoeModelOutputWithPast,
41
+ SequenceClassifierOutputWithPast,
42
+ TokenClassifierOutput,
43
+ )
44
+ from transformers.modeling_utils import PreTrainedModel
45
+ from transformers.pytorch_utils import is_torch_greater_or_equal_than_1_13
46
+ from transformers.utils import (
47
+ add_start_docstrings,
48
+ add_start_docstrings_to_model_forward,
49
+ is_flash_attn_2_available,
50
+ is_flash_attn_greater_or_equal_2_10,
51
+ logging,
52
+ replace_return_docstrings,
53
+ )
54
+ from transformers.utils.import_utils import is_torch_fx_available
55
+ from .configuration_tanuki import TanukiConfig
56
+
57
+
58
+ if is_flash_attn_2_available():
59
+ #from transformers.modeling_flash_attention_utils import _flash_attention_forward
60
+ from flash_attn import flash_attn_func, flash_attn_varlen_func
61
+ from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input # noqa
62
+
63
+ _flash_supports_window_size = "window_size" in list(inspect.signature(flash_attn_func).parameters)
64
+
65
+ # This makes `_prepare_4d_causal_attention_mask` a leaf function in the FX graph.
66
+ # It means that the function will not be traced through and simply appear as a node in the graph.
67
+ if is_torch_fx_available():
68
+ if not is_torch_greater_or_equal_than_1_13:
69
+ import torch.fx
70
+
71
+ _prepare_4d_causal_attention_mask = torch.fx.wrap(_prepare_4d_causal_attention_mask)
72
+
73
+
74
+ logger = logging.get_logger(__name__)
75
+
76
+ _CONFIG_FOR_DOC = "TanukiConfig"
77
+
78
+
79
+ def load_balancing_loss_func(
80
+ gate_logits: torch.Tensor, num_experts: torch.Tensor = None, top_k=2, attention_mask: Optional[torch.Tensor] = None
81
+ ) -> float:
82
+ r"""
83
+ Computes auxiliary load balancing loss as in Switch Transformer - implemented in Pytorch.
84
+
85
+ See Switch Transformer (https://arxiv.org/abs/2101.03961) for more details. This function implements the loss
86
+ function presented in equations (4) - (6) of the paper. It aims at penalizing cases where the routing between
87
+ experts is too unbalanced.
88
+
89
+ Args:
90
+ gate_logits (Union[`torch.Tensor`, Tuple[torch.Tensor]):
91
+ Logits from the `gate`, should be a tuple of model.config.num_hidden_layers tensors of
92
+ shape [batch_size X sequence_length, num_experts].
93
+ attention_mask (`torch.Tensor`, None):
94
+ The attention_mask used in forward function
95
+ shape [batch_size X sequence_length] if not None.
96
+ num_experts (`int`, *optional*):
97
+ Number of experts
98
+
99
+ Returns:
100
+ The auxiliary loss.
101
+ """
102
+ if gate_logits is None or not isinstance(gate_logits, tuple):
103
+ return 0
104
+
105
+ if isinstance(gate_logits, tuple):
106
+ compute_device = gate_logits[0].device
107
+ concatenated_gate_logits = torch.cat([layer_gate.to(compute_device) for layer_gate in gate_logits], dim=0)
108
+
109
+ routing_weights = torch.nn.functional.softmax(concatenated_gate_logits, dim=-1)
110
+
111
+ _, selected_experts = torch.topk(routing_weights, top_k, dim=-1)
112
+
113
+ expert_mask = torch.nn.functional.one_hot(selected_experts, num_experts)
114
+
115
+ if attention_mask is None:
116
+ # Compute the percentage of tokens routed to each experts
117
+ tokens_per_expert = torch.mean(expert_mask.float(), dim=0)
118
+
119
+ # Compute the average probability of routing to these experts
120
+ router_prob_per_expert = torch.mean(routing_weights, dim=0)
121
+ else:
122
+ batch_size, sequence_length = attention_mask.shape
123
+ num_hidden_layers = concatenated_gate_logits.shape[0] // (batch_size * sequence_length)
124
+
125
+ # Compute the mask that masks all padding tokens as 0 with the same shape of expert_mask
126
+ expert_attention_mask = (
127
+ attention_mask[None, :, :, None, None]
128
+ .expand((num_hidden_layers, batch_size, sequence_length, top_k, num_experts))
129
+ .reshape(-1, top_k, num_experts)
130
+ .to(compute_device)
131
+ )
132
+
133
+ # Compute the percentage of tokens routed to each experts
134
+ tokens_per_expert = torch.sum(expert_mask.float() * expert_attention_mask, dim=0) / torch.sum(
135
+ expert_attention_mask, dim=0
136
+ )
137
+
138
+ # Compute the mask that masks all padding tokens as 0 with the same shape of tokens_per_expert
139
+ router_per_expert_attention_mask = (
140
+ attention_mask[None, :, :, None]
141
+ .expand((num_hidden_layers, batch_size, sequence_length, num_experts))
142
+ .reshape(-1, num_experts)
143
+ .to(compute_device)
144
+ )
145
+
146
+ # Compute the average probability of routing to these experts
147
+ router_prob_per_expert = torch.sum(routing_weights * router_per_expert_attention_mask, dim=0) / torch.sum(
148
+ router_per_expert_attention_mask, dim=0
149
+ )
150
+
151
+ overall_loss = torch.sum(tokens_per_expert * router_prob_per_expert.unsqueeze(0))
152
+ return overall_loss * num_experts
153
+
154
+ # Copied from transformers.models.llama.modeling_llama._get_unpad_data
155
+ def _get_unpad_data(attention_mask):
156
+ seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32)
157
+ indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()
158
+ max_seqlen_in_batch = seqlens_in_batch.max().item()
159
+ cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0))
160
+ return (
161
+ indices,
162
+ cu_seqlens,
163
+ max_seqlen_in_batch,
164
+ )
165
+
166
+
167
+ # Copied from transformers.models.llama.modeling_llama.LlamaRMSNorm with Llama->Mixtral
168
+ class TanukiRMSNorm(nn.Module):
169
+ def __init__(self, hidden_size, eps=1e-6):
170
+ """
171
+ TanukiRMSNorm is equivalent to T5LayerNorm
172
+ """
173
+ super().__init__()
174
+ self.weight = nn.Parameter(torch.ones(hidden_size))
175
+ self.variance_epsilon = eps
176
+
177
+ def forward(self, hidden_states):
178
+ input_dtype = hidden_states.dtype
179
+ hidden_states = hidden_states.to(torch.float32)
180
+ variance = hidden_states.pow(2).mean(-1, keepdim=True)
181
+ hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
182
+ return self.weight * hidden_states.to(input_dtype)
183
+
184
+
185
+ # copied from transformers.models.mistral.modeling_mistral.MistralRotaryEmbedding with Mistral->Mixtral
186
+ # TODO @longjie no longer copied from Mistral after static cache
187
+ class TanukiRotaryEmbedding(nn.Module):
188
+ def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None):
189
+ super().__init__()
190
+
191
+ self.dim = dim
192
+ self.max_position_embeddings = max_position_embeddings
193
+ self.base = base
194
+ inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2, dtype=torch.int64).float().to(device) / self.dim))
195
+ self.register_buffer("inv_freq", inv_freq, persistent=False)
196
+
197
+ # Build here to make `torch.jit.trace` work.
198
+ self._set_cos_sin_cache(
199
+ seq_len=max_position_embeddings, device=self.inv_freq.device, dtype=torch.get_default_dtype()
200
+ )
201
+
202
+ def _set_cos_sin_cache(self, seq_len, device, dtype):
203
+ self.max_seq_len_cached = seq_len
204
+ t = torch.arange(self.max_seq_len_cached, device=device, dtype=torch.int64).type_as(self.inv_freq)
205
+
206
+ freqs = torch.outer(t, self.inv_freq)
207
+ # Different from paper, but it uses a different permutation in order to obtain the same calculation
208
+ emb = torch.cat((freqs, freqs), dim=-1)
209
+ self.register_buffer("cos_cached", emb.cos().to(dtype), persistent=False)
210
+ self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False)
211
+
212
+ def forward(self, x, seq_len=None):
213
+ # x: [bs, num_attention_heads, seq_len, head_size]
214
+ if seq_len > self.max_seq_len_cached:
215
+ self._set_cos_sin_cache(seq_len=seq_len, device=x.device, dtype=x.dtype)
216
+
217
+ return (
218
+ self.cos_cached[:seq_len].to(dtype=x.dtype),
219
+ self.sin_cached[:seq_len].to(dtype=x.dtype),
220
+ )
221
+
222
+
223
+ # Copied from transformers.models.llama.modeling_llama.rotate_half
224
+ def rotate_half(x):
225
+ """Rotates half the hidden dims of the input."""
226
+ x1 = x[..., : x.shape[-1] // 2]
227
+ x2 = x[..., x.shape[-1] // 2 :]
228
+ return torch.cat((-x2, x1), dim=-1)
229
+
230
+
231
+ # copied from transformers.models.mistral.modeling_mistral.apply_rotary_pos_emb
232
+ # TODO @longjie no longer copied from Mistral after static cache
233
+ def apply_rotary_pos_emb(q, k, cos, sin, position_ids, unsqueeze_dim=1):
234
+ """Applies Rotary Position Embedding to the query and key tensors.
235
+
236
+ Args:
237
+ q (`torch.Tensor`): The query tensor.
238
+ k (`torch.Tensor`): The key tensor.
239
+ cos (`torch.Tensor`): The cosine part of the rotary embedding.
240
+ sin (`torch.Tensor`): The sine part of the rotary embedding.
241
+ position_ids (`torch.Tensor`):
242
+ The position indices of the tokens corresponding to the query and key tensors. For example, this can be
243
+ used to pass offsetted position ids when working with a KV-cache.
244
+ unsqueeze_dim (`int`, *optional*, defaults to 1):
245
+ The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
246
+ sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
247
+ that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
248
+ k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
249
+ cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
250
+ the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
251
+ Returns:
252
+ `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
253
+ """
254
+ cos = cos[position_ids].unsqueeze(unsqueeze_dim)
255
+ sin = sin[position_ids].unsqueeze(unsqueeze_dim)
256
+ q_embed = (q * cos) + (rotate_half(q) * sin)
257
+ k_embed = (k * cos) + (rotate_half(k) * sin)
258
+ return q_embed, k_embed
259
+
260
+
261
+ # Copied from transformers.models.llama.modeling_llama.repeat_kv
262
+ def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
263
+ """
264
+ This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
265
+ num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
266
+ """
267
+ batch, num_key_value_heads, slen, head_dim = hidden_states.shape
268
+ if n_rep == 1:
269
+ return hidden_states
270
+ hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
271
+ return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
272
+
273
+
274
+ # copied from transformers.models.mistral.modeling_mistral.MistralAttention with Mistral->Mixtral
275
+ # TODO @longjie no longer copied from Mistral after static cache
276
+ class TanukiAttention(nn.Module):
277
+ """
278
+ Multi-headed attention from 'Attention Is All You Need' paper. Modified to use sliding window attention: Longformer
279
+ and "Generating Long Sequences with Sparse Transformers".
280
+ """
281
+
282
+ def __init__(self, config: TanukiConfig, layer_idx: Optional[int] = None):
283
+ super().__init__()
284
+ self.config = config
285
+ self.layer_idx = layer_idx
286
+ if layer_idx is None:
287
+ logger.warning_once(
288
+ f"Instantiating {self.__class__.__name__} without passing a `layer_idx` is not recommended and will "
289
+ "lead to errors during the forward call if caching is used. Please make sure to provide a `layer_idx` "
290
+ "when creating this class."
291
+ )
292
+
293
+ self.hidden_size = config.hidden_size
294
+ self.num_heads = config.num_attention_heads
295
+ self.head_dim = self.hidden_size // self.num_heads
296
+ self.num_key_value_heads = config.num_key_value_heads
297
+ self.num_key_value_groups = self.num_heads // self.num_key_value_heads
298
+ self.max_position_embeddings = config.max_position_embeddings
299
+ self.rope_theta = config.rope_theta
300
+ self.is_causal = True
301
+ self.attention_dropout = config.attention_dropout
302
+
303
+ if (self.head_dim * self.num_heads) != self.hidden_size:
304
+ raise ValueError(
305
+ f"hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size}"
306
+ f" and `num_heads`: {self.num_heads})."
307
+ )
308
+ self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=False)
309
+ self.k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
310
+ self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
311
+ self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False)
312
+
313
+ self.rotary_emb = TanukiRotaryEmbedding(
314
+ self.head_dim,
315
+ max_position_embeddings=self.max_position_embeddings,
316
+ base=self.rope_theta,
317
+ )
318
+
319
+ def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int):
320
+ return tensor.view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2).contiguous()
321
+
322
+ def forward(
323
+ self,
324
+ hidden_states: torch.Tensor,
325
+ attention_mask: Optional[torch.Tensor] = None,
326
+ position_ids: Optional[torch.LongTensor] = None,
327
+ past_key_value: Optional[Cache] = None,
328
+ output_attentions: bool = False,
329
+ use_cache: bool = False,
330
+ cache_position: Optional[torch.LongTensor] = None,
331
+ ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
332
+ bsz, q_len, _ = hidden_states.size()
333
+
334
+ query_states = self.q_proj(hidden_states)
335
+ key_states = self.k_proj(hidden_states)
336
+ value_states = self.v_proj(hidden_states)
337
+
338
+ query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
339
+ key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
340
+ value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
341
+
342
+ kv_seq_len = key_states.shape[-2]
343
+ if past_key_value is not None:
344
+ if self.layer_idx is None:
345
+ raise ValueError(
346
+ f"The cache structure has changed since version v4.36. If you are using {self.__class__.__name__} "
347
+ "for auto-regressive decoding with k/v caching, please make sure to initialize the attention class "
348
+ "with a layer index."
349
+ )
350
+ kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx)
351
+ cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
352
+ query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)
353
+
354
+ if past_key_value is not None:
355
+ cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} # Specific to RoPE models
356
+ key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
357
+
358
+ # repeat k/v heads if n_kv_heads < n_heads
359
+ key_states = repeat_kv(key_states, self.num_key_value_groups)
360
+ value_states = repeat_kv(value_states, self.num_key_value_groups)
361
+
362
+ attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim)
363
+
364
+ if attn_weights.size() != (bsz, self.num_heads, q_len, kv_seq_len):
365
+ raise ValueError(
366
+ f"Attention weights should be of size {(bsz, self.num_heads, q_len, kv_seq_len)}, but is"
367
+ f" {attn_weights.size()}"
368
+ )
369
+
370
+ if attention_mask is not None: # no matter the length, we just slice it
371
+ causal_mask = attention_mask[:, :, :, : key_states.shape[-2]]
372
+ attn_weights = attn_weights + causal_mask
373
+
374
+ # upcast attention to fp32
375
+ attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype)
376
+ attn_weights = nn.functional.dropout(attn_weights, p=self.attention_dropout, training=self.training)
377
+ attn_output = torch.matmul(attn_weights, value_states)
378
+
379
+ if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim):
380
+ raise ValueError(
381
+ f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.head_dim)}, but is"
382
+ f" {attn_output.size()}"
383
+ )
384
+
385
+ attn_output = attn_output.transpose(1, 2).contiguous()
386
+ attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)
387
+
388
+ attn_output = self.o_proj(attn_output)
389
+
390
+ if not output_attentions:
391
+ attn_weights = None
392
+
393
+ return attn_output, attn_weights, past_key_value
394
+
395
+
396
+ # copied from transformers.models.mistral.modeling_mistral.MistralFlashAttention2 with Mistral->Mixtral
397
+ # TODO @longjie no longer copied from Mistral after static cache
398
+ class TanukiFlashAttention2(TanukiAttention):
399
+ """
400
+ Tanuki flash attention module. This module inherits from `TanukiAttention` as the weights of the module stays
401
+ untouched. The only required change would be on the forward pass where it needs to correctly call the public API of
402
+ flash attention and deal with padding tokens in case the input contains any of them.
403
+ """
404
+
405
+ # Copied from transformers.models.llama.modeling_llama.LlamaFlashAttention2.__init__
406
+ def __init__(self, *args, **kwargs):
407
+ super().__init__(*args, **kwargs)
408
+
409
+ # TODO: Should be removed once Flash Attention for RoCm is bumped to 2.1.
410
+ # flash_attn<2.1 generates top-left aligned causal mask, while what is needed here is bottom-right alignement, that was made default for flash_attn>=2.1. This attribute is used to handle this difference. Reference: https://github.com/Dao-AILab/flash-attention/releases/tag/v2.1.0.
411
+ # Beware that with flash_attn<2.1, using q_seqlen != k_seqlen (except for the case q_seqlen == 1) produces a wrong mask (top-left).
412
+ self._flash_attn_uses_top_left_mask = not is_flash_attn_greater_or_equal_2_10()
413
+
414
+ def forward(
415
+ self,
416
+ hidden_states: torch.Tensor,
417
+ attention_mask: Optional[torch.Tensor] = None,
418
+ position_ids: Optional[torch.LongTensor] = None,
419
+ past_key_value: Optional[Cache] = None,
420
+ output_attentions: bool = False,
421
+ use_cache: bool = False,
422
+ cache_position: Optional[torch.LongTensor] = None,
423
+ ):
424
+ bsz, q_len, _ = hidden_states.size()
425
+
426
+ query_states = self.q_proj(hidden_states)
427
+ key_states = self.k_proj(hidden_states)
428
+ value_states = self.v_proj(hidden_states)
429
+
430
+ query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
431
+ key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
432
+ value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
433
+
434
+ kv_seq_len = key_states.shape[-2]
435
+ if past_key_value is not None:
436
+ if self.layer_idx is None:
437
+ raise ValueError(
438
+ f"The cache structure has changed since version v4.36. If you are using {self.__class__.__name__} "
439
+ "for auto-regressive decoding with k/v caching, please make sure to initialize the attention class "
440
+ "with a layer index."
441
+ )
442
+ kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx)
443
+
444
+ # Because the input can be padded, the absolute sequence length depends on the max position id.
445
+ rotary_seq_len = max(kv_seq_len, position_ids[:, -1].max().item()) + 1
446
+ cos, sin = self.rotary_emb(value_states, seq_len=rotary_seq_len)
447
+
448
+ query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)
449
+
450
+ use_sliding_windows = (
451
+ _flash_supports_window_size
452
+ and getattr(self.config, "sliding_window", None) is not None
453
+ and kv_seq_len > self.config.sliding_window
454
+ )
455
+
456
+ if not _flash_supports_window_size:
457
+ logger.warning_once(
458
+ "The current flash attention version does not support sliding window attention, for a more memory efficient implementation"
459
+ " make sure to upgrade flash-attn library."
460
+ )
461
+
462
+ if past_key_value is not None:
463
+ # Activate slicing cache only if the config has a value `sliding_windows` attribute
464
+ cache_has_contents = past_key_value.get_seq_length(self.layer_idx) > 0
465
+ if (
466
+ getattr(self.config, "sliding_window", None) is not None
467
+ and kv_seq_len > self.config.sliding_window
468
+ and cache_has_contents
469
+ ):
470
+ slicing_tokens = 1 - self.config.sliding_window
471
+
472
+ past_key = past_key_value[self.layer_idx][0]
473
+ past_value = past_key_value[self.layer_idx][1]
474
+
475
+ past_key = past_key[:, :, slicing_tokens:, :].contiguous()
476
+ past_value = past_value[:, :, slicing_tokens:, :].contiguous()
477
+
478
+ if past_key.shape[-2] != self.config.sliding_window - 1:
479
+ raise ValueError(
480
+ f"past key must have a shape of (`batch_size, num_heads, self.config.sliding_window-1, head_dim`), got"
481
+ f" {past_key.shape}"
482
+ )
483
+
484
+ if attention_mask is not None:
485
+ attention_mask = attention_mask[:, slicing_tokens:]
486
+ attention_mask = torch.cat([attention_mask, torch.ones_like(attention_mask[:, -1:])], dim=-1)
487
+
488
+ cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} # Specific to RoPE models
489
+ key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
490
+
491
+ # repeat k/v heads if n_kv_heads < n_heads
492
+ key_states = repeat_kv(key_states, self.num_key_value_groups)
493
+ value_states = repeat_kv(value_states, self.num_key_value_groups)
494
+ dropout_rate = 0.0 if not self.training else self.attention_dropout
495
+
496
+ # In PEFT, usually we cast the layer norms in float32 for training stability reasons
497
+ # therefore the input hidden states gets silently casted in float32. Hence, we need
498
+ # cast them back in float16 just to be sure everything works as expected.
499
+ input_dtype = query_states.dtype
500
+ if input_dtype == torch.float32:
501
+ if torch.is_autocast_enabled():
502
+ target_dtype = torch.get_autocast_gpu_dtype()
503
+ # Handle the case where the model is quantized
504
+ elif hasattr(self.config, "_pre_quantization_dtype"):
505
+ target_dtype = self.config._pre_quantization_dtype
506
+ else:
507
+ target_dtype = self.q_proj.weight.dtype
508
+
509
+ logger.warning_once(
510
+ f"The input hidden states seems to be silently casted in float32, this might be related to"
511
+ f" the fact you have upcasted embedding or layer norm layers in float32. We will cast back the input in"
512
+ f" {target_dtype}."
513
+ )
514
+
515
+ query_states = query_states.to(target_dtype)
516
+ key_states = key_states.to(target_dtype)
517
+ value_states = value_states.to(target_dtype)
518
+
519
+ # Reashape to the expected shape for Flash Attention
520
+ query_states = query_states.transpose(1, 2)
521
+ key_states = key_states.transpose(1, 2)
522
+ value_states = value_states.transpose(1, 2)
523
+
524
+ attn_output = self._flash_attention_forward(
525
+ query_states,
526
+ key_states,
527
+ value_states,
528
+ attention_mask,
529
+ q_len,
530
+ dropout=dropout_rate,
531
+ use_sliding_windows=use_sliding_windows,
532
+ )
533
+
534
+ attn_output = attn_output.reshape(bsz, q_len, self.hidden_size).contiguous()
535
+ attn_output = self.o_proj(attn_output)
536
+
537
+ if not output_attentions:
538
+ attn_weights = None
539
+
540
+ return attn_output, attn_weights, past_key_value
541
+
542
+ def _flash_attention_forward(
543
+ self,
544
+ query_states,
545
+ key_states,
546
+ value_states,
547
+ attention_mask,
548
+ query_length,
549
+ dropout=0.0,
550
+ softmax_scale=None,
551
+ use_sliding_windows=False,
552
+ ):
553
+ """
554
+ Calls the forward method of Flash Attention - if the input hidden states contain at least one padding token
555
+ first unpad the input, then computes the attention scores and pad the final attention scores.
556
+
557
+ Args:
558
+ query_states (`torch.Tensor`):
559
+ Input query states to be passed to Flash Attention API
560
+ key_states (`torch.Tensor`):
561
+ Input key states to be passed to Flash Attention API
562
+ value_states (`torch.Tensor`):
563
+ Input value states to be passed to Flash Attention API
564
+ attention_mask (`torch.Tensor`):
565
+ The padding mask - corresponds to a tensor of size `(batch_size, seq_len)` where 0 stands for the
566
+ position of padding tokens and 1 for the position of non-padding tokens.
567
+ dropout (`float`):
568
+ Attention dropout
569
+ softmax_scale (`float`, *optional*):
570
+ The scaling of QK^T before applying softmax. Default to 1 / sqrt(head_dim)
571
+ use_sliding_windows (`bool`, *optional*):
572
+ Whether to activate sliding window attention.
573
+ """
574
+ if not self._flash_attn_uses_top_left_mask:
575
+ causal = self.is_causal
576
+ else:
577
+ # TODO: Remove the `query_length != 1` check once Flash Attention for RoCm is bumped to 2.1. For details, please see the comment in LlamaFlashAttention2 __init__.
578
+ causal = self.is_causal and query_length != 1
579
+
580
+ # Contains at least one padding token in the sequence
581
+ if attention_mask is not None:
582
+ batch_size = query_states.shape[0]
583
+ query_states, key_states, value_states, indices_q, cu_seq_lens, max_seq_lens = self._upad_input(
584
+ query_states, key_states, value_states, attention_mask, query_length
585
+ )
586
+
587
+ cu_seqlens_q, cu_seqlens_k = cu_seq_lens
588
+ max_seqlen_in_batch_q, max_seqlen_in_batch_k = max_seq_lens
589
+
590
+ if not use_sliding_windows:
591
+ attn_output_unpad = flash_attn_varlen_func(
592
+ query_states,
593
+ key_states,
594
+ value_states,
595
+ cu_seqlens_q=cu_seqlens_q,
596
+ cu_seqlens_k=cu_seqlens_k,
597
+ max_seqlen_q=max_seqlen_in_batch_q,
598
+ max_seqlen_k=max_seqlen_in_batch_k,
599
+ dropout_p=dropout,
600
+ softmax_scale=softmax_scale,
601
+ causal=causal,
602
+ )
603
+ else:
604
+ attn_output_unpad = flash_attn_varlen_func(
605
+ query_states,
606
+ key_states,
607
+ value_states,
608
+ cu_seqlens_q=cu_seqlens_q,
609
+ cu_seqlens_k=cu_seqlens_k,
610
+ max_seqlen_q=max_seqlen_in_batch_q,
611
+ max_seqlen_k=max_seqlen_in_batch_k,
612
+ dropout_p=dropout,
613
+ softmax_scale=softmax_scale,
614
+ causal=causal,
615
+ window_size=(self.config.sliding_window, self.config.sliding_window),
616
+ )
617
+
618
+ attn_output = pad_input(attn_output_unpad, indices_q, batch_size, query_length)
619
+ else:
620
+ if not use_sliding_windows:
621
+ attn_output = flash_attn_func(
622
+ query_states,
623
+ key_states,
624
+ value_states,
625
+ dropout,
626
+ softmax_scale=softmax_scale,
627
+ causal=causal,
628
+ )
629
+ else:
630
+ attn_output = flash_attn_func(
631
+ query_states,
632
+ key_states,
633
+ value_states,
634
+ dropout,
635
+ softmax_scale=softmax_scale,
636
+ causal=causal,
637
+ window_size=(self.config.sliding_window, self.config.sliding_window),
638
+ )
639
+
640
+ return attn_output
641
+
642
+ def _upad_input(self, query_layer, key_layer, value_layer, attention_mask, query_length):
643
+ batch_size, kv_seq_len, num_heads, head_dim = key_layer.shape
644
+
645
+ # On the first iteration we need to properly re-create the padding mask
646
+ # by slicing it on the proper place
647
+ if kv_seq_len != attention_mask.shape[-1]:
648
+ attention_mask_num_tokens = attention_mask.shape[-1]
649
+ attention_mask = attention_mask[:, attention_mask_num_tokens - kv_seq_len :]
650
+
651
+ indices_k, cu_seqlens_k, max_seqlen_in_batch_k = _get_unpad_data(attention_mask)
652
+
653
+ key_layer = index_first_axis(key_layer.reshape(batch_size * kv_seq_len, num_heads, head_dim), indices_k)
654
+ value_layer = index_first_axis(value_layer.reshape(batch_size * kv_seq_len, num_heads, head_dim), indices_k)
655
+
656
+ if query_length == kv_seq_len:
657
+ query_layer = index_first_axis(
658
+ query_layer.reshape(batch_size * kv_seq_len, num_heads, head_dim), indices_k
659
+ )
660
+ cu_seqlens_q = cu_seqlens_k
661
+ max_seqlen_in_batch_q = max_seqlen_in_batch_k
662
+ indices_q = indices_k
663
+ elif query_length == 1:
664
+ max_seqlen_in_batch_q = 1
665
+ cu_seqlens_q = torch.arange(
666
+ batch_size + 1, dtype=torch.int32, device=query_layer.device
667
+ ) # There is a memcpy here, that is very bad.
668
+ indices_q = cu_seqlens_q[:-1]
669
+ query_layer = query_layer.squeeze(1)
670
+ else:
671
+ # The -q_len: slice assumes left padding.
672
+ attention_mask = attention_mask[:, -query_length:]
673
+ query_layer, indices_q, cu_seqlens_q, max_seqlen_in_batch_q = unpad_input(query_layer, attention_mask)
674
+
675
+ return (
676
+ query_layer,
677
+ key_layer,
678
+ value_layer,
679
+ indices_q,
680
+ (cu_seqlens_q, cu_seqlens_k),
681
+ (max_seqlen_in_batch_q, max_seqlen_in_batch_k),
682
+ )
683
+
684
+
685
+ # copied from transformers.models.mistral.modeling_mistral.MistralSdpaAttention with Mistral->Mixtral
686
+ # TODO @longjie no longer copied from Mistral after static cache
687
+ class TanukiSdpaAttention(TanukiAttention):
688
+ """
689
+ Tanuki attention module using torch.nn.functional.scaled_dot_product_attention. This module inherits from
690
+ `TanukiAttention` as the weights of the module stays untouched. The only changes are on the forward pass to adapt to
691
+ SDPA API.
692
+ """
693
+
694
+ # Adapted from TanukiAttention.forward
695
+ def forward(
696
+ self,
697
+ hidden_states: torch.Tensor,
698
+ attention_mask: Optional[torch.Tensor] = None,
699
+ position_ids: Optional[torch.LongTensor] = None,
700
+ past_key_value: Optional[Cache] = None,
701
+ output_attentions: bool = False,
702
+ use_cache: bool = False,
703
+ cache_position: Optional[torch.LongTensor] = None,
704
+ ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
705
+ if output_attentions:
706
+ # TODO: Improve this warning with e.g. `model.config.attn_implementation = "manual"` once this is implemented.
707
+ logger.warning_once(
708
+ "TanukiModel is using TanukiSdpaAttention, but `torch.nn.functional.scaled_dot_product_attention` does not support `output_attentions=True`. Falling back to the manual attention implementation, "
709
+ 'but specifying the manual implementation will be required from Transformers version v5.0.0 onwards. This warning can be removed using the argument `attn_implementation="eager"` when loading the model.'
710
+ )
711
+ return super().forward(
712
+ hidden_states=hidden_states,
713
+ attention_mask=attention_mask,
714
+ position_ids=position_ids,
715
+ past_key_value=past_key_value,
716
+ output_attentions=output_attentions,
717
+ use_cache=use_cache,
718
+ )
719
+
720
+ bsz, q_len, _ = hidden_states.size()
721
+
722
+ query_states = self.q_proj(hidden_states)
723
+ key_states = self.k_proj(hidden_states)
724
+ value_states = self.v_proj(hidden_states)
725
+
726
+ query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
727
+ key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
728
+ value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
729
+
730
+ kv_seq_len = key_states.shape[-2]
731
+ if past_key_value is not None:
732
+ kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx)
733
+ cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
734
+
735
+ query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)
736
+
737
+ if past_key_value is not None:
738
+ cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} # Specific to RoPE models
739
+ key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
740
+
741
+ key_states = repeat_kv(key_states, self.num_key_value_groups)
742
+ value_states = repeat_kv(value_states, self.num_key_value_groups)
743
+
744
+ causal_mask = attention_mask
745
+ if attention_mask is not None: # no matter the length, we just slice it
746
+ causal_mask = attention_mask[:, :, :, : key_states.shape[-2]]
747
+
748
+ # SDPA with memory-efficient backend is currently (torch==2.1.2) bugged with non-contiguous inputs with custom attn_mask,
749
+ # Reference: https://github.com/pytorch/pytorch/issues/112577.
750
+ if query_states.device.type == "cuda" and attention_mask is not None:
751
+ query_states = query_states.contiguous()
752
+ key_states = key_states.contiguous()
753
+ value_states = value_states.contiguous()
754
+
755
+ # We dispatch to SDPA's Flash Attention or Efficient kernels via this `is_causal` if statement instead of an inline conditional assignment
756
+ # in SDPA to support both torch.compile's dynamic shapes and full graph options. An inline conditional prevents dynamic shapes from compiling.
757
+ # The q_len > 1 is necessary to match with AttentionMaskConverter.to_causal_4d that does not create a causal mask in case q_len == 1.
758
+ is_causal = True if causal_mask is None and q_len > 1 else False
759
+
760
+ attn_output = torch.nn.functional.scaled_dot_product_attention(
761
+ query_states,
762
+ key_states,
763
+ value_states,
764
+ attn_mask=causal_mask,
765
+ dropout_p=self.attention_dropout if self.training else 0.0,
766
+ is_causal=is_causal,
767
+ )
768
+
769
+ attn_output = attn_output.transpose(1, 2).contiguous()
770
+ attn_output = attn_output.view(bsz, q_len, self.hidden_size)
771
+
772
+ attn_output = self.o_proj(attn_output)
773
+
774
+ return attn_output, None, past_key_value
775
+
776
+
777
+ TANUKI_ATTENTION_CLASSES = {
778
+ "eager": TanukiAttention,
779
+ "flash_attention_2": TanukiFlashAttention2,
780
+ "sdpa": TanukiSdpaAttention,
781
+ }
782
+
783
+
784
+ class TanukiBlockSparseTop2MLP(nn.Module):
785
+ def __init__(self, config: TanukiConfig):
786
+ super().__init__()
787
+ self.ffn_dim = config.intermediate_size
788
+ self.hidden_dim = config.hidden_size
789
+
790
+ self.w1 = nn.Linear(self.hidden_dim, self.ffn_dim, bias=False)
791
+ self.w2 = nn.Linear(self.ffn_dim, self.hidden_dim, bias=False)
792
+ self.w3 = nn.Linear(self.hidden_dim, self.ffn_dim, bias=False)
793
+
794
+ self.act_fn = ACT2FN[config.hidden_act]
795
+
796
+ def forward(self, hidden_states):
797
+ current_hidden_states = self.act_fn(self.w1(hidden_states)) * self.w3(hidden_states)
798
+ current_hidden_states = self.w2(current_hidden_states)
799
+ return current_hidden_states
800
+
801
+
802
+ class TanukiSparseMoeBlock(nn.Module):
803
+ """
804
+ This implementation is
805
+ strictly equivalent to standard MoE with full capacity (no
806
+ dropped tokens). It's faster since it formulates MoE operations
807
+ in terms of block-sparse operations to accomodate imbalanced
808
+ assignments of tokens to experts, whereas standard MoE either
809
+ (1) drop tokens at the cost of reduced performance or (2) set
810
+ capacity factor to number of experts and thus waste computation
811
+ and memory on padding.
812
+ """
813
+
814
+ def __init__(self, config):
815
+ super().__init__()
816
+ self.hidden_dim = config.hidden_size
817
+ self.ffn_dim = config.intermediate_size
818
+ self.num_experts = config.num_local_experts
819
+ self.top_k = config.num_experts_per_tok
820
+
821
+ # gating
822
+ self.gate = nn.Linear(self.hidden_dim, self.num_experts, bias=False)
823
+
824
+ self.experts = nn.ModuleList([TanukiBlockSparseTop2MLP(config) for _ in range(self.num_experts)])
825
+
826
+ # Jitter parameters
827
+ self.jitter_noise = config.router_jitter_noise
828
+
829
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
830
+ """ """
831
+ batch_size, sequence_length, hidden_dim = hidden_states.shape
832
+ if self.training and self.jitter_noise > 0:
833
+ hidden_states *= torch.empty_like(hidden_states).uniform_(1.0 - self.jitter_noise, 1.0 + self.jitter_noise)
834
+ hidden_states = hidden_states.view(-1, hidden_dim)
835
+ # router_logits: (batch * sequence_length, n_experts)
836
+ router_logits = self.gate(hidden_states)
837
+
838
+ mean = router_logits.mean(dim=-1, keepdim=True)
839
+ std = router_logits.std(dim=-1, keepdim=True)
840
+ router_logits = (router_logits - mean) / (std + 1e-5)
841
+
842
+ routing_weights = F.softmax(router_logits, dim=1, dtype=torch.float)
843
+ routing_weights, selected_experts = torch.topk(routing_weights, self.top_k, dim=-1)
844
+ routing_weights /= routing_weights.sum(dim=-1, keepdim=True)
845
+ # we cast back to the input dtype
846
+ routing_weights = routing_weights.to(hidden_states.dtype)
847
+
848
+ final_hidden_states = torch.zeros(
849
+ (batch_size * sequence_length, hidden_dim), dtype=hidden_states.dtype, device=hidden_states.device
850
+ )
851
+
852
+ # One hot encode the selected experts to create an expert mask
853
+ # this will be used to easily index which expert is going to be sollicitated
854
+ expert_mask = torch.nn.functional.one_hot(selected_experts, num_classes=self.num_experts).permute(2, 1, 0)
855
+
856
+ # Loop over all available experts in the model and perform the computation on each expert
857
+ for expert_idx in range(self.num_experts):
858
+ expert_layer = self.experts[expert_idx]
859
+ idx, top_x = torch.where(expert_mask[expert_idx])
860
+
861
+ # Index the correct hidden states and compute the expert hidden state for
862
+ # the current expert. We need to make sure to multiply the output hidden
863
+ # states by `routing_weights` on the corresponding tokens (top-1 and top-2)
864
+ current_state = hidden_states[None, top_x].reshape(-1, hidden_dim)
865
+ current_hidden_states = expert_layer(current_state) * routing_weights[top_x, idx, None]
866
+
867
+ # However `index_add_` only support torch tensors for indexing so we'll use
868
+ # the `top_x` tensor here.
869
+ final_hidden_states.index_add_(0, top_x, current_hidden_states.to(hidden_states.dtype))
870
+ final_hidden_states = final_hidden_states.reshape(batch_size, sequence_length, hidden_dim)
871
+ return final_hidden_states, router_logits
872
+
873
+
874
+ class TanukiDecoderLayer(nn.Module):
875
+ def __init__(self, config: TanukiConfig, layer_idx: int):
876
+ super().__init__()
877
+ self.hidden_size = config.hidden_size
878
+
879
+ self.self_attn = TANUKI_ATTENTION_CLASSES[config._attn_implementation](config, layer_idx)
880
+
881
+ self.block_sparse_moe = TanukiSparseMoeBlock(config)
882
+ self.input_layernorm = TanukiRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
883
+ self.post_attention_layernorm = TanukiRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
884
+
885
+ def forward(
886
+ self,
887
+ hidden_states: torch.Tensor,
888
+ attention_mask: Optional[torch.Tensor] = None,
889
+ position_ids: Optional[torch.LongTensor] = None,
890
+ past_key_value: Optional[Tuple[torch.Tensor]] = None,
891
+ output_attentions: Optional[bool] = False,
892
+ output_router_logits: Optional[bool] = False,
893
+ use_cache: Optional[bool] = False,
894
+ cache_position: Optional[torch.LongTensor] = None,
895
+ **kwargs,
896
+ ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]:
897
+ """
898
+ Args:
899
+ hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`
900
+ attention_mask (`torch.FloatTensor`, *optional*): attention mask of size
901
+ `(batch, sequence_length)` where padding elements are indicated by 0.
902
+ past_key_value (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states
903
+ output_attentions (`bool`, *optional*):
904
+ Whether or not to return the attentions tensors of all attention layers. See `attentions` under
905
+ returned tensors for more detail.
906
+ output_router_logits (`bool`, *optional*):
907
+ Whether or not to return the logits of all the routers. They are useful for computing the router loss, and
908
+ should not be returned during inference.
909
+ use_cache (`bool`, *optional*):
910
+ If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding
911
+ (see `past_key_values`).
912
+ cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*):
913
+ Indices depicting the position of the input sequence tokens in the sequence.
914
+ kwargs (`dict`, *optional*):
915
+ Arbitrary kwargs to be ignored, used for FSDP and other methods that injects code
916
+ into the model
917
+ """
918
+
919
+ residual = hidden_states
920
+
921
+ hidden_states = self.input_layernorm(hidden_states)
922
+
923
+ # Self Attention
924
+ hidden_states, self_attn_weights, present_key_value = self.self_attn(
925
+ hidden_states=hidden_states,
926
+ attention_mask=attention_mask,
927
+ position_ids=position_ids,
928
+ past_key_value=past_key_value,
929
+ output_attentions=output_attentions,
930
+ use_cache=use_cache,
931
+ cache_position=cache_position,
932
+ )
933
+ hidden_states = residual + hidden_states
934
+
935
+ # Fully Connected
936
+ residual = hidden_states
937
+ hidden_states = self.post_attention_layernorm(hidden_states)
938
+ hidden_states, router_logits = self.block_sparse_moe(hidden_states)
939
+ hidden_states = residual + hidden_states
940
+
941
+ outputs = (hidden_states,)
942
+
943
+ if output_attentions:
944
+ outputs += (self_attn_weights,)
945
+
946
+ if use_cache:
947
+ outputs += (present_key_value,)
948
+
949
+ if output_router_logits:
950
+ outputs += (router_logits,)
951
+
952
+ return outputs
953
+
954
+ TANUKI_START_DOCSTRING = r"""
955
+ This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the
956
+ library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
957
+ etc.)
958
+
959
+ This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.
960
+ Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage
961
+ and behavior.
962
+
963
+ Parameters:
964
+ config ([`TanukiConfig`]):
965
+ Model configuration class with all the parameters of the model. Initializing with a config file does not
966
+ load the weights associated with the model, only the configuration. Check out the
967
+ [`~PreTrainedModel.from_pretrained`] method to load the model weights.
968
+ """
969
+
970
+
971
+ @add_start_docstrings(
972
+ "The bare Tanuki Model outputting raw hidden-states without any specific head on top.",
973
+ TANUKI_START_DOCSTRING,
974
+ )
975
+ # Copied from transformers.models.qwen2.modeling_qwen2.Qwen2PreTrainedModel with Qwen2->Mixtral
976
+ class TanukiPreTrainedModel(PreTrainedModel):
977
+ config_class = TanukiConfig
978
+ base_model_prefix = "model"
979
+ supports_gradient_checkpointing = True
980
+ _no_split_modules = ["TanukiDecoderLayer"]
981
+ _skip_keys_device_placement = "past_key_values"
982
+ _supports_flash_attn_2 = True
983
+ _supports_sdpa = True
984
+ _supports_cache_class = True
985
+
986
+ def _init_weights(self, module):
987
+ std = self.config.initializer_range
988
+ if isinstance(module, nn.Linear):
989
+ module.weight.data.normal_(mean=0.0, std=std)
990
+ if module.bias is not None:
991
+ module.bias.data.zero_()
992
+ elif isinstance(module, nn.Embedding):
993
+ module.weight.data.normal_(mean=0.0, std=std)
994
+ if module.padding_idx is not None:
995
+ module.weight.data[module.padding_idx].zero_()
996
+
997
+
998
+ TANUKI_INPUTS_DOCSTRING = r"""
999
+ Args:
1000
+ input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
1001
+ Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide
1002
+ it.
1003
+
1004
+ Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
1005
+ [`PreTrainedTokenizer.__call__`] for details.
1006
+
1007
+ [What are input IDs?](../glossary#input-ids)
1008
+ attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
1009
+ Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
1010
+
1011
+ - 1 for tokens that are **not masked**,
1012
+ - 0 for tokens that are **masked**.
1013
+
1014
+ [What are attention masks?](../glossary#attention-mask)
1015
+
1016
+ Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
1017
+ [`PreTrainedTokenizer.__call__`] for details.
1018
+
1019
+ If `past_key_values` is used, optionally only the last `decoder_input_ids` have to be input (see
1020
+ `past_key_values`).
1021
+
1022
+ If you want to change padding behavior, you should read [`modeling_opt._prepare_decoder_attention_mask`]
1023
+ and modify to your needs. See diagram 1 in [the paper](https://arxiv.org/abs/1910.13461) for more
1024
+ information on the default strategy.
1025
+
1026
+ - 1 indicates the head is **not masked**,
1027
+ - 0 indicates the head is **masked**.
1028
+ position_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
1029
+ Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,
1030
+ config.n_positions - 1]`.
1031
+
1032
+ [What are position IDs?](../glossary#position-ids)
1033
+ past_key_values (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):
1034
+ Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of shape
1035
+ `(batch_size, num_heads, sequence_length, embed_size_per_head)`) and 2 additional tensors of shape
1036
+ `(batch_size, num_heads, encoder_sequence_length, embed_size_per_head)`.
1037
+
1038
+ Contains pre-computed hidden-states (key and values in the self-attention blocks and in the cross-attention
1039
+ blocks) that can be used (see `past_key_values` input) to speed up sequential decoding.
1040
+
1041
+ If `past_key_values` are used, the user can optionally input only the last `decoder_input_ids` (those that
1042
+ don't have their past key value states given to this model) of shape `(batch_size, 1)` instead of all
1043
+ `decoder_input_ids` of shape `(batch_size, sequence_length)`.
1044
+ inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
1045
+ Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This
1046
+ is useful if you want more control over how to convert `input_ids` indices into associated vectors than the
1047
+ model's internal embedding lookup matrix.
1048
+ use_cache (`bool`, *optional*):
1049
+ If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see
1050
+ `past_key_values`).
1051
+ output_attentions (`bool`, *optional*):
1052
+ Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
1053
+ tensors for more detail.
1054
+ output_hidden_states (`bool`, *optional*):
1055
+ Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
1056
+ more detail.
1057
+ output_router_logits (`bool`, *optional*):
1058
+ Whether or not to return the logits of all the routers. They are useful for computing the router loss, and
1059
+ should not be returned during inference.
1060
+ return_dict (`bool`, *optional*):
1061
+ Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
1062
+ cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*):
1063
+ Indices depicting the position of the input sequence tokens in the sequence. Contrarily to `position_ids`,
1064
+ this tensor is not affected by padding. It is used to update the cache in the correct position and to infer
1065
+ the complete sequence length.
1066
+ """
1067
+
1068
+
1069
+ @add_start_docstrings(
1070
+ "The bare Tanuki Model outputting raw hidden-states without any specific head on top.",
1071
+ TANUKI_START_DOCSTRING,
1072
+ )
1073
+ # copied from transformers.models.mistral.modeling_mistral.MistralModel with MISTRAL->MIXTRAL,Mistral->Mixtral
1074
+ # TODO @longjie no longer copied from Mistral after static cache
1075
+ class TanukiModel(TanukiPreTrainedModel):
1076
+ """
1077
+ Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`TanukiDecoderLayer`]
1078
+
1079
+ Args:
1080
+ config: TanukiConfig
1081
+ """
1082
+
1083
+ def __init__(self, config: TanukiConfig):
1084
+ super().__init__(config)
1085
+ self.padding_idx = config.pad_token_id
1086
+ self.vocab_size = config.vocab_size
1087
+
1088
+ self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
1089
+ self.layers = nn.ModuleList(
1090
+ [TanukiDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
1091
+ )
1092
+ self._attn_implementation = config._attn_implementation
1093
+ self.norm = TanukiRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
1094
+
1095
+ self.gradient_checkpointing = False
1096
+ # Initialize weights and apply final processing
1097
+ self.post_init()
1098
+
1099
+ def get_input_embeddings(self):
1100
+ return self.embed_tokens
1101
+
1102
+ def set_input_embeddings(self, value):
1103
+ self.embed_tokens = value
1104
+
1105
+ # Ignore copy
1106
+ @add_start_docstrings_to_model_forward(TANUKI_INPUTS_DOCSTRING)
1107
+ def forward(
1108
+ self,
1109
+ input_ids: torch.LongTensor = None,
1110
+ attention_mask: Optional[torch.Tensor] = None,
1111
+ position_ids: Optional[torch.LongTensor] = None,
1112
+ past_key_values: Optional[List[torch.FloatTensor]] = None,
1113
+ inputs_embeds: Optional[torch.FloatTensor] = None,
1114
+ use_cache: Optional[bool] = None,
1115
+ output_attentions: Optional[bool] = None,
1116
+ output_hidden_states: Optional[bool] = None,
1117
+ output_router_logits: Optional[bool] = None,
1118
+ return_dict: Optional[bool] = None,
1119
+ cache_position: Optional[torch.LongTensor] = None,
1120
+ ) -> Union[Tuple, MoeModelOutputWithPast]:
1121
+ output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
1122
+ output_router_logits = (
1123
+ output_router_logits if output_router_logits is not None else self.config.output_router_logits
1124
+ )
1125
+ output_hidden_states = (
1126
+ output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
1127
+ )
1128
+ use_cache = use_cache if use_cache is not None else self.config.use_cache
1129
+
1130
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
1131
+
1132
+ if (input_ids is None) ^ (inputs_embeds is not None):
1133
+ raise ValueError(
1134
+ "You cannot specify both input_ids and inputs_embeds at the same time, and must specify either one"
1135
+ )
1136
+
1137
+ if self.gradient_checkpointing and self.training:
1138
+ if use_cache:
1139
+ logger.warning_once(
1140
+ "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
1141
+ )
1142
+ use_cache = False
1143
+
1144
+ use_legacy_cache = False
1145
+ if use_cache and not isinstance(past_key_values, Cache):
1146
+ use_legacy_cache = True
1147
+ past_key_values = DynamicCache.from_legacy_cache(past_key_values)
1148
+ logger.warning_once(
1149
+ "We detected that you are passing `past_key_values` as a tuple and this is deprecated and will be removed in v4.43. "
1150
+ "Please use an appropriate `Cache` class (https://huggingface.co/docs/transformers/v4.41.3/en/internal/generation_utils#transformers.Cache)"
1151
+ )
1152
+
1153
+ if inputs_embeds is None:
1154
+ inputs_embeds = self.embed_tokens(input_ids)
1155
+
1156
+ if cache_position is None:
1157
+ past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
1158
+ cache_position = torch.arange(
1159
+ past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device
1160
+ )
1161
+ if position_ids is None:
1162
+ position_ids = cache_position.unsqueeze(0)
1163
+
1164
+ causal_mask = self._update_causal_mask(
1165
+ attention_mask, inputs_embeds, cache_position, past_key_values, output_attentions
1166
+ )
1167
+
1168
+ hidden_states = inputs_embeds
1169
+
1170
+ # decoder layers
1171
+ all_hidden_states = () if output_hidden_states else None
1172
+ all_self_attns = () if output_attentions else None
1173
+ all_router_logits = () if output_router_logits else None
1174
+ next_decoder_cache = None
1175
+
1176
+ for decoder_layer in self.layers:
1177
+ if output_hidden_states:
1178
+ all_hidden_states += (hidden_states,)
1179
+
1180
+ if self.gradient_checkpointing and self.training:
1181
+ layer_outputs = self._gradient_checkpointing_func(
1182
+ decoder_layer.__call__,
1183
+ hidden_states,
1184
+ causal_mask,
1185
+ position_ids,
1186
+ past_key_values,
1187
+ output_attentions,
1188
+ output_router_logits,
1189
+ use_cache,
1190
+ cache_position,
1191
+ )
1192
+ else:
1193
+ layer_outputs = decoder_layer(
1194
+ hidden_states,
1195
+ attention_mask=causal_mask,
1196
+ position_ids=position_ids,
1197
+ past_key_value=past_key_values,
1198
+ output_attentions=output_attentions,
1199
+ output_router_logits=output_router_logits,
1200
+ use_cache=use_cache,
1201
+ cache_position=cache_position,
1202
+ )
1203
+
1204
+ hidden_states = layer_outputs[0]
1205
+
1206
+ if use_cache:
1207
+ next_decoder_cache = layer_outputs[2 if output_attentions else 1]
1208
+
1209
+ if output_attentions:
1210
+ all_self_attns += (layer_outputs[1],)
1211
+
1212
+ if output_router_logits:
1213
+ all_router_logits += (layer_outputs[-1],)
1214
+
1215
+ hidden_states = self.norm(hidden_states)
1216
+
1217
+ # add hidden states from the last decoder layer
1218
+ if output_hidden_states:
1219
+ all_hidden_states += (hidden_states,)
1220
+
1221
+ next_cache = None
1222
+ if use_cache:
1223
+ next_cache = next_decoder_cache.to_legacy_cache() if use_legacy_cache else next_decoder_cache
1224
+
1225
+ if not return_dict:
1226
+ return tuple(
1227
+ v
1228
+ for v in [hidden_states, next_cache, all_hidden_states, all_self_attns, all_router_logits]
1229
+ if v is not None
1230
+ )
1231
+ return MoeModelOutputWithPast(
1232
+ last_hidden_state=hidden_states,
1233
+ past_key_values=next_cache,
1234
+ hidden_states=all_hidden_states,
1235
+ attentions=all_self_attns,
1236
+ router_logits=all_router_logits,
1237
+ )
1238
+
1239
+ # Copied from transformers.models.llama.modeling_llama.LlamaModel._update_causal_mask
1240
+ def _update_causal_mask(
1241
+ self,
1242
+ attention_mask: torch.Tensor,
1243
+ input_tensor: torch.Tensor,
1244
+ cache_position: torch.Tensor,
1245
+ past_key_values: Cache,
1246
+ output_attentions: bool,
1247
+ ):
1248
+ # TODO: As of torch==2.2.0, the `attention_mask` passed to the model in `generate` is 2D and of dynamic length even when the static
1249
+ # KV cache is used. This is an issue for torch.compile which then recaptures cudagraphs at each decode steps due to the dynamic shapes.
1250
+ # (`recording cudagraph tree for symint key 13`, etc.), which is VERY slow. A workaround is `@torch.compiler.disable`, but this prevents using
1251
+ # `fullgraph=True`. See more context in https://github.com/huggingface/transformers/pull/29114
1252
+
1253
+ if self.config._attn_implementation == "flash_attention_2":
1254
+ if attention_mask is not None and 0.0 in attention_mask:
1255
+ return attention_mask
1256
+ return None
1257
+
1258
+ # For SDPA, when possible, we will rely on its `is_causal` argument instead of its `attn_mask` argument, in
1259
+ # order to dispatch on Flash Attention 2. This feature is not compatible with static cache, as SDPA will fail
1260
+ # to infer the attention mask.
1261
+ past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
1262
+ using_static_cache = isinstance(past_key_values, StaticCache)
1263
+
1264
+ # When output attentions is True, sdpa implementation's forward method calls the eager implementation's forward
1265
+ if self.config._attn_implementation == "sdpa" and not using_static_cache and not output_attentions:
1266
+ if AttentionMaskConverter._ignore_causal_mask_sdpa(
1267
+ attention_mask,
1268
+ inputs_embeds=input_tensor,
1269
+ past_key_values_length=past_seen_tokens,
1270
+ is_training=self.training,
1271
+ ):
1272
+ return None
1273
+
1274
+ dtype, device = input_tensor.dtype, input_tensor.device
1275
+ min_dtype = torch.finfo(dtype).min
1276
+ sequence_length = input_tensor.shape[1]
1277
+ if using_static_cache:
1278
+ target_length = past_key_values.get_max_length()
1279
+ else:
1280
+ target_length = (
1281
+ attention_mask.shape[-1]
1282
+ if isinstance(attention_mask, torch.Tensor)
1283
+ else past_seen_tokens + sequence_length + 1
1284
+ )
1285
+
1286
+ if attention_mask is not None and attention_mask.dim() == 4:
1287
+ # in this case we assume that the mask comes already in inverted form and requires no inversion or slicing
1288
+ if attention_mask.max() != 0:
1289
+ raise ValueError("Custom 4D attention mask should be passed in inverted form with max==0`")
1290
+ causal_mask = attention_mask
1291
+ else:
1292
+ causal_mask = torch.full(
1293
+ (sequence_length, target_length), fill_value=min_dtype, dtype=dtype, device=device
1294
+ )
1295
+ if sequence_length != 1:
1296
+ causal_mask = torch.triu(causal_mask, diagonal=1)
1297
+ causal_mask *= torch.arange(target_length, device=device) > cache_position.reshape(-1, 1)
1298
+ causal_mask = causal_mask[None, None, :, :].expand(input_tensor.shape[0], 1, -1, -1)
1299
+ if attention_mask is not None:
1300
+ causal_mask = causal_mask.clone() # copy to contiguous memory for in-place edit
1301
+ mask_length = attention_mask.shape[-1]
1302
+ padding_mask = causal_mask[:, :, :, :mask_length] + attention_mask[:, None, None, :]
1303
+ padding_mask = padding_mask == 0
1304
+ causal_mask[:, :, :, :mask_length] = causal_mask[:, :, :, :mask_length].masked_fill(
1305
+ padding_mask, min_dtype
1306
+ )
1307
+ if (
1308
+ self.config._attn_implementation == "sdpa"
1309
+ and attention_mask is not None
1310
+ and attention_mask.device.type == "cuda"
1311
+ and not output_attentions
1312
+ ):
1313
+ # Attend to all tokens in fully masked rows in the causal_mask, for example the relevant first rows when
1314
+ # using left padding. This is required by F.scaled_dot_product_attention memory-efficient attention path.
1315
+ # Details: https://github.com/pytorch/pytorch/issues/110213
1316
+ causal_mask = AttentionMaskConverter._unmask_unattended(causal_mask, min_dtype)
1317
+
1318
+ return causal_mask
1319
+
1320
+
1321
+ class TanukiForCausalLM(TanukiPreTrainedModel):
1322
+ _tied_weights_keys = ["lm_head.weight"]
1323
+
1324
+ def __init__(self, config):
1325
+ super().__init__(config)
1326
+ self.model = TanukiModel(config)
1327
+ self.vocab_size = config.vocab_size
1328
+ self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
1329
+ self.router_aux_loss_coef = config.router_aux_loss_coef
1330
+ self.num_experts = config.num_local_experts
1331
+ self.num_experts_per_tok = config.num_experts_per_tok
1332
+ # Initialize weights and apply final processing
1333
+ self.post_init()
1334
+
1335
+ def get_input_embeddings(self):
1336
+ return self.model.embed_tokens
1337
+
1338
+ def set_input_embeddings(self, value):
1339
+ self.model.embed_tokens = value
1340
+
1341
+ def get_output_embeddings(self):
1342
+ return self.lm_head
1343
+
1344
+ def set_output_embeddings(self, new_embeddings):
1345
+ self.lm_head = new_embeddings
1346
+
1347
+ def set_decoder(self, decoder):
1348
+ self.model = decoder
1349
+
1350
+ def get_decoder(self):
1351
+ return self.model
1352
+
1353
+ @add_start_docstrings_to_model_forward(TANUKI_INPUTS_DOCSTRING)
1354
+ @replace_return_docstrings(output_type=MoeCausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC)
1355
+ # Ignore copy
1356
+ def forward(
1357
+ self,
1358
+ input_ids: torch.LongTensor = None,
1359
+ attention_mask: Optional[torch.Tensor] = None,
1360
+ position_ids: Optional[torch.LongTensor] = None,
1361
+ past_key_values: Optional[List[torch.FloatTensor]] = None,
1362
+ inputs_embeds: Optional[torch.FloatTensor] = None,
1363
+ labels: Optional[torch.LongTensor] = None,
1364
+ use_cache: Optional[bool] = None,
1365
+ output_attentions: Optional[bool] = None,
1366
+ output_hidden_states: Optional[bool] = None,
1367
+ output_router_logits: Optional[bool] = None,
1368
+ return_dict: Optional[bool] = None,
1369
+ cache_position: Optional[torch.LongTensor] = None,
1370
+ ) -> Union[Tuple, MoeCausalLMOutputWithPast]:
1371
+ r"""
1372
+ Args:
1373
+ labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
1374
+ Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
1375
+ config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
1376
+ (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
1377
+
1378
+ Returns:
1379
+
1380
+ Example:
1381
+
1382
+ ```python
1383
+ >>> from transformers import AutoTokenizer, TanukiForCausalLM
1384
+
1385
+ >>> model = TanukiForCausalLM.from_pretrained("mistralai/Tanuki-8x7B-v0.1")
1386
+ >>> tokenizer = AutoTokenizer.from_pretrained("mistralai/Tanuki-8x7B-v0.1")
1387
+
1388
+ >>> prompt = "Hey, are you conscious? Can you talk to me?"
1389
+ >>> inputs = tokenizer(prompt, return_tensors="pt")
1390
+
1391
+ >>> # Generate
1392
+ >>> generate_ids = model.generate(inputs.input_ids, max_length=30)
1393
+ >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
1394
+ "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you."
1395
+ ```"""
1396
+
1397
+ output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
1398
+ output_router_logits = (
1399
+ output_router_logits if output_router_logits is not None else self.config.output_router_logits
1400
+ )
1401
+
1402
+ output_hidden_states = (
1403
+ output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
1404
+ )
1405
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
1406
+
1407
+ # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
1408
+ outputs = self.model(
1409
+ input_ids=input_ids,
1410
+ attention_mask=attention_mask,
1411
+ position_ids=position_ids,
1412
+ past_key_values=past_key_values,
1413
+ inputs_embeds=inputs_embeds,
1414
+ use_cache=use_cache,
1415
+ output_attentions=output_attentions,
1416
+ output_hidden_states=output_hidden_states,
1417
+ output_router_logits=output_router_logits,
1418
+ return_dict=return_dict,
1419
+ cache_position=cache_position,
1420
+ )
1421
+
1422
+ hidden_states = outputs[0]
1423
+ logits = self.lm_head(hidden_states)
1424
+ logits = logits.float()
1425
+
1426
+ loss = None
1427
+ if labels is not None:
1428
+ # Shift so that tokens < n predict n
1429
+ shift_logits = logits[..., :-1, :].contiguous()
1430
+ shift_labels = labels[..., 1:].contiguous()
1431
+ # Flatten the tokens
1432
+ loss_fct = CrossEntropyLoss()
1433
+ shift_logits = shift_logits.view(-1, self.config.vocab_size)
1434
+ shift_labels = shift_labels.view(-1)
1435
+ # Enable model parallelism
1436
+ shift_labels = shift_labels.to(shift_logits.device)
1437
+ loss = loss_fct(shift_logits, shift_labels)
1438
+
1439
+ aux_loss = None
1440
+ if output_router_logits:
1441
+ aux_loss = load_balancing_loss_func(
1442
+ outputs.router_logits if return_dict else outputs[-1],
1443
+ self.num_experts,
1444
+ self.num_experts_per_tok,
1445
+ attention_mask,
1446
+ )
1447
+ if labels is not None:
1448
+ loss += self.router_aux_loss_coef * aux_loss.to(loss.device) # make sure to reside in the same device
1449
+
1450
+ if not return_dict:
1451
+ output = (logits,) + outputs[1:]
1452
+ if output_router_logits:
1453
+ output = (aux_loss,) + output
1454
+ return (loss,) + output if loss is not None else output
1455
+
1456
+ return MoeCausalLMOutputWithPast(
1457
+ loss=loss,
1458
+ aux_loss=aux_loss,
1459
+ logits=logits,
1460
+ past_key_values=outputs.past_key_values,
1461
+ hidden_states=outputs.hidden_states,
1462
+ attentions=outputs.attentions,
1463
+ router_logits=outputs.router_logits,
1464
+ )
1465
+
1466
+ def prepare_inputs_for_generation(
1467
+ self,
1468
+ input_ids,
1469
+ past_key_values=None,
1470
+ attention_mask=None,
1471
+ inputs_embeds=None,
1472
+ output_router_logits=False,
1473
+ cache_position=None,
1474
+ use_cache=True,
1475
+ **kwargs,
1476
+ ):
1477
+ past_length = 0
1478
+ # Omit tokens covered by past_key_values
1479
+ if past_key_values is not None:
1480
+ # Past key values are always initialized with a `Cache` object -> no need for if-else anymore
1481
+ past_length = cache_position[0] if cache_position is not None else past_key_values.get_seq_length()
1482
+ max_cache_length = (
1483
+ torch.tensor(past_key_values.get_max_length(), device=input_ids.device)
1484
+ if past_key_values.get_max_length() is not None
1485
+ else None
1486
+ )
1487
+ cache_length = past_length if max_cache_length is None else torch.min(max_cache_length, past_length)
1488
+
1489
+ # Keep only the unprocessed tokens:
1490
+ # 1 - If the length of the attention_mask exceeds the length of input_ids, then we are in a setting where
1491
+ # some of the inputs are exclusively passed as part of the cache (e.g. when passing input_embeds as
1492
+ # input)
1493
+ if attention_mask is not None and attention_mask.shape[1] > input_ids.shape[1]:
1494
+ input_ids = input_ids[:, -(attention_mask.shape[1] - past_length) :]
1495
+ # 2 - If the past_length is smaller than input_ids', then input_ids holds all input tokens. We can discard
1496
+ # input_ids based on the past_length.
1497
+ elif past_length < input_ids.shape[1]:
1498
+ input_ids = input_ids[:, past_length:]
1499
+ # 3 - Otherwise (past_length >= input_ids.shape[1]), let's assume input_ids only has unprocessed tokens.
1500
+
1501
+ # If we are about to go beyond the maximum cache length, we need to crop the input attention mask.
1502
+ if (
1503
+ max_cache_length is not None
1504
+ and attention_mask is not None
1505
+ and cache_length + input_ids.shape[1] > max_cache_length
1506
+ ):
1507
+ attention_mask = attention_mask[:, -max_cache_length:]
1508
+
1509
+ position_ids = kwargs.get("position_ids", None)
1510
+ if attention_mask is not None and position_ids is None:
1511
+ # create position_ids on the fly for batch generation
1512
+ position_ids = attention_mask.long().cumsum(-1) - 1
1513
+ position_ids.masked_fill_(attention_mask == 0, 1)
1514
+ if past_key_values:
1515
+ position_ids = position_ids[:, -input_ids.shape[1] :]
1516
+
1517
+ # if `inputs_embeds` are passed, we only want to use them in the 1st generation step
1518
+ if inputs_embeds is not None and past_length == 0:
1519
+ model_inputs = {"inputs_embeds": inputs_embeds}
1520
+ else:
1521
+ model_inputs = {"input_ids": input_ids}
1522
+
1523
+ input_length = position_ids.shape[-1] if position_ids is not None else input_ids.shape[-1]
1524
+ if cache_position is None:
1525
+ cache_position = torch.arange(past_length, past_length + input_length, device=input_ids.device)
1526
+ elif use_cache:
1527
+ cache_position = cache_position[-input_length:]
1528
+
1529
+ model_inputs.update(
1530
+ {
1531
+ "position_ids": position_ids,
1532
+ "past_key_values": past_key_values,
1533
+ "use_cache": use_cache,
1534
+ "attention_mask": attention_mask,
1535
+ "output_router_logits": output_router_logits,
1536
+ "cache_position": cache_position,
1537
+ }
1538
+ )
1539
+ return model_inputs
1540
+
1541
+ @staticmethod
1542
+ def _reorder_cache(past_key_values, beam_idx):
1543
+ reordered_past = ()
1544
+ for layer_past in past_key_values:
1545
+ reordered_past += (
1546
+ tuple(past_state.index_select(0, beam_idx.to(past_state.device)) for past_state in layer_past),
1547
+ )
1548
+ return reordered_past
1549
+
1550
+
1551
+ @add_start_docstrings(
1552
+ """
1553
+ The Tanuki Model transformer with a sequence classification head on top (linear layer).
1554
+
1555
+ [`TanukiForSequenceClassification`] uses the last token in order to do the classification, as other causal models
1556
+ (e.g. GPT-2) do.
1557
+
1558
+ Since it does classification on the last token, it requires to know the position of the last token. If a
1559
+ `pad_token_id` is defined in the configuration, it finds the last token that is not a padding token in each row. If
1560
+ no `pad_token_id` is defined, it simply takes the last value in each row of the batch. Since it cannot guess the
1561
+ padding tokens when `inputs_embeds` are passed instead of `input_ids`, it does the same (take the last value in
1562
+ each row of the batch).
1563
+ """,
1564
+ TANUKI_START_DOCSTRING,
1565
+ )
1566
+ # Copied from transformers.models.llama.modeling_llama.LlamaForSequenceClassification with Llama->Mixtral, LLAMA->MIXTRAL
1567
+ class TanukiForSequenceClassification(TanukiPreTrainedModel):
1568
+ def __init__(self, config):
1569
+ super().__init__(config)
1570
+ self.num_labels = config.num_labels
1571
+ self.model = TanukiModel(config)
1572
+ self.score = nn.Linear(config.hidden_size, self.num_labels, bias=False)
1573
+
1574
+ # Initialize weights and apply final processing
1575
+ self.post_init()
1576
+
1577
+ def get_input_embeddings(self):
1578
+ return self.model.embed_tokens
1579
+
1580
+ def set_input_embeddings(self, value):
1581
+ self.model.embed_tokens = value
1582
+
1583
+ @add_start_docstrings_to_model_forward(TANUKI_INPUTS_DOCSTRING)
1584
+ def forward(
1585
+ self,
1586
+ input_ids: torch.LongTensor = None,
1587
+ attention_mask: Optional[torch.Tensor] = None,
1588
+ position_ids: Optional[torch.LongTensor] = None,
1589
+ past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = None,
1590
+ inputs_embeds: Optional[torch.FloatTensor] = None,
1591
+ labels: Optional[torch.LongTensor] = None,
1592
+ use_cache: Optional[bool] = None,
1593
+ output_attentions: Optional[bool] = None,
1594
+ output_hidden_states: Optional[bool] = None,
1595
+ return_dict: Optional[bool] = None,
1596
+ ) -> Union[Tuple, SequenceClassifierOutputWithPast]:
1597
+ r"""
1598
+ labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
1599
+ Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
1600
+ config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
1601
+ `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
1602
+ """
1603
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
1604
+
1605
+ transformer_outputs = self.model(
1606
+ input_ids,
1607
+ attention_mask=attention_mask,
1608
+ position_ids=position_ids,
1609
+ past_key_values=past_key_values,
1610
+ inputs_embeds=inputs_embeds,
1611
+ use_cache=use_cache,
1612
+ output_attentions=output_attentions,
1613
+ output_hidden_states=output_hidden_states,
1614
+ return_dict=return_dict,
1615
+ )
1616
+ hidden_states = transformer_outputs[0]
1617
+ logits = self.score(hidden_states)
1618
+
1619
+ if input_ids is not None:
1620
+ batch_size = input_ids.shape[0]
1621
+ else:
1622
+ batch_size = inputs_embeds.shape[0]
1623
+
1624
+ if self.config.pad_token_id is None and batch_size != 1:
1625
+ raise ValueError("Cannot handle batch sizes > 1 if no padding token is defined.")
1626
+ if self.config.pad_token_id is None:
1627
+ sequence_lengths = -1
1628
+ else:
1629
+ if input_ids is not None:
1630
+ # if no pad token found, use modulo instead of reverse indexing for ONNX compatibility
1631
+ sequence_lengths = torch.eq(input_ids, self.config.pad_token_id).int().argmax(-1) - 1
1632
+ sequence_lengths = sequence_lengths % input_ids.shape[-1]
1633
+ sequence_lengths = sequence_lengths.to(logits.device)
1634
+ else:
1635
+ sequence_lengths = -1
1636
+
1637
+ pooled_logits = logits[torch.arange(batch_size, device=logits.device), sequence_lengths]
1638
+
1639
+ loss = None
1640
+ if labels is not None:
1641
+ labels = labels.to(logits.device)
1642
+ if self.config.problem_type is None:
1643
+ if self.num_labels == 1:
1644
+ self.config.problem_type = "regression"
1645
+ elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
1646
+ self.config.problem_type = "single_label_classification"
1647
+ else:
1648
+ self.config.problem_type = "multi_label_classification"
1649
+
1650
+ if self.config.problem_type == "regression":
1651
+ loss_fct = MSELoss()
1652
+ if self.num_labels == 1:
1653
+ loss = loss_fct(pooled_logits.squeeze(), labels.squeeze())
1654
+ else:
1655
+ loss = loss_fct(pooled_logits, labels)
1656
+ elif self.config.problem_type == "single_label_classification":
1657
+ loss_fct = CrossEntropyLoss()
1658
+ loss = loss_fct(pooled_logits.view(-1, self.num_labels), labels.view(-1))
1659
+ elif self.config.problem_type == "multi_label_classification":
1660
+ loss_fct = BCEWithLogitsLoss()
1661
+ loss = loss_fct(pooled_logits, labels)
1662
+ if not return_dict:
1663
+ output = (pooled_logits,) + transformer_outputs[1:]
1664
+ return ((loss,) + output) if loss is not None else output
1665
+
1666
+ return SequenceClassifierOutputWithPast(
1667
+ loss=loss,
1668
+ logits=pooled_logits,
1669
+ past_key_values=transformer_outputs.past_key_values,
1670
+ hidden_states=transformer_outputs.hidden_states,
1671
+ attentions=transformer_outputs.attentions,
1672
+ )
1673
+
1674
+
1675
+ @add_start_docstrings(
1676
+ """
1677
+ The Tanuki Model transformer with a token classification head on top (a linear layer on top of the hidden-states
1678
+ output) e.g. for Named-Entity-Recognition (NER) tasks.
1679
+ """,
1680
+ TANUKI_START_DOCSTRING,
1681
+ )
1682
+ # Copied from transformers.models.llama.modeling_llama.LlamaForTokenClassification with Llama->Mixtral, LLAMA->MIXTRAL
1683
+ class TanukiForTokenClassification(TanukiPreTrainedModel):
1684
+ def __init__(self, config):
1685
+ super().__init__(config)
1686
+ self.num_labels = config.num_labels
1687
+ self.model = TanukiModel(config)
1688
+ if getattr(config, "classifier_dropout", None) is not None:
1689
+ classifier_dropout = config.classifier_dropout
1690
+ elif getattr(config, "hidden_dropout", None) is not None:
1691
+ classifier_dropout = config.hidden_dropout
1692
+ else:
1693
+ classifier_dropout = 0.1
1694
+ self.dropout = nn.Dropout(classifier_dropout)
1695
+ self.score = nn.Linear(config.hidden_size, config.num_labels)
1696
+
1697
+ # Initialize weights and apply final processing
1698
+ self.post_init()
1699
+
1700
+ def get_input_embeddings(self):
1701
+ return self.model.embed_tokens
1702
+
1703
+ def set_input_embeddings(self, value):
1704
+ self.model.embed_tokens = value
1705
+
1706
+ @add_start_docstrings_to_model_forward(TANUKI_INPUTS_DOCSTRING)
1707
+ def forward(
1708
+ self,
1709
+ input_ids: Optional[torch.LongTensor] = None,
1710
+ attention_mask: Optional[torch.Tensor] = None,
1711
+ position_ids: Optional[torch.LongTensor] = None,
1712
+ past_key_values: Optional[List[torch.FloatTensor]] = None,
1713
+ inputs_embeds: Optional[torch.FloatTensor] = None,
1714
+ labels: Optional[torch.LongTensor] = None,
1715
+ use_cache: Optional[bool] = None,
1716
+ output_attentions: Optional[bool] = None,
1717
+ output_hidden_states: Optional[bool] = None,
1718
+ return_dict: Optional[bool] = None,
1719
+ ) -> Union[Tuple, TokenClassifierOutput]:
1720
+ r"""
1721
+ labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
1722
+ Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
1723
+ config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
1724
+ `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
1725
+ """
1726
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
1727
+
1728
+ outputs = self.model(
1729
+ input_ids,
1730
+ attention_mask=attention_mask,
1731
+ position_ids=position_ids,
1732
+ past_key_values=past_key_values,
1733
+ inputs_embeds=inputs_embeds,
1734
+ use_cache=use_cache,
1735
+ output_attentions=output_attentions,
1736
+ output_hidden_states=output_hidden_states,
1737
+ return_dict=return_dict,
1738
+ )
1739
+ sequence_output = outputs[0]
1740
+ sequence_output = self.dropout(sequence_output)
1741
+ logits = self.score(sequence_output)
1742
+
1743
+ loss = None
1744
+ if labels is not None:
1745
+ loss_fct = CrossEntropyLoss()
1746
+ loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
1747
+
1748
+ if not return_dict:
1749
+ output = (logits,) + outputs[2:]
1750
+ return ((loss,) + output) if loss is not None else output
1751
+
1752
+ return TokenClassifierOutput(
1753
+ loss=loss,
1754
+ logits=logits,
1755
+ hidden_states=outputs.hidden_states,
1756
+ attentions=outputs.attentions,
1757
+ )
special_tokens_map.json ADDED
@@ -0,0 +1,45 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "bos_token": {
3
+ "content": "<s>",
4
+ "lstrip": false,
5
+ "normalized": false,
6
+ "rstrip": false,
7
+ "single_word": false
8
+ },
9
+ "cls_token": {
10
+ "content": "<CLS>",
11
+ "lstrip": false,
12
+ "normalized": false,
13
+ "rstrip": false,
14
+ "single_word": false
15
+ },
16
+ "eos_token": {
17
+ "content": "</s>",
18
+ "lstrip": false,
19
+ "normalized": false,
20
+ "rstrip": false,
21
+ "single_word": false
22
+ },
23
+ "mask_token": {
24
+ "content": "<MASK>",
25
+ "lstrip": false,
26
+ "normalized": false,
27
+ "rstrip": false,
28
+ "single_word": false
29
+ },
30
+ "pad_token": "<unk>",
31
+ "sep_token": {
32
+ "content": "<SEP>",
33
+ "lstrip": false,
34
+ "normalized": false,
35
+ "rstrip": false,
36
+ "single_word": false
37
+ },
38
+ "unk_token": {
39
+ "content": "<unk>",
40
+ "lstrip": false,
41
+ "normalized": false,
42
+ "rstrip": false,
43
+ "single_word": false
44
+ }
45
+ }
tokenizer.json ADDED
The diff for this file is too large to render. See raw diff
 
tokenizer_config.json ADDED
@@ -0,0 +1,127 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "added_tokens_decoder": {
3
+ "0": {
4
+ "content": "<unk>",
5
+ "lstrip": false,
6
+ "normalized": false,
7
+ "rstrip": false,
8
+ "single_word": false,
9
+ "special": true
10
+ },
11
+ "1": {
12
+ "content": "<s>",
13
+ "lstrip": false,
14
+ "normalized": false,
15
+ "rstrip": false,
16
+ "single_word": false,
17
+ "special": true
18
+ },
19
+ "2": {
20
+ "content": "</s>",
21
+ "lstrip": false,
22
+ "normalized": false,
23
+ "rstrip": false,
24
+ "single_word": false,
25
+ "special": true
26
+ },
27
+ "3": {
28
+ "content": "<pad>",
29
+ "lstrip": false,
30
+ "normalized": false,
31
+ "rstrip": false,
32
+ "single_word": false,
33
+ "special": true
34
+ },
35
+ "4": {
36
+ "content": "<CLS>",
37
+ "lstrip": false,
38
+ "normalized": false,
39
+ "rstrip": false,
40
+ "single_word": false,
41
+ "special": true
42
+ },
43
+ "5": {
44
+ "content": "<SEP>",
45
+ "lstrip": false,
46
+ "normalized": false,
47
+ "rstrip": false,
48
+ "single_word": false,
49
+ "special": true
50
+ },
51
+ "6": {
52
+ "content": "<EOD>",
53
+ "lstrip": false,
54
+ "normalized": false,
55
+ "rstrip": false,
56
+ "single_word": false,
57
+ "special": true
58
+ },
59
+ "7": {
60
+ "content": "<MASK>",
61
+ "lstrip": false,
62
+ "normalized": false,
63
+ "rstrip": false,
64
+ "single_word": false,
65
+ "special": true
66
+ },
67
+ "9": {
68
+ "content": "<llm-code>",
69
+ "lstrip": false,
70
+ "normalized": false,
71
+ "rstrip": false,
72
+ "single_word": false,
73
+ "special": true
74
+ },
75
+ "10": {
76
+ "content": "</llm-code>",
77
+ "lstrip": false,
78
+ "normalized": false,
79
+ "rstrip": false,
80
+ "single_word": false,
81
+ "special": true
82
+ },
83
+ "11": {
84
+ "content": "<llm-code-output>",
85
+ "lstrip": false,
86
+ "normalized": false,
87
+ "rstrip": false,
88
+ "single_word": false,
89
+ "special": true
90
+ },
91
+ "12": {
92
+ "content": "</llm-code-output>",
93
+ "lstrip": false,
94
+ "normalized": false,
95
+ "rstrip": false,
96
+ "single_word": false,
97
+ "special": true
98
+ },
99
+ "65000": {
100
+ "content": "[PAD]",
101
+ "lstrip": false,
102
+ "normalized": false,
103
+ "rstrip": false,
104
+ "single_word": false,
105
+ "special": true
106
+ }
107
+ },
108
+ "additional_special_tokens": [],
109
+ "bos_token": "<s>",
110
+ "chat_template": "{% for message in messages %}{% if message['role'] == 'user' %}{{ '\n\n### 指示:\n' + message['content'] }}{% elif message['role'] == 'system' %}{{ message['content'] }}{% elif message['role'] == 'assistant' %}{{ '\n\n### 応答:\n' + message['content'] + eos_token }}{% endif %}{% if loop.last and add_generation_prompt %}{{ '\n\n### 応答:\n' }}{% endif %}{% endfor %}",
111
+ "clean_up_tokenization_spaces": false,
112
+ "cls_token": "<CLS>",
113
+ "eos_token": "</s>",
114
+ "extra_ids": 0,
115
+ "mask_token": "<MASK>",
116
+ "max_length": 1000000000000000019884624838656,
117
+ "model_max_length": 8192,
118
+ "pad_token": "<unk>",
119
+ "padding_side": "right",
120
+ "sep_token": "<SEP>",
121
+ "sp_model_kwargs": {},
122
+ "stride": 0,
123
+ "tokenizer_class": "PreTrainedTokenizerFast",
124
+ "truncation_side": "right",
125
+ "truncation_strategy": "longest_first",
126
+ "unk_token": "<unk>"
127
+ }