jon-tow commited on
Commit
9c915ec
•
1 Parent(s): 5d3a767

refactor: use `transformers.StableLm` impl

Browse files
arcade100k.tiktoken DELETED
The diff for this file is too large to render. See raw diff
 
config.json CHANGED
@@ -1,33 +1,36 @@
1
  {
2
- "_name_or_path": "StableLM 2 12B Chat",
3
  "architectures": [
4
- "StableLMEpochForCausalLM"
5
  ],
6
  "attention_dropout": 0.0,
7
  "auto_map": {
8
- "AutoConfig": "configuration_stablelm_epoch.StableLMEpochConfig",
9
- "AutoModelForCausalLM": "modeling_stablelm_epoch_2.StableLMEpochForCausalLM"
10
  },
11
- "bos_token_id": null,
12
  "eos_token_id": 100257,
13
  "hidden_act": "silu",
 
14
  "hidden_size": 5120,
15
  "initializer_range": 0.01,
16
  "intermediate_size": 13824,
 
17
  "max_position_embeddings": 4096,
18
- "model_type": "stablelm_epoch",
19
- "norm_eps": 1e-05,
20
  "num_attention_heads": 32,
21
  "num_hidden_layers": 40,
22
  "num_key_value_heads": 8,
23
- "rope_pct": 0.25,
 
 
24
  "rope_theta": 10000,
25
  "rotary_scaling_factor": 1.0,
26
  "tie_word_embeddings": false,
27
- "torch_dtype": "float32",
28
- "transformers_version": "4.38.2",
29
  "use_cache": true,
30
  "use_norm_bias": false,
 
31
  "use_qkv_bias": false,
32
  "vocab_size": 100352
33
  }
 
1
  {
 
2
  "architectures": [
3
+ "StableLmForCausalLM"
4
  ],
5
  "attention_dropout": 0.0,
6
  "auto_map": {
7
+ "AutoConfig": "configuration_stablelm.StableLmConfig",
8
+ "AutoModelForCausalLM": "modeling_stablelm.StableLmForCausalLM"
9
  },
10
+ "bos_token_id": 100257,
11
  "eos_token_id": 100257,
12
  "hidden_act": "silu",
13
+ "hidden_dropout": 0.0,
14
  "hidden_size": 5120,
15
  "initializer_range": 0.01,
16
  "intermediate_size": 13824,
17
+ "layer_norm_eps": 1e-05,
18
  "max_position_embeddings": 4096,
19
+ "model_type": "stablelm",
 
20
  "num_attention_heads": 32,
21
  "num_hidden_layers": 40,
22
  "num_key_value_heads": 8,
23
+ "partial_rotary_factor": 0.25,
24
+ "qk_layernorm": true,
25
+ "rope_scaling": null,
26
  "rope_theta": 10000,
27
  "rotary_scaling_factor": 1.0,
28
  "tie_word_embeddings": false,
29
+ "torch_dtype": "bfloat16",
30
+ "transformers_version": "4.39.0.dev0",
31
  "use_cache": true,
32
  "use_norm_bias": false,
33
+ "use_parallel_residual": true,
34
  "use_qkv_bias": false,
35
  "vocab_size": 100352
36
  }
configuration_stablelm_epoch.py → configuration_stablelm.py RENAMED
@@ -1,4 +1,4 @@
1
- # Copyright 2023 Stability and The HuggingFace Inc. team. All rights reserved.
2
  #
3
  # Licensed under the Apache License, Version 2.0 (the "License");
4
  # you may not use this file except in compliance with the License.
@@ -11,32 +11,45 @@
11
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
  # See the License for the specific language governing permissions and
13
  # limitations under the License.
14
- """ StableLM Epoch model configuration"""
15
- from transformers import PretrainedConfig
 
16
  from transformers.utils import logging
17
 
18
 
19
  logger = logging.get_logger(__name__)
20
 
 
 
 
 
 
21
 
22
- class StableLMEpochConfig(PretrainedConfig):
23
  r"""
24
- Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
25
- documentation from [`PretrainedConfig`] for more information.
 
 
 
 
 
 
 
26
 
27
  Args:
28
- vocab_size (`int`, *optional*, defaults to 50_304):
29
  Vocabulary size of the StableLM model. Defines the number of different tokens that
30
- can be represented by the `inputs_ids` passed when calling [`StableLMEpochModel`].
31
  intermediate_size (`int`, *optional*, defaults to 6912):
32
  Dimension of the MLP representations.
33
  hidden_size (`int`, *optional*, defaults to 2560):
34
- Dimension of the decoder layers and the pooler layer.
35
  num_hidden_layers (`int`, *optional*, defaults to 32):
36
  Number of hidden layers in the Transformer decoder.
37
  num_attention_heads (`int`, *optional*, defaults to 32):
38
  Number of attention heads for each attention layer in the Transformer encoder.
39
- num_key_value_heads (`int`, *optional*):
40
  This is the number of key_value heads that should be used to implement Grouped Query Attention. If
41
  `num_key_value_heads=num_attention_heads`, the model will use Multi Head Attention (MHA), if
42
  `num_key_value_heads=1 the model will use Multi Query Attention (MQA) otherwise GQA is used. When
@@ -46,72 +59,139 @@ class StableLMEpochConfig(PretrainedConfig):
46
  `num_attention_heads`.
47
  hidden_act (`str` or `function`, *optional*, defaults to `"silu"`):
48
  The non-linear activation function (function or string).
49
- rope_pct (`float`, *optional*, defaults to 1.0):
50
- Percentage of hidden dimensions to allocate to rotary embeddings.
51
- rope_theta (`float`, *optional*, defaults to 10000.0):
52
- The base period of the RoPE embeddings.
53
- max_position_embeddings (`int`, *optional*, defaults to 2048):
54
  The maximum sequence length that this model might ever be used with.
55
  Typically set this to something large just in case (e.g., 512 or 1024 or 2048).
56
- initializer_range (`float`, *optional*, defaults to 1e-5):
57
  The standard deviation of the truncated_normal_initializer for initializing
58
  all weight matrices.
59
- norm_eps (`float`, *optional*, defaults to 1e-8):
60
  The epsilon used by the normalization layers.
61
  use_cache (`bool`, *optional*, defaults to `True`):
62
  Whether or not the model should return the last key/values attentions
63
  (not used by all models). Only relevant if `config.is_decoder=True`.
64
- use_qkv_bias (`bool`, *optional*, defaults to `True`):
 
 
 
 
 
 
 
 
 
 
 
 
65
  Whether or not the model should use bias for qkv layers.
66
- tie_word_embeddings(`bool`, *optional*, defaults to `False`):
67
- Whether to tie weight embeddings
 
 
 
 
 
68
  attention_dropout (`float`, *optional*, defaults to 0.0):
69
  The dropout ratio for the attention probabilities.
70
- """
71
- model_type = "stablelm_epoch"
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
72
  keys_to_ignore_at_inference = ["past_key_values"]
73
 
74
  def __init__(
75
  self,
76
- vocab_size=50_304,
77
  intermediate_size=6912,
78
  hidden_size=2560,
79
  num_hidden_layers=32,
80
  num_attention_heads=32,
81
  num_key_value_heads=32,
82
  hidden_act="silu",
83
- rope_pct=0.25,
84
- rope_theta=10_000,
85
  max_position_embeddings=4096,
86
  initializer_range=0.02,
87
- norm_eps=1.0e-5,
88
  use_cache=True,
89
- use_qkv_bias=True,
90
- bos_token_id=0,
91
- eos_token_id=2,
92
  tie_word_embeddings=False,
93
- attention_dropout: float = 0.0,
 
 
 
 
 
 
 
 
 
94
  **kwargs,
95
  ):
96
  self.vocab_size = vocab_size
97
  self.max_position_embeddings = max_position_embeddings
98
- self.intermediate_size = intermediate_size
99
  self.hidden_size = hidden_size
 
100
  self.num_hidden_layers = num_hidden_layers
101
  self.num_attention_heads = num_attention_heads
102
  self.num_key_value_heads = num_key_value_heads
103
  self.hidden_act = hidden_act
104
- self.rope_pct = rope_pct
105
- self.rope_theta = rope_theta
106
  self.initializer_range = initializer_range
107
- self.norm_eps = norm_eps
108
  self.use_cache = use_cache
 
 
109
  self.use_qkv_bias = use_qkv_bias
110
- self.tie_word_embeddings = tie_word_embeddings
 
 
111
  self.attention_dropout = attention_dropout
 
 
 
112
  super().__init__(
113
  bos_token_id=bos_token_id,
114
  eos_token_id=eos_token_id,
115
  tie_word_embeddings=tie_word_embeddings,
116
  **kwargs,
117
  )
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright 2024 Stability AI and The HuggingFace Inc. team. All rights reserved.
2
  #
3
  # Licensed under the Apache License, Version 2.0 (the "License");
4
  # you may not use this file except in compliance with the License.
 
11
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
  # See the License for the specific language governing permissions and
13
  # limitations under the License.
14
+ """ StableLM model configuration """
15
+
16
+ from transformers.configuration_utils import PretrainedConfig
17
  from transformers.utils import logging
18
 
19
 
20
  logger = logging.get_logger(__name__)
21
 
22
+ STABLELM_PRETRAINED_CONFIG_ARCHIVE_MAP = {
23
+ "stabilityai/stablelm-3b-4e1t": "https://huggingface.co/stabilityai/stablelm-3b-4e1t/resolve/main/config.json",
24
+ # See all StableLM models at https://huggingface.co/models?filter=stablelm
25
+ }
26
+
27
 
28
+ class StableLmConfig(PretrainedConfig):
29
  r"""
30
+ This is the configuration class to store the configuration of a [`~StableLmModel`].
31
+ It is used to instantiate an StableLM model according to the specified arguments, defining the model
32
+ architecture. Instantiating a configuration with the defaults will yield a similar configuration to that of
33
+ the StableLM [stabilityai/stablelm-3b-4e1t](https://huggingface.co/stabilityai/stablelm-3b-4e1t) architecture.
34
+
35
+ Configuration objects inherit from [`PretrainedConfig`] and can be used
36
+ to control the model outputs. Read the documentation from [`PretrainedConfig`]
37
+ for more information.
38
+
39
 
40
  Args:
41
+ vocab_size (`int`, *optional*, defaults to 50304):
42
  Vocabulary size of the StableLM model. Defines the number of different tokens that
43
+ can be represented by the `inputs_ids` passed when calling [`StableLmModel`].
44
  intermediate_size (`int`, *optional*, defaults to 6912):
45
  Dimension of the MLP representations.
46
  hidden_size (`int`, *optional*, defaults to 2560):
47
+ Number of hidden layers in the Transformer decoder.
48
  num_hidden_layers (`int`, *optional*, defaults to 32):
49
  Number of hidden layers in the Transformer decoder.
50
  num_attention_heads (`int`, *optional*, defaults to 32):
51
  Number of attention heads for each attention layer in the Transformer encoder.
52
+ num_key_value_heads (`int`, *optional*, defaults to 32):
53
  This is the number of key_value heads that should be used to implement Grouped Query Attention. If
54
  `num_key_value_heads=num_attention_heads`, the model will use Multi Head Attention (MHA), if
55
  `num_key_value_heads=1 the model will use Multi Query Attention (MQA) otherwise GQA is used. When
 
59
  `num_attention_heads`.
60
  hidden_act (`str` or `function`, *optional*, defaults to `"silu"`):
61
  The non-linear activation function (function or string).
62
+ max_position_embeddings (`int`, *optional*, defaults to 4096):
 
 
 
 
63
  The maximum sequence length that this model might ever be used with.
64
  Typically set this to something large just in case (e.g., 512 or 1024 or 2048).
65
+ initializer_range (`float`, *optional*, defaults to 0.02):
66
  The standard deviation of the truncated_normal_initializer for initializing
67
  all weight matrices.
68
+ layer_norm_eps (`float`, *optional*, defaults to 1e-05):
69
  The epsilon used by the normalization layers.
70
  use_cache (`bool`, *optional*, defaults to `True`):
71
  Whether or not the model should return the last key/values attentions
72
  (not used by all models). Only relevant if `config.is_decoder=True`.
73
+ tie_word_embeddings (`bool`, *optional*, defaults to `False`):
74
+ Whether the model's input and output word embeddings should be tied.
75
+ rope_theta (`float`, *optional*, defaults to `10000.0`):
76
+ The base period of the RoPE embeddings.
77
+ rope_scaling (`Dict`, *optional*):
78
+ Dictionary containing the scaling configuration for the RoPE embeddings. Currently supports two scaling
79
+ strategies: linear and dynamic. Their scaling factor must be a float greater than 1. The expected format is
80
+ `{"type": strategy name, "factor": scaling factor}`. When using this flag, don't update
81
+ `max_position_embeddings` to the expected new maximum. See the following thread for more information on how
82
+ these scaling strategies behave:
83
+ https://www.reddit.com/r/LocalLLaMA/comments/14mrgpr/dynamically_scaled_rope_further_increases/. This
84
+ is an experimental feature, subject to breaking API changes in future versions.
85
+ use_qkv_bias (`bool`, *optional*, defaults to `False`):
86
  Whether or not the model should use bias for qkv layers.
87
+ qk_layernorm (`bool`, *optional*, defaults to `False`):
88
+ Whether or not to normalize, per head, the Queries and Keys after projecting the hidden states.
89
+ use_parallel_residual (`bool`, *optional*, defaults to `False`):
90
+ Whether to use a "parallel" formulation in each Transformer layer, which can provide a slight training
91
+ speedup at large scales.
92
+ hidden_dropout (`float`, *optional*, defaults to 0.0):
93
+ The dropout ratio after applying the MLP to the hidden states.
94
  attention_dropout (`float`, *optional*, defaults to 0.0):
95
  The dropout ratio for the attention probabilities.
96
+ partial_rotary_factor (`float`, *optional*, defaults to 0.25):
97
+ Percentage of the query and keys which will have rotary embedding.
98
+ bos_token_id (int, *optional*, defaults to 0):
99
+ The id of the `BOS` token in the vocabulary.
100
+ eos_token_id (int, *optional*, defaults to 0):
101
+ The id of the `EOS` token in the vocabulary.
102
+
103
+ Example:
104
+
105
+ ```python
106
+ >>> from transformers import StableLmModel, StableLmConfig
107
+
108
+ >>> # Initializing a StableLM stablelm-3b style configuration
109
+ >>> configuration = StableLmConfig()
110
+ ```"""
111
+
112
+ model_type = "stablelm"
113
  keys_to_ignore_at_inference = ["past_key_values"]
114
 
115
  def __init__(
116
  self,
117
+ vocab_size=50304,
118
  intermediate_size=6912,
119
  hidden_size=2560,
120
  num_hidden_layers=32,
121
  num_attention_heads=32,
122
  num_key_value_heads=32,
123
  hidden_act="silu",
 
 
124
  max_position_embeddings=4096,
125
  initializer_range=0.02,
126
+ layer_norm_eps=1.0e-5,
127
  use_cache=True,
 
 
 
128
  tie_word_embeddings=False,
129
+ rope_theta=10_000,
130
+ rope_scaling=None,
131
+ use_qkv_bias=False,
132
+ qk_layernorm=False,
133
+ use_parallel_residual=False,
134
+ hidden_dropout=0.0,
135
+ attention_dropout=0.0,
136
+ partial_rotary_factor=0.25,
137
+ bos_token_id=0,
138
+ eos_token_id=0,
139
  **kwargs,
140
  ):
141
  self.vocab_size = vocab_size
142
  self.max_position_embeddings = max_position_embeddings
143
+
144
  self.hidden_size = hidden_size
145
+ self.intermediate_size = intermediate_size
146
  self.num_hidden_layers = num_hidden_layers
147
  self.num_attention_heads = num_attention_heads
148
  self.num_key_value_heads = num_key_value_heads
149
  self.hidden_act = hidden_act
150
+
 
151
  self.initializer_range = initializer_range
152
+ self.layer_norm_eps = layer_norm_eps
153
  self.use_cache = use_cache
154
+ self.rope_theta = rope_theta
155
+ self.rope_scaling = rope_scaling
156
  self.use_qkv_bias = use_qkv_bias
157
+ self.qk_layernorm = qk_layernorm
158
+ self.use_parallel_residual = use_parallel_residual
159
+ self.hidden_dropout = hidden_dropout
160
  self.attention_dropout = attention_dropout
161
+ self.partial_rotary_factor = partial_rotary_factor
162
+ self._rope_scaling_validation()
163
+
164
  super().__init__(
165
  bos_token_id=bos_token_id,
166
  eos_token_id=eos_token_id,
167
  tie_word_embeddings=tie_word_embeddings,
168
  **kwargs,
169
  )
170
+
171
+ # Copied from transformers.models.llama.configuration_llama.LlamaConfig._rope_scaling_validation
172
+ def _rope_scaling_validation(self):
173
+ """
174
+ Validate the `rope_scaling` configuration.
175
+ """
176
+ if self.rope_scaling is None:
177
+ return
178
+
179
+ if not isinstance(self.rope_scaling, dict) or len(self.rope_scaling) != 2:
180
+ raise ValueError(
181
+ "`rope_scaling` must be a dictionary with with two fields, `type` and `factor`, "
182
+ f"got {self.rope_scaling}"
183
+ )
184
+ rope_scaling_type = self.rope_scaling.get("type", None)
185
+ rope_scaling_factor = self.rope_scaling.get("factor", None)
186
+ if rope_scaling_type is None or rope_scaling_type not in ["linear", "dynamic"]:
187
+ raise ValueError(
188
+ f"`rope_scaling`'s type field must be one of ['linear', 'dynamic'], got {rope_scaling_type}"
189
+ )
190
+ if (
191
+ rope_scaling_factor is None
192
+ or not isinstance(rope_scaling_factor, float)
193
+ or rope_scaling_factor <= 1.0
194
+ ):
195
+ raise ValueError(
196
+ f"`rope_scaling`'s factor field must be a float > 1, got {rope_scaling_factor}"
197
+ )
generation_config.json CHANGED
@@ -1,6 +1,7 @@
1
  {
2
  "_from_model_config": true,
 
3
  "eos_token_id": 100257,
4
- "transformers_version": "4.38.2",
5
- "use_cache": false
6
  }
 
1
  {
2
  "_from_model_config": true,
3
+ "bos_token_id": 100257,
4
  "eos_token_id": 100257,
5
+ "pad_token_id": 100257,
6
+ "transformers_version": "4.39.0.dev0"
7
  }
merges.txt ADDED
The diff for this file is too large to render. See raw diff
 
model-00003-of-00011.safetensors → model-00001-of-00005.safetensors RENAMED
@@ -1,3 +1,3 @@
1
  version https://git-lfs.github.com/spec/v1
2
- oid sha256:9753f9ffcb3c923f74134da0286a92b4a760ce1e6125917577841f35588490df
3
- size 4991480240
 
1
  version https://git-lfs.github.com/spec/v1
2
+ oid sha256:f7b82912e530ed52a6c86e02926d2f2aaff637c4661237b1046d8d951d0cae9e
3
+ size 4996740184
model-00004-of-00011.safetensors → model-00002-of-00005.safetensors RENAMED
@@ -1,3 +1,3 @@
1
  version https://git-lfs.github.com/spec/v1
2
- oid sha256:870952cf822b6eab8ddff063bf15c471af7d3afd421053a142d9b97c9875462e
3
- size 4729285024
 
1
  version https://git-lfs.github.com/spec/v1
2
+ oid sha256:d4892fa9d48fec59b7b965f517a6753f40514095c41604dcc519acd74838bae6
3
+ size 4988954032
model-00001-of-00011.safetensors → model-00003-of-00005.safetensors RENAMED
@@ -1,3 +1,3 @@
1
  version https://git-lfs.github.com/spec/v1
2
- oid sha256:c23adb51b0165de13a01370c8be51647eb47efcc962b48a49b463dbf0761e778
3
- size 4823584128
 
1
  version https://git-lfs.github.com/spec/v1
2
+ oid sha256:4dcf8d0fad1c259a5a87db3403a869007832c2de0a8b251da405e8ec101d9a4b
3
+ size 4988954112
model-00002-of-00011.safetensors → model-00004-of-00005.safetensors RENAMED
@@ -1,3 +1,3 @@
1
  version https://git-lfs.github.com/spec/v1
2
- oid sha256:c19dfb81bba890a457860e01b4d30882c01e9859db22f18bc1706df9a5a6898a
3
- size 4729284832
 
1
  version https://git-lfs.github.com/spec/v1
2
+ oid sha256:38cfb1c7d31b72065eaa30f3fac0fc61d68d5b585681ea1e0c3e0ef62ba77176
3
+ size 4949632512
model-00005-of-00005.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:9be8630013080c8f9d0fd32f1596bd69b1cbf65252aac42f0c8149b45012cf99
3
+ size 4362333000
model-00005-of-00011.safetensors DELETED
@@ -1,3 +0,0 @@
1
- version https://git-lfs.github.com/spec/v1
2
- oid sha256:aeac7e78918f081de151b2d815b944af61c496c0271df23526ae163b4f9c232d
3
- size 4729285024
 
 
 
 
model-00006-of-00011.safetensors DELETED
@@ -1,3 +0,0 @@
1
- version https://git-lfs.github.com/spec/v1
2
- oid sha256:4eba80232438e4a3686cbf8ca2f3cc4c46dc8ac86ca36091d07d0ebd47884bb1
3
- size 4991480472
 
 
 
 
model-00007-of-00011.safetensors DELETED
@@ -1,3 +0,0 @@
1
- version https://git-lfs.github.com/spec/v1
2
- oid sha256:19104b4e5135ca62ecf14e78236dfc4dff274a8d9c6c8c92840d6025f8ab7209
3
- size 4729285024
 
 
 
 
model-00008-of-00011.safetensors DELETED
@@ -1,3 +0,0 @@
1
- version https://git-lfs.github.com/spec/v1
2
- oid sha256:c83b0864c77b5b6cbd8ea16cd32bc4f8173dc2fa655ca9c85d9d07129dcd285d
3
- size 4729285024
 
 
 
 
model-00009-of-00011.safetensors DELETED
@@ -1,3 +0,0 @@
1
- version https://git-lfs.github.com/spec/v1
2
- oid sha256:6989452d45de58e4f9cf1f8d4ea58ae20f9fa40263661f229f6755eb1434e834
3
- size 4991480472
 
 
 
 
model-00010-of-00011.safetensors DELETED
@@ -1,3 +0,0 @@
1
- version https://git-lfs.github.com/spec/v1
2
- oid sha256:68b0db26738c4395980f9c3b3fa93939f67f402094f855c341321cd0ce19a6d2
3
- size 3072493336
 
 
 
 
model-00011-of-00011.safetensors DELETED
@@ -1,3 +0,0 @@
1
- version https://git-lfs.github.com/spec/v1
2
- oid sha256:255fd9c0259c6774d81e64dc95730c6e30942e346b5865a6b890d2f5f8bacf8e
3
- size 2055209088
 
 
 
 
model.safetensors.index.json CHANGED
The diff for this file is too large to render. See raw diff
 
modeling_stablelm_epoch_2.py → modeling_stablelm.py RENAMED
@@ -1,4 +1,9 @@
1
- # Copyright 2023 Stability AI, EleutherAI, and The HuggingFace Inc. team. All rights reserved.
 
 
 
 
 
2
  #
3
  # Licensed under the Apache License, Version 2.0 (the "License");
4
  # you may not use this file except in compliance with the License.
@@ -11,30 +16,38 @@
11
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
  # See the License for the specific language governing permissions and
13
  # limitations under the License.
14
- #
15
- # This code is based off the following work:
16
- # https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/modeling_llama.py
17
- # https://github.com/huggingface/transformers/blob/main/src/transformers/models/gpt_neox/modeling_gpt_neox.py
18
- """ PyTorch StableLM Epoch model. """
19
- from typing import Optional, Tuple, Union
20
  import math
21
- import warnings
22
 
23
  import torch
24
  import torch.nn.functional as F
25
  import torch.utils.checkpoint
26
  from torch import nn
27
- from torch.nn import CrossEntropyLoss
28
 
29
- from transformers.cache_utils import Cache
 
 
 
 
 
30
  from transformers.modeling_outputs import (
31
  BaseModelOutputWithPast,
32
  CausalLMOutputWithPast,
 
33
  )
34
  from transformers.modeling_utils import PreTrainedModel
35
- from transformers.utils import logging, is_flash_attn_greater_or_equal_2_10
 
 
 
 
 
 
 
 
36
 
37
- from .configuration_stablelm_epoch import StableLMEpochConfig
38
 
39
  try:
40
  from flash_attn import flash_attn_func, flash_attn_varlen_func
@@ -46,15 +59,15 @@ except:
46
 
47
  logger = logging.get_logger(__name__)
48
 
 
 
49
 
50
  # Copied from transformers.models.llama.modeling_llama._get_unpad_data
51
  def _get_unpad_data(attention_mask):
52
  seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32)
53
  indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()
54
  max_seqlen_in_batch = seqlens_in_batch.max().item()
55
- cu_seqlens = F.pad(
56
- torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.torch.int32), (1, 0)
57
- )
58
  return (
59
  indices,
60
  cu_seqlens,
@@ -62,58 +75,9 @@ def _get_unpad_data(attention_mask):
62
  )
63
 
64
 
65
- # Copied from transformers.models.bart.modeling_bart._make_causal_mask
66
- def _make_causal_mask(
67
- input_ids_shape: torch.Size,
68
- dtype: torch.dtype,
69
- device: torch.device,
70
- past_key_values_length: int = 0,
71
- ):
72
- """Make causal mask used for bi-directional self-attention."""
73
- batch_size, tgt_len = input_ids_shape
74
- mask = torch.full((tgt_len, tgt_len), torch.finfo(torch.float16).min, device=device)
75
- mask_cond = torch.arange(mask.size(-1), device=device)
76
- mask.masked_fill_(mask_cond < (mask_cond + 1).view(mask.size(-1), 1), 0)
77
- mask = mask.to(dtype)
78
- if past_key_values_length > 0:
79
- mask = torch.cat(
80
- [
81
- torch.zeros(
82
- tgt_len, past_key_values_length, dtype=dtype, device=device
83
- ),
84
- mask,
85
- ],
86
- dim=-1,
87
- )
88
- return mask[None, None, :, :].expand(
89
- batch_size, 1, tgt_len, tgt_len + past_key_values_length
90
- )
91
-
92
-
93
- # Copied from transformers.models.bart.modeling_bart._expand_mask
94
- def _expand_mask(mask: torch.Tensor, dtype: torch.dtype, tgt_len: Optional[int] = None):
95
- """Expands attention_mask from `[batch_size, seq_len]` to `[batch_size, 1, tgt_seq_len, src_seq_len]`."""
96
- batch_size, src_len = mask.size()
97
- tgt_len = tgt_len if tgt_len is not None else src_len
98
-
99
- expanded_mask = (
100
- mask[:, None, None, :].expand(batch_size, 1, tgt_len, src_len).to(dtype)
101
- )
102
- inverted_mask = 1.0 - expanded_mask
103
-
104
- return inverted_mask.masked_fill(
105
- inverted_mask.to(torch.bool), torch.finfo(dtype).min
106
- )
107
-
108
-
109
- class RotaryEmbedding(nn.Module):
110
- def __init__(
111
- self,
112
- dim: int,
113
- max_position_embeddings: int,
114
- base: int = 10_000,
115
- device: Optional[torch.device] = None,
116
- ):
117
  super().__init__()
118
 
119
  self.dim = dim
@@ -122,7 +86,7 @@ class RotaryEmbedding(nn.Module):
122
  inv_freq = 1.0 / (
123
  self.base
124
  ** (
125
- torch.arange(0, self.dim, 2, device=device, dtype=torch.float32)
126
  / self.dim
127
  )
128
  )
@@ -135,97 +99,178 @@ class RotaryEmbedding(nn.Module):
135
  dtype=torch.get_default_dtype(),
136
  )
137
 
138
- def _set_cos_sin_cache(
139
- self, seq_len: int, device: torch.device, dtype: torch.dtype
140
- ):
141
  self.max_seq_len_cached = seq_len
142
- t = torch.arange(self.max_seq_len_cached, device=device, dtype=torch.float32)
 
 
143
 
144
- # Don't do einsum, it converts fp32 to fp16 under AMP
145
- # freqs = torch.einsum("i,j->ij", t, self.inv_freq)
146
  freqs = torch.outer(t, self.inv_freq)
147
  # Different from paper, but it uses a different permutation in order to obtain the same calculation
148
  emb = torch.cat((freqs, freqs), dim=-1)
149
- self.register_buffer(
150
- "cos_cached", emb.cos()[None, None, :, :].to(dtype), persistent=False
151
- )
152
- self.register_buffer(
153
- "sin_cached", emb.sin()[None, None, :, :].to(dtype), persistent=False
154
- )
155
 
156
- def forward(self, x: torch.Tensor, seq_len: Optional[int] = None):
157
- # x: [batch_size, num_heads, seq_len, head_size]
158
  if seq_len > self.max_seq_len_cached:
159
- self._set_cos_sin_cache(
160
- seq_len=seq_len, device=x.device, dtype=torch.get_default_dtype()
161
- )
162
  return (
163
- self.cos_cached[:, :, :seq_len, ...].to(dtype=x.dtype),
164
- self.sin_cached[:, :, :seq_len, ...].to(dtype=x.dtype),
165
  )
166
 
167
 
168
- def rotate_half(x: torch.Tensor):
169
- """Rotates half the hidden dims of the input."""
170
- x1, x2 = torch.chunk(x, 2, dim=-1)
171
- return torch.cat((-x2, x1), dim=-1)
172
 
 
 
 
 
 
 
 
 
 
 
173
 
174
- def apply_rotary_pos_emb(q, k, cos, sin, position_ids):
175
- # The first two dimensions of cos and sin are always 1, so we can `squeeze` them.
176
- cos = cos.squeeze(1).squeeze(0) # [seq_len, dim]
177
- sin = sin.squeeze(1).squeeze(0) # [seq_len, dim]
178
- cos = cos[position_ids].unsqueeze(1) # [batch_size, 1, seq_len, dim]
179
- sin = sin[position_ids].unsqueeze(1) # [batch_size, 1, seq_len, dim]
180
- q_embed = (q * cos) + (rotate_half(q) * sin)
181
- k_embed = (k * cos) + (rotate_half(k) * sin)
182
- return q_embed, k_embed
 
 
 
183
 
184
 
185
- class LayerNormPerHead(torch.nn.Module):
 
 
 
186
  def __init__(
187
  self,
188
- head_dim: int,
189
- num_heads: int,
190
- eps: float = 1e-5,
191
- bias: bool = False,
 
192
  ):
193
- super().__init__()
194
- self.head_dim = head_dim
195
- self.num_heads = num_heads
196
- self.norms = torch.torch.nn.ModuleList(
197
- [nn.LayerNorm(head_dim, eps=eps, bias=bias) for _ in range(self.num_heads)]
198
- )
199
 
200
- def forward(self, x: torch.Tensor):
201
- # Split along the num_heads axis to get per-head inputs
202
- # [batch_size, num_heads, seq_len, head_dim] -> [batch_size, 1, seq_len, head_dim] * num_heads
203
- heads = torch.split(x, 1, dim=1)
204
- # Normalize and put the heads back together
205
- return torch.cat([norm(x) for norm, x in zip(self.norms, heads)], dim=1)
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
206
 
207
 
208
- class MLP(nn.Module):
209
- def __init__(self, config: StableLMEpochConfig):
 
210
  super().__init__()
211
  self.config = config
212
  self.hidden_size = config.hidden_size
213
  self.intermediate_size = config.intermediate_size
214
- self.gate_proj = nn.Linear(
215
- config.hidden_size, config.intermediate_size, bias=False
216
- )
217
- self.up_proj = nn.Linear(
218
- config.hidden_size, config.intermediate_size, bias=False
219
- )
220
- self.down_proj = nn.Linear(
221
- config.intermediate_size, config.hidden_size, bias=False
222
- )
223
- self.act_fn = nn.SiLU()
224
 
225
- def forward(self, x: torch.Tensor) -> torch.Tensor:
226
  return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
227
 
228
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
229
  def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
230
  """
231
  This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
@@ -240,25 +285,35 @@ def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
240
  return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
241
 
242
 
243
- class Attention(nn.Module):
244
- def __init__(self, config: StableLMEpochConfig):
 
 
245
  super().__init__()
246
  self.config = config
 
 
 
 
 
 
 
 
247
  self.hidden_size = config.hidden_size
248
  self.num_heads = config.num_attention_heads
249
  self.head_dim = self.hidden_size // self.num_heads
250
  self.num_key_value_heads = config.num_key_value_heads
251
  self.num_key_value_groups = self.num_heads // self.num_key_value_heads
252
  self.max_position_embeddings = config.max_position_embeddings
 
 
253
  self.is_causal = True
254
- self.attention_dropout = config.attention_dropout
255
 
256
  if (self.head_dim * self.num_heads) != self.hidden_size:
257
  raise ValueError(
258
  f"hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size}"
259
  f" and `num_heads`: {self.num_heads})."
260
  )
261
-
262
  self.q_proj = nn.Linear(
263
  self.hidden_size, self.num_heads * self.head_dim, bias=config.use_qkv_bias
264
  )
@@ -274,31 +329,54 @@ class Attention(nn.Module):
274
  )
275
  self.o_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=False)
276
 
277
- self.q_norm = LayerNormPerHead(
278
- self.head_dim, self.num_heads, eps=config.norm_eps, bias=False
279
- )
280
- self.k_norm = LayerNormPerHead(
281
- self.head_dim, self.num_key_value_heads, eps=config.norm_eps, bias=False
282
- )
 
 
283
 
 
284
  self._init_rope()
285
 
 
286
  def _init_rope(self):
287
- self.rotary_ndims = int(self.head_dim * self.config.rope_pct)
288
- self.rotary_emb = RotaryEmbedding(
289
- self.rotary_ndims,
290
- max_position_embeddings=self.config.max_position_embeddings,
291
- base=self.config.rope_theta,
292
- )
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
293
 
294
  def forward(
295
  self,
296
- hidden_states: torch.FloatTensor,
297
- attention_mask: torch.FloatTensor,
298
- position_ids: torch.LongTensor,
299
- past_key_value: Optional[Tuple[torch.Tensor]] = None,
300
- output_attentions: Optional[bool] = False,
301
- use_cache: Optional[bool] = False,
302
  ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
303
  bsz, q_len, _ = hidden_states.size()
304
 
@@ -316,33 +394,49 @@ class Attention(nn.Module):
316
  bsz, q_len, self.num_key_value_heads, self.head_dim
317
  ).transpose(1, 2)
318
 
319
- # [batch_size, num_heads, seq_len, head_dim]
320
- query_states = self.q_norm(query_states)
321
- key_states = self.k_norm(key_states)
322
-
323
- query_rot = query_states[..., : self.rotary_ndims]
324
- query_pass = query_states[..., self.rotary_ndims :]
325
- key_rot = key_states[..., : self.rotary_ndims]
326
- key_pass = key_states[..., self.rotary_ndims :]
327
 
328
  kv_seq_len = key_states.shape[-2]
329
  if past_key_value is not None:
330
- kv_seq_len += past_key_value[0].shape[-2]
 
 
 
 
 
 
331
  cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
332
- query_states, key_states = apply_rotary_pos_emb(
 
 
 
 
 
 
 
 
 
 
 
333
  query_rot, key_rot, cos, sin, position_ids
334
  )
335
 
336
- # [batch_size, num_heads, seq_len, head_dim]
337
- query_states = torch.cat((query_states, query_pass), dim=-1)
338
- key_states = torch.cat((key_states, key_pass), dim=-1)
339
 
340
  if past_key_value is not None:
341
- # Reuse k, v, self_attention
342
- key_states = torch.cat((past_key_value[0], key_states), dim=2)
343
- value_states = torch.cat((past_key_value[1], value_states), dim=2)
344
-
345
- past_key_value = (key_states, value_states) if use_cache else None
 
 
 
 
346
 
347
  # Repeat k/v heads if n_kv_heads < n_heads
348
  key_states = repeat_kv(key_states, self.num_key_value_groups)
@@ -365,13 +459,12 @@ class Attention(nn.Module):
365
  )
366
  attn_weights = attn_weights + attention_mask
367
 
368
- # Upcast attention to fp32
369
  attn_weights = nn.functional.softmax(
370
- attn_weights, dim=-1, dtype=torch.float32
371
  ).to(query_states.dtype)
372
- attn_weights = nn.functional.dropout(
373
- attn_weights, p=self.attention_dropout, training=self.training
374
- )
375
  attn_output = torch.matmul(attn_weights, value_states)
376
 
377
  if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim):
@@ -380,11 +473,9 @@ class Attention(nn.Module):
380
  f" {attn_output.size()}"
381
  )
382
 
383
- # Merge heads
384
  attn_output = attn_output.transpose(1, 2).contiguous()
385
  attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)
386
 
387
- # Final linear projection
388
  attn_output = self.o_proj(attn_output)
389
 
390
  if not output_attentions:
@@ -393,16 +484,133 @@ class Attention(nn.Module):
393
  return attn_output, attn_weights, past_key_value
394
 
395
 
396
- class FlashAttention2(Attention):
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
397
  """
398
- Reference: https://github.com/huggingface/transformers/blob/5d36025ca13d05151b7a0c761e90d429c4644a30/src/transformers/models/llama/modeling_llama.py#L456
 
 
399
  """
400
 
 
401
  def __init__(self, *args, **kwargs):
402
  super().__init__(*args, **kwargs)
403
 
404
  # TODO: Should be removed once Flash Attention for RoCm is bumped to 2.1.
405
- # 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.
406
  # 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).
407
  self._flash_attn_uses_top_left_mask = not is_flash_attn_greater_or_equal_2_10()
408
 
@@ -416,14 +624,7 @@ class FlashAttention2(Attention):
416
  use_cache: bool = False,
417
  **kwargs,
418
  ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
419
- # FlashAttention2 attention does not support output_attentions
420
- if "padding_mask" in kwargs:
421
- warnings.warn(
422
- "Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use `attention_mask` instead.`"
423
- )
424
-
425
- # overwrite attention_mask with padding_mask
426
- attention_mask = kwargs.pop("padding_mask")
427
 
428
  output_attentions = False
429
 
@@ -446,33 +647,47 @@ class FlashAttention2(Attention):
446
  bsz, q_len, self.num_key_value_heads, self.head_dim
447
  ).transpose(1, 2)
448
 
449
- # [batch_size, num_heads, seq_len, head_dim]
450
- query_states = self.q_norm(query_states)
451
- key_states = self.k_norm(key_states)
452
-
453
- query_rot = query_states[..., : self.rotary_ndims]
454
- query_pass = query_states[..., self.rotary_ndims :]
455
- key_rot = key_states[..., : self.rotary_ndims]
456
- key_pass = key_states[..., self.rotary_ndims :]
457
 
458
  kv_seq_len = key_states.shape[-2]
459
  if past_key_value is not None:
460
- kv_seq_len += past_key_value[0].shape[-2]
 
 
 
 
 
 
461
  cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
462
- query_states, key_states = apply_rotary_pos_emb(
 
 
 
 
 
 
 
 
 
 
463
  query_rot, key_rot, cos, sin, position_ids
464
  )
465
 
466
- # [batch_size, num_heads, seq_len, head_dim]
467
- query_states = torch.cat((query_states, query_pass), dim=-1)
468
- key_states = torch.cat((key_states, key_pass), dim=-1)
469
 
470
  if past_key_value is not None:
471
- # Reuse k, v, self_attention
472
- key_states = torch.cat((past_key_value[0], key_states), dim=2)
473
- value_states = torch.cat((past_key_value[1], value_states), dim=2)
474
-
475
- past_key_value = (key_states, value_states) if use_cache else None
 
 
 
476
 
477
  # TODO: These transpose are quite inefficient but Flash Attention requires the layout [batch_size, sequence_length, num_heads, head_dim]. We would need to refactor the KV cache
478
  # to be able to avoid many of these transpose/reshape/view.
@@ -480,7 +695,7 @@ class FlashAttention2(Attention):
480
  key_states = key_states.transpose(1, 2)
481
  value_states = value_states.transpose(1, 2)
482
 
483
- dropout_rate = self.attention_dropout if self.training else 0.0
484
 
485
  attn_output = self._flash_attention_forward(
486
  query_states,
@@ -490,6 +705,7 @@ class FlashAttention2(Attention):
490
  q_len,
491
  dropout=dropout_rate,
492
  )
 
493
  attn_output = attn_output.reshape(bsz, q_len, self.hidden_size).contiguous()
494
  attn_output = self.o_proj(attn_output)
495
 
@@ -498,6 +714,7 @@ class FlashAttention2(Attention):
498
 
499
  return attn_output, attn_weights, past_key_value
500
 
 
501
  def _flash_attention_forward(
502
  self,
503
  query_states,
@@ -522,7 +739,7 @@ class FlashAttention2(Attention):
522
  attention_mask (`torch.Tensor`):
523
  The padding mask - corresponds to a tensor of size `(batch_size, seq_len)` where 0 stands for the
524
  position of padding tokens and 1 for the position of non-padding tokens.
525
- dropout (`int`, *optional*):
526
  Attention dropout
527
  softmax_scale (`float`, *optional*):
528
  The scaling of QK^T before applying softmax. Default to 1 / sqrt(head_dim)
@@ -530,7 +747,7 @@ class FlashAttention2(Attention):
530
  if not self._flash_attn_uses_top_left_mask:
531
  causal = self.is_causal
532
  else:
533
- # TODO: Remove the `query_length != 1` check once Flash Attention for RoCm is bumped to 2.1. For details, please see the comment in FlashAttention2 __init__.
534
  causal = self.is_causal and query_length != 1
535
 
536
  # Contains at least one padding token in the sequence
@@ -578,6 +795,7 @@ class FlashAttention2(Attention):
578
 
579
  return attn_output
580
 
 
581
  def _upad_input(
582
  self, query_layer, key_layer, value_layer, attention_mask, query_length
583
  ):
@@ -625,32 +843,62 @@ class FlashAttention2(Attention):
625
 
626
 
627
  ATTENTION_CLASSES = {
628
- "eager": Attention,
629
- "flash_attention_2": FlashAttention2,
 
630
  }
631
 
632
 
633
- class DecoderLayer(nn.Module):
634
- def __init__(self, config: StableLMEpochConfig):
635
  super().__init__()
636
- self.self_attn = ATTENTION_CLASSES[config._attn_implementation](config=config)
637
- self.mlp = MLP(config)
 
 
 
 
638
  self.input_layernorm = nn.LayerNorm(
639
- config.hidden_size, eps=config.norm_eps, bias=config.use_norm_bias
640
  )
 
 
 
 
 
 
641
 
642
  def forward(
643
  self,
644
- hidden_states: Optional[torch.FloatTensor],
645
- attention_mask: Optional[torch.FloatTensor] = None,
646
  position_ids: Optional[torch.LongTensor] = None,
647
  past_key_value: Optional[Tuple[torch.Tensor]] = None,
648
  output_attentions: Optional[bool] = False,
649
  use_cache: Optional[bool] = False,
650
- ) -> Union[
651
- Tuple[torch.Tensor],
652
- Optional[Tuple[torch.Tensor, Tuple[torch.FloatTensor, ...]]],
653
  ]:
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
654
  residual = hidden_states
655
 
656
  hidden_states = self.input_layernorm(hidden_states)
@@ -665,11 +913,21 @@ class DecoderLayer(nn.Module):
665
  use_cache=use_cache,
666
  )
667
 
668
- # Fully Connected
669
- mlp_output = self.mlp(hidden_states)
670
-
671
- # Parallel Residual
672
- hidden_states = residual + self_attn_output + mlp_output
 
 
 
 
 
 
 
 
 
 
673
 
674
  outputs = (hidden_states,)
675
 
@@ -682,50 +940,148 @@ class DecoderLayer(nn.Module):
682
  return outputs
683
 
684
 
685
- class StableLMEpochPreTrainedModel(PreTrainedModel):
686
- """An abstract class to handle weights initialization and a simple interface
687
- for downloading and loading pretrained models.
688
- """
 
 
 
 
 
 
 
 
 
 
 
689
 
690
- config_class = StableLMEpochConfig
691
- base_model_prefix = "transformer"
 
 
 
 
 
 
692
  supports_gradient_checkpointing = True
693
- _no_split_modules = ["DecoderLayer"]
694
  _skip_keys_device_placement = "past_key_values"
695
  _supports_flash_attn_2 = True
 
 
696
 
697
- def _init_weights(self, module: nn.Module):
698
- """Initialize the weights"""
699
  if isinstance(module, nn.Linear):
700
- module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
701
  if module.bias is not None:
702
  module.bias.data.zero_()
703
  elif isinstance(module, nn.Embedding):
704
- module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
705
  if module.padding_idx is not None:
706
  module.weight.data[module.padding_idx].zero_()
707
- elif isinstance(module, nn.LayerNorm):
708
- if module.bias is not None:
709
- module.bias.data.zero_()
710
- module.weight.data.fill_(1.0)
711
 
712
- def _set_gradient_checkpointing(self, module: nn.Module, value=False):
713
- if isinstance(module, StableLMEpochModel):
714
- module.gradient_checkpointing = value
715
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
716
 
717
- class StableLMEpochModel(StableLMEpochPreTrainedModel):
718
- def __init__(self, config: StableLMEpochConfig):
 
 
 
719
  super().__init__(config)
 
 
 
720
  self.embed_tokens = nn.Embedding(
721
- config.vocab_size, config.hidden_size, config.pad_token_id
722
  )
723
  self.layers = nn.ModuleList(
724
- [DecoderLayer(config) for _ in range(config.num_hidden_layers)]
 
 
 
725
  )
726
- self.norm = nn.LayerNorm(config.hidden_size, eps=config.norm_eps)
727
 
728
- self._use_flash_attention_2 = config._attn_implementation == "flash_attention_2"
729
  self.gradient_checkpointing = False
730
  # Initialize weights and apply final processing
731
  self.post_init()
@@ -733,47 +1089,16 @@ class StableLMEpochModel(StableLMEpochPreTrainedModel):
733
  def get_input_embeddings(self):
734
  return self.embed_tokens
735
 
736
- def set_input_embeddings(self, value: nn.Module):
737
  self.embed_tokens = value
738
 
739
- # Copied from transformers.models.bart.modeling_bart.BartDecoder._prepare_decoder_attention_mask
740
- def _prepare_decoder_attention_mask(
741
- self,
742
- attention_mask: torch.Tensor,
743
- input_shape: torch.Size,
744
- inputs_embeds: torch.Tensor,
745
- past_key_values_length: int,
746
- ):
747
- # Create causal mask
748
- # [batch_size, seq_len] -> [batch_size, 1, tgt_seq_len, src_seq_len]
749
- combined_attention_mask = None
750
- if input_shape[-1] > 1:
751
- combined_attention_mask = _make_causal_mask(
752
- input_shape,
753
- inputs_embeds.dtype,
754
- device=inputs_embeds.device,
755
- past_key_values_length=past_key_values_length,
756
- )
757
-
758
- if attention_mask is not None:
759
- # [batch_size, seq_len] -> [batch_size, 1, tgt_seq_len, src_seq_len]
760
- expanded_attn_mask = _expand_mask(
761
- attention_mask, inputs_embeds.dtype, tgt_len=input_shape[-1]
762
- ).to(inputs_embeds.device)
763
- combined_attention_mask = (
764
- expanded_attn_mask
765
- if combined_attention_mask is None
766
- else expanded_attn_mask + combined_attention_mask
767
- )
768
-
769
- return combined_attention_mask
770
-
771
  def forward(
772
  self,
773
- input_ids: Optional[torch.LongTensor] = None,
774
- attention_mask: Optional[torch.FloatTensor] = None,
775
  position_ids: Optional[torch.LongTensor] = None,
776
- past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
777
  inputs_embeds: Optional[torch.FloatTensor] = None,
778
  use_cache: Optional[bool] = None,
779
  output_attentions: Optional[bool] = None,
@@ -796,7 +1121,7 @@ class StableLMEpochModel(StableLMEpochPreTrainedModel):
796
  return_dict if return_dict is not None else self.config.use_return_dict
797
  )
798
 
799
- # Retrieve input_ids and inputs_embeds
800
  if input_ids is not None and inputs_embeds is not None:
801
  raise ValueError(
802
  "You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time"
@@ -813,6 +1138,20 @@ class StableLMEpochModel(StableLMEpochPreTrainedModel):
813
  seq_length_with_past = seq_length
814
  past_key_values_length = 0
815
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
816
  if position_ids is None:
817
  device = input_ids.device if input_ids is not None else inputs_embeds.device
818
  position_ids = torch.arange(
@@ -821,28 +1160,29 @@ class StableLMEpochModel(StableLMEpochPreTrainedModel):
821
  dtype=torch.long,
822
  device=device,
823
  )
824
- position_ids = position_ids.unsqueeze(0).view(-1, seq_length)
825
- else:
826
- position_ids = position_ids.view(-1, seq_length).long()
827
 
828
  if inputs_embeds is None:
829
  inputs_embeds = self.embed_tokens(input_ids)
830
- # Embed positions
831
- if self._use_flash_attention_2:
832
  # 2d mask is passed through the layers
833
  attention_mask = (
834
  attention_mask
835
  if (attention_mask is not None and 0 in attention_mask)
836
  else None
837
  )
 
 
 
 
 
 
 
 
838
  else:
839
- if attention_mask is None:
840
- attention_mask = torch.ones(
841
- (batch_size, seq_length_with_past),
842
- dtype=torch.bool,
843
- device=inputs_embeds.device,
844
- )
845
- attention_mask = self._prepare_decoder_attention_mask(
846
  attention_mask,
847
  (batch_size, seq_length),
848
  inputs_embeds,
@@ -851,47 +1191,30 @@ class StableLMEpochModel(StableLMEpochPreTrainedModel):
851
 
852
  hidden_states = inputs_embeds
853
 
854
- if self.gradient_checkpointing and self.training:
855
- if use_cache:
856
- logger.warning(
857
- "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
858
- )
859
- use_cache = False
860
-
861
- # Decoder layers
862
  all_hidden_states = () if output_hidden_states else None
863
  all_self_attns = () if output_attentions else None
864
- next_decoder_cache = () if use_cache else None
865
 
866
- for idx, decoder_layer in enumerate(self.layers):
867
  if output_hidden_states:
868
  all_hidden_states += (hidden_states,)
869
 
870
- past_key_value = (
871
- past_key_values[idx] if past_key_values is not None else None
872
- )
873
-
874
  if self.gradient_checkpointing and self.training:
875
-
876
- def create_custom_forward(module):
877
- def custom_forward(*inputs):
878
- # None for past_key_value
879
- return module(*inputs, past_key_value, output_attentions)
880
-
881
- return custom_forward
882
-
883
- layer_outputs = torch.utils.checkpoint.checkpoint(
884
- create_custom_forward(decoder_layer),
885
  hidden_states,
886
  attention_mask,
887
  position_ids,
 
 
888
  )
889
  else:
890
  layer_outputs = decoder_layer(
891
  hidden_states,
892
  attention_mask=attention_mask,
893
  position_ids=position_ids,
894
- past_key_value=past_key_value,
895
  output_attentions=output_attentions,
896
  use_cache=use_cache,
897
  )
@@ -899,18 +1222,25 @@ class StableLMEpochModel(StableLMEpochPreTrainedModel):
899
  hidden_states = layer_outputs[0]
900
 
901
  if use_cache:
902
- next_decoder_cache += (layer_outputs[2 if output_attentions else 1],)
903
 
904
  if output_attentions:
905
  all_self_attns += (layer_outputs[1],)
906
 
907
  hidden_states = self.norm(hidden_states)
908
 
909
- # Add hidden states from the last decoder layer
910
  if output_hidden_states:
911
  all_hidden_states += (hidden_states,)
912
 
913
- next_cache = next_decoder_cache if use_cache else None
 
 
 
 
 
 
 
914
  if not return_dict:
915
  return tuple(
916
  v
@@ -925,42 +1255,55 @@ class StableLMEpochModel(StableLMEpochPreTrainedModel):
925
  )
926
 
927
 
928
- class StableLMEpochForCausalLM(StableLMEpochPreTrainedModel):
 
929
  _tied_weights_keys = ["lm_head.weight"]
930
 
931
- def __init__(self, config: StableLMEpochConfig):
 
932
  super().__init__(config)
933
-
934
- self.model = StableLMEpochModel(config)
935
  self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
936
 
937
  # Initialize weights and apply final processing
938
  self.post_init()
939
 
 
940
  def get_input_embeddings(self):
941
  return self.model.embed_tokens
942
 
 
943
  def set_input_embeddings(self, value):
944
  self.model.embed_tokens = value
945
 
 
946
  def get_output_embeddings(self):
947
  return self.lm_head
948
 
949
- def set_output_embeddings(self, new_embeddings: nn.Module):
 
950
  self.lm_head = new_embeddings
951
 
952
- def get_decoder(self):
953
- return self.model
954
-
955
  def set_decoder(self, decoder):
956
  self.model = decoder
957
 
 
 
 
 
 
 
 
 
 
958
  def forward(
959
  self,
960
- input_ids: Optional[torch.LongTensor] = None,
961
- attention_mask: Optional[torch.FloatTensor] = None,
962
  position_ids: Optional[torch.LongTensor] = None,
963
- past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
964
  inputs_embeds: Optional[torch.FloatTensor] = None,
965
  labels: Optional[torch.LongTensor] = None,
966
  use_cache: Optional[bool] = None,
@@ -968,6 +1311,32 @@ class StableLMEpochForCausalLM(StableLMEpochPreTrainedModel):
968
  output_hidden_states: Optional[bool] = None,
969
  return_dict: Optional[bool] = None,
970
  ) -> Union[Tuple, CausalLMOutputWithPast]:
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
971
  output_attentions = (
972
  output_attentions
973
  if output_attentions is not None
@@ -982,9 +1351,8 @@ class StableLMEpochForCausalLM(StableLMEpochPreTrainedModel):
982
  return_dict if return_dict is not None else self.config.use_return_dict
983
  )
984
 
985
- # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
986
  outputs = self.model(
987
- input_ids,
988
  attention_mask=attention_mask,
989
  position_ids=position_ids,
990
  past_key_values=past_key_values,
@@ -996,7 +1364,7 @@ class StableLMEpochForCausalLM(StableLMEpochPreTrainedModel):
996
  )
997
 
998
  hidden_states = outputs[0]
999
- logits = self.lm_head(hidden_states).float()
1000
 
1001
  loss = None
1002
  if labels is not None:
@@ -1026,33 +1394,52 @@ class StableLMEpochForCausalLM(StableLMEpochPreTrainedModel):
1026
  def prepare_inputs_for_generation(
1027
  self,
1028
  input_ids,
1029
- past_key_values: Optional[torch.Tensor] = None,
1030
- attention_mask: Optional[torch.Tensor] = None,
1031
- inputs_embeds: Optional[torch.Tensor] = None,
1032
  **kwargs,
1033
  ):
1034
- # Trim decoder_input_ids if past is used
1035
  if past_key_values is not None:
1036
- past_length = past_key_values[0][0].shape[2]
1037
-
1038
- # Some generation methods already pass only the last input ID
1039
- if input_ids.shape[1] > past_length:
1040
- remove_prefix_length = past_length
1041
  else:
1042
- # Default to old behavior: keep only final ID
1043
- remove_prefix_length = input_ids.shape[1] - 1
1044
-
1045
- input_ids = input_ids[:, remove_prefix_length:]
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1046
 
1047
  position_ids = kwargs.get("position_ids", None)
1048
  if attention_mask is not None and position_ids is None:
1049
- # Create position_ids on the fly for batch generation
1050
  position_ids = attention_mask.long().cumsum(-1) - 1
1051
  position_ids.masked_fill_(attention_mask == 0, 1)
1052
  if past_key_values:
1053
- position_ids = position_ids[:, -1].unsqueeze(-1)
1054
 
1055
- # If `inputs_embeds` are passed, we only want to use them in the 1st generation step
1056
  if inputs_embeds is not None and past_key_values is None:
1057
  model_inputs = {"inputs_embeds": inputs_embeds}
1058
  else:
@@ -1060,10 +1447,10 @@ class StableLMEpochForCausalLM(StableLMEpochPreTrainedModel):
1060
 
1061
  model_inputs.update(
1062
  {
1063
- "attention_mask": attention_mask,
1064
  "past_key_values": past_key_values,
1065
  "use_cache": kwargs.get("use_cache"),
1066
- "position_ids": position_ids,
1067
  }
1068
  )
1069
  return model_inputs
@@ -1081,5 +1468,141 @@ class StableLMEpochForCausalLM(StableLMEpochPreTrainedModel):
1081
  return reordered_past
1082
 
1083
 
1084
- StableLMEpochConfig.register_for_auto_class()
1085
- StableLMEpochForCausalLM.register_for_auto_class("AutoModelForCausalLM")
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright 2024 EleutherAI and the HuggingFace Inc. team. All rights reserved.
2
+ #
3
+ # This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
4
+ # and OPT implementations in this library. It has been modified from its
5
+ # original forms to accommodate minor architectural differences compared
6
+ # to GPT-NeoX and OPT used by the Meta AI team that trained the model.
7
  #
8
  # Licensed under the Apache License, Version 2.0 (the "License");
9
  # you may not use this file except in compliance with the License.
 
16
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
17
  # See the License for the specific language governing permissions and
18
  # limitations under the License.
19
+ """ PyTorch StableLM model."""
 
 
 
 
 
20
  import math
21
+ from typing import List, Optional, Tuple, Union
22
 
23
  import torch
24
  import torch.nn.functional as F
25
  import torch.utils.checkpoint
26
  from torch import nn
27
+ from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
28
 
29
+ from transformers.activations import ACT2FN
30
+ from transformers.cache_utils import Cache, DynamicCache
31
+ from transformers.modeling_attn_mask_utils import (
32
+ _prepare_4d_causal_attention_mask,
33
+ _prepare_4d_causal_attention_mask_for_sdpa,
34
+ )
35
  from transformers.modeling_outputs import (
36
  BaseModelOutputWithPast,
37
  CausalLMOutputWithPast,
38
+ SequenceClassifierOutputWithPast,
39
  )
40
  from transformers.modeling_utils import PreTrainedModel
41
+ from transformers.utils import (
42
+ add_start_docstrings,
43
+ add_start_docstrings_to_model_forward,
44
+ is_flash_attn_2_available,
45
+ is_flash_attn_greater_or_equal_2_10,
46
+ logging,
47
+ replace_return_docstrings,
48
+ )
49
+ from .configuration_stablelm import StableLmConfig
50
 
 
51
 
52
  try:
53
  from flash_attn import flash_attn_func, flash_attn_varlen_func
 
59
 
60
  logger = logging.get_logger(__name__)
61
 
62
+ _CONFIG_FOR_DOC = "StableLmConfig"
63
+
64
 
65
  # Copied from transformers.models.llama.modeling_llama._get_unpad_data
66
  def _get_unpad_data(attention_mask):
67
  seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32)
68
  indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()
69
  max_seqlen_in_batch = seqlens_in_batch.max().item()
70
+ cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0))
 
 
71
  return (
72
  indices,
73
  cu_seqlens,
 
75
  )
76
 
77
 
78
+ # Copied from transformers.models.mistral.modeling_mistral.MistralRotaryEmbedding with Mistral->StableLm
79
+ class StableLmRotaryEmbedding(nn.Module):
80
+ def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None):
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
81
  super().__init__()
82
 
83
  self.dim = dim
 
86
  inv_freq = 1.0 / (
87
  self.base
88
  ** (
89
+ torch.arange(0, self.dim, 2, dtype=torch.int64).float().to(device)
90
  / self.dim
91
  )
92
  )
 
99
  dtype=torch.get_default_dtype(),
100
  )
101
 
102
+ def _set_cos_sin_cache(self, seq_len, device, dtype):
 
 
103
  self.max_seq_len_cached = seq_len
104
+ t = torch.arange(
105
+ self.max_seq_len_cached, device=device, dtype=torch.int64
106
+ ).type_as(self.inv_freq)
107
 
 
 
108
  freqs = torch.outer(t, self.inv_freq)
109
  # Different from paper, but it uses a different permutation in order to obtain the same calculation
110
  emb = torch.cat((freqs, freqs), dim=-1)
111
+ self.register_buffer("cos_cached", emb.cos().to(dtype), persistent=False)
112
+ self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False)
 
 
 
 
113
 
114
+ def forward(self, x, seq_len=None):
115
+ # x: [bs, num_attention_heads, seq_len, head_size]
116
  if seq_len > self.max_seq_len_cached:
117
+ self._set_cos_sin_cache(seq_len=seq_len, device=x.device, dtype=x.dtype)
118
+
 
119
  return (
120
+ self.cos_cached[:seq_len].to(dtype=x.dtype),
121
+ self.sin_cached[:seq_len].to(dtype=x.dtype),
122
  )
123
 
124
 
125
+ # Copied from transformers.models.falcon.modeling_falcon.FalconLinearScalingRotaryEmbedding with Falcon->StableLm
126
+ class StableLmLinearScalingRotaryEmbedding(StableLmRotaryEmbedding):
127
+ """StableLmRotaryEmbedding extended with linear scaling. Credits to the Reddit user /u/kaiokendev"""
 
128
 
129
+ def __init__(
130
+ self,
131
+ dim,
132
+ max_position_embeddings=2048,
133
+ base=10000,
134
+ device=None,
135
+ scaling_factor=1.0,
136
+ ):
137
+ self.scaling_factor = scaling_factor
138
+ super().__init__(dim, max_position_embeddings, base, device)
139
 
140
+ def _set_cos_sin_cache(self, seq_len, device, dtype):
141
+ self.max_seq_len_cached = seq_len
142
+ t = torch.arange(
143
+ self.max_seq_len_cached, device=device, dtype=torch.int64
144
+ ).type_as(self.inv_freq)
145
+ t = t / self.scaling_factor
146
+
147
+ freqs = torch.outer(t, self.inv_freq)
148
+ # Different from paper, but it uses a different permutation in order to obtain the same calculation
149
+ emb = torch.cat((freqs, freqs), dim=-1)
150
+ self.register_buffer("cos_cached", emb.cos().to(dtype), persistent=False)
151
+ self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False)
152
 
153
 
154
+ # Copied from transformers.models.falcon.modeling_falcon.FalconDynamicNTKScalingRotaryEmbedding with Falcon->StableLm
155
+ class StableLmDynamicNTKScalingRotaryEmbedding(StableLmRotaryEmbedding):
156
+ """StableLmRotaryEmbedding extended with Dynamic NTK scaling. Credits to the Reddit users /u/bloc97 and /u/emozilla"""
157
+
158
  def __init__(
159
  self,
160
+ dim,
161
+ max_position_embeddings=2048,
162
+ base=10000,
163
+ device=None,
164
+ scaling_factor=1.0,
165
  ):
166
+ self.scaling_factor = scaling_factor
167
+ super().__init__(dim, max_position_embeddings, base, device)
 
 
 
 
168
 
169
+ def _set_cos_sin_cache(self, seq_len, device, dtype):
170
+ self.max_seq_len_cached = seq_len
171
+
172
+ if seq_len > self.max_position_embeddings:
173
+ base = self.base * (
174
+ (self.scaling_factor * seq_len / self.max_position_embeddings)
175
+ - (self.scaling_factor - 1)
176
+ ) ** (self.dim / (self.dim - 2))
177
+ inv_freq = 1.0 / (
178
+ base
179
+ ** (
180
+ torch.arange(0, self.dim, 2, dtype=torch.int64).float().to(device)
181
+ / self.dim
182
+ )
183
+ )
184
+ self.register_buffer("inv_freq", inv_freq, persistent=False)
185
+
186
+ t = torch.arange(
187
+ self.max_seq_len_cached, device=device, dtype=torch.int64
188
+ ).type_as(self.inv_freq)
189
+
190
+ freqs = torch.outer(t, self.inv_freq)
191
+ # Different from paper, but it uses a different permutation in order to obtain the same calculation
192
+ emb = torch.cat((freqs, freqs), dim=-1)
193
+ self.register_buffer("cos_cached", emb.cos().to(dtype), persistent=False)
194
+ self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False)
195
+
196
+
197
+ # Copied from transformers.models.llama.modeling_llama.rotate_half
198
+ def rotate_half(x):
199
+ """Rotates half the hidden dims of the input."""
200
+ x1 = x[..., : x.shape[-1] // 2]
201
+ x2 = x[..., x.shape[-1] // 2 :]
202
+ return torch.cat((-x2, x1), dim=-1)
203
+
204
+
205
+ # Copied from transformers.models.mistral.modeling_mistral.apply_rotary_pos_emb
206
+ def apply_rotary_pos_emb(q, k, cos, sin, position_ids, unsqueeze_dim=1):
207
+ """Applies Rotary Position Embedding to the query and key tensors.
208
+
209
+ Args:
210
+ q (`torch.Tensor`): The query tensor.
211
+ k (`torch.Tensor`): The key tensor.
212
+ cos (`torch.Tensor`): The cosine part of the rotary embedding.
213
+ sin (`torch.Tensor`): The sine part of the rotary embedding.
214
+ position_ids (`torch.Tensor`):
215
+ The position indices of the tokens corresponding to the query and key tensors. For example, this can be
216
+ used to pass offsetted position ids when working with a KV-cache.
217
+ unsqueeze_dim (`int`, *optional*, defaults to 1):
218
+ The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
219
+ sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
220
+ that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
221
+ k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
222
+ cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
223
+ the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
224
+ Returns:
225
+ `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
226
+ """
227
+ cos = cos[position_ids].unsqueeze(unsqueeze_dim)
228
+ sin = sin[position_ids].unsqueeze(unsqueeze_dim)
229
+ q_embed = (q * cos) + (rotate_half(q) * sin)
230
+ k_embed = (k * cos) + (rotate_half(k) * sin)
231
+ return q_embed, k_embed
232
 
233
 
234
+ # Copied from transformers.models.mistral.modeling_mistral.MistralMLP with Mistral->StableLm
235
+ class StableLmMLP(nn.Module):
236
+ def __init__(self, config):
237
  super().__init__()
238
  self.config = config
239
  self.hidden_size = config.hidden_size
240
  self.intermediate_size = config.intermediate_size
241
+ self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
242
+ self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
243
+ self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
244
+ self.act_fn = ACT2FN[config.hidden_act]
 
 
 
 
 
 
245
 
246
+ def forward(self, x):
247
  return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
248
 
249
 
250
+ class StableLmLayerNormPerHead(nn.Module):
251
+ def __init__(self, dim, num_heads, eps=1e-5, bias=False):
252
+ super().__init__()
253
+ self.dim = dim
254
+ self.num_heads = num_heads
255
+ self.norms = nn.ModuleList(
256
+ [nn.LayerNorm(dim, eps=eps, bias=bias) for _ in range(self.num_heads)]
257
+ )
258
+
259
+ def forward(self, hidden_states: torch.Tensor):
260
+ # Split along the num_heads axis to get per-head inputs
261
+ # [batch_size, num_heads, seq_len, head_dim] -> [batch_size, 1, seq_len, head_dim] * num_heads
262
+ states_per_heads = torch.split(hidden_states, 1, dim=1)
263
+ # Normalize and merge the heads back together
264
+ return torch.cat(
265
+ [
266
+ norm(hidden_states)
267
+ for norm, hidden_states in zip(self.norms, states_per_heads)
268
+ ],
269
+ dim=1,
270
+ )
271
+
272
+
273
+ # Copied from transformers.models.llama.modeling_llama.repeat_kv
274
  def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
275
  """
276
  This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
 
285
  return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
286
 
287
 
288
+ class StableLmAttention(nn.Module):
289
+ """Multi-headed attention from 'Attention Is All You Need' paper"""
290
+
291
+ def __init__(self, config: StableLmConfig, layer_idx: Optional[int] = None):
292
  super().__init__()
293
  self.config = config
294
+ self.layer_idx = layer_idx
295
+ if layer_idx is None:
296
+ logger.warning_once(
297
+ f"Instantiating {self.__class__.__name__} without passing a `layer_idx` is not recommended and will "
298
+ "lead to errors during the forward call if caching is used. Please make sure to provide a `layer_idx` "
299
+ "when creating this class."
300
+ )
301
+
302
  self.hidden_size = config.hidden_size
303
  self.num_heads = config.num_attention_heads
304
  self.head_dim = self.hidden_size // self.num_heads
305
  self.num_key_value_heads = config.num_key_value_heads
306
  self.num_key_value_groups = self.num_heads // self.num_key_value_heads
307
  self.max_position_embeddings = config.max_position_embeddings
308
+ self.rope_theta = config.rope_theta
309
+ self.partial_rotary_factor = config.partial_rotary_factor
310
  self.is_causal = True
 
311
 
312
  if (self.head_dim * self.num_heads) != self.hidden_size:
313
  raise ValueError(
314
  f"hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size}"
315
  f" and `num_heads`: {self.num_heads})."
316
  )
 
317
  self.q_proj = nn.Linear(
318
  self.hidden_size, self.num_heads * self.head_dim, bias=config.use_qkv_bias
319
  )
 
329
  )
330
  self.o_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=False)
331
 
332
+ self.qk_layernorm = config.qk_layernorm
333
+ if self.qk_layernorm:
334
+ self.q_layernorm = StableLmLayerNormPerHead(
335
+ self.head_dim, self.num_heads, eps=config.layer_norm_eps
336
+ )
337
+ self.k_layernorm = StableLmLayerNormPerHead(
338
+ self.head_dim, self.num_key_value_heads, eps=config.layer_norm_eps
339
+ )
340
 
341
+ self.attention_dropout = nn.Dropout(config.attention_dropout)
342
  self._init_rope()
343
 
344
+ # Copied from transformers.models.persimmon.modeling_persimmon.PersimmonAttention._init_rope with Persimmon->StableLm
345
  def _init_rope(self):
346
+ if self.config.rope_scaling is None:
347
+ self.rotary_emb = StableLmRotaryEmbedding(
348
+ int(self.partial_rotary_factor * self.head_dim),
349
+ max_position_embeddings=self.max_position_embeddings,
350
+ base=self.rope_theta,
351
+ )
352
+ else:
353
+ scaling_type = self.config.rope_scaling["type"]
354
+ scaling_factor = self.config.rope_scaling["factor"]
355
+ if scaling_type == "linear":
356
+ self.rotary_emb = StableLmLinearScalingRotaryEmbedding(
357
+ int(self.partial_rotary_factor * self.head_dim),
358
+ max_position_embeddings=self.max_position_embeddings,
359
+ scaling_factor=scaling_factor,
360
+ base=self.rope_theta,
361
+ )
362
+ elif scaling_type == "dynamic":
363
+ self.rotary_emb = StableLmDynamicNTKScalingRotaryEmbedding(
364
+ int(self.partial_rotary_factor * self.head_dim),
365
+ max_position_embeddings=self.max_position_embeddings,
366
+ scaling_factor=scaling_factor,
367
+ base=self.rope_theta,
368
+ )
369
+ else:
370
+ raise ValueError(f"Unknown RoPE scaling type {scaling_type}")
371
 
372
  def forward(
373
  self,
374
+ hidden_states: torch.Tensor,
375
+ attention_mask: Optional[torch.Tensor] = None,
376
+ position_ids: Optional[torch.LongTensor] = None,
377
+ past_key_value: Optional[Cache] = None,
378
+ output_attentions: bool = False,
379
+ use_cache: bool = False,
380
  ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
381
  bsz, q_len, _ = hidden_states.size()
382
 
 
394
  bsz, q_len, self.num_key_value_heads, self.head_dim
395
  ).transpose(1, 2)
396
 
397
+ if self.qk_layernorm:
398
+ query_states = self.q_layernorm(query_states)
399
+ key_states = self.k_layernorm(key_states)
 
 
 
 
 
400
 
401
  kv_seq_len = key_states.shape[-2]
402
  if past_key_value is not None:
403
+ if self.layer_idx is None:
404
+ raise ValueError(
405
+ f"The cache structure has changed since version v4.36. If you are using {self.__class__.__name__} "
406
+ "for auto-regressive decoding with k/v caching, please make sure to initialize the attention class "
407
+ "with a layer index."
408
+ )
409
+ kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx)
410
  cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
411
+
412
+ # Partial rotary embedding
413
+ query_rot, query_pass = (
414
+ query_states[..., : self.rotary_emb.dim],
415
+ query_states[..., self.rotary_emb.dim :],
416
+ )
417
+ key_rot, key_pass = (
418
+ key_states[..., : self.rotary_emb.dim],
419
+ key_states[..., self.rotary_emb.dim :],
420
+ )
421
+ # [batch_size, seq_length, num_heads, head_dim // config.partial_rotary_factor]
422
+ query_rot, key_rot = apply_rotary_pos_emb(
423
  query_rot, key_rot, cos, sin, position_ids
424
  )
425
 
426
+ # [batch_size, seq_length, num_heads, head_dim]
427
+ query_states = torch.cat((query_rot, query_pass), dim=-1)
428
+ key_states = torch.cat((key_rot, key_pass), dim=-1)
429
 
430
  if past_key_value is not None:
431
+ # Specific to RoPE models with partial rotation
432
+ cache_kwargs = {
433
+ "sin": sin,
434
+ "cos": cos,
435
+ "partial_rotation_size": self.rotary_emb.dim,
436
+ }
437
+ key_states, value_states = past_key_value.update(
438
+ key_states, value_states, self.layer_idx, cache_kwargs
439
+ )
440
 
441
  # Repeat k/v heads if n_kv_heads < n_heads
442
  key_states = repeat_kv(key_states, self.num_key_value_groups)
 
459
  )
460
  attn_weights = attn_weights + attention_mask
461
 
462
+ # upcast attention to fp32
463
  attn_weights = nn.functional.softmax(
464
+ attn_weights, dtype=torch.float32, dim=-1
465
  ).to(query_states.dtype)
466
+ attn_weights = self.attention_dropout(attn_weights)
467
+
 
468
  attn_output = torch.matmul(attn_weights, value_states)
469
 
470
  if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim):
 
473
  f" {attn_output.size()}"
474
  )
475
 
 
476
  attn_output = attn_output.transpose(1, 2).contiguous()
477
  attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)
478
 
 
479
  attn_output = self.o_proj(attn_output)
480
 
481
  if not output_attentions:
 
484
  return attn_output, attn_weights, past_key_value
485
 
486
 
487
+ class StableLmSdpaAttention(StableLmAttention):
488
+ def forward(
489
+ self,
490
+ hidden_states: torch.Tensor,
491
+ attention_mask: Optional[torch.Tensor] = None,
492
+ position_ids: Optional[torch.LongTensor] = None,
493
+ past_key_value: Optional[Cache] = None,
494
+ output_attentions: bool = False,
495
+ use_cache: bool = False,
496
+ ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
497
+ if output_attentions:
498
+ # TODO: Improve this warning with e.g. `model.config.attn_implementation = "manual"` once this is implemented.
499
+ logger.warning_once(
500
+ "StableLmModel is using StableLmSdpaAttention, but `torch.nn.functional.scaled_dot_product_attention` does not support `output_attentions=True`. Falling back to the manual attention implementation, "
501
+ '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.'
502
+ )
503
+ return super().forward(
504
+ hidden_states=hidden_states,
505
+ attention_mask=attention_mask,
506
+ position_ids=position_ids,
507
+ past_key_value=past_key_value,
508
+ output_attentions=output_attentions,
509
+ use_cache=use_cache,
510
+ )
511
+
512
+ bsz, q_len, _ = hidden_states.size()
513
+
514
+ query_states = self.q_proj(hidden_states)
515
+ key_states = self.k_proj(hidden_states)
516
+ value_states = self.v_proj(hidden_states)
517
+
518
+ query_states = query_states.view(
519
+ bsz, q_len, self.num_heads, self.head_dim
520
+ ).transpose(1, 2)
521
+ key_states = key_states.view(
522
+ bsz, q_len, self.num_key_value_heads, self.head_dim
523
+ ).transpose(1, 2)
524
+ value_states = value_states.view(
525
+ bsz, q_len, self.num_key_value_heads, self.head_dim
526
+ ).transpose(1, 2)
527
+
528
+ if self.qk_layernorm:
529
+ query_states = self.q_layernorm(query_states)
530
+ key_states = self.k_layernorm(key_states)
531
+
532
+ kv_seq_len = key_states.shape[-2]
533
+ if past_key_value is not None:
534
+ if self.layer_idx is None:
535
+ raise ValueError(
536
+ f"The cache structure has changed since version v4.36. If you are using {self.__class__.__name__} "
537
+ "for auto-regressive decoding with k/v caching, please make sure to initialize the attention class "
538
+ "with a layer index."
539
+ )
540
+ kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx)
541
+ cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
542
+
543
+ # Partial rotary embedding
544
+ query_rot, query_pass = (
545
+ query_states[..., : self.rotary_emb.dim],
546
+ query_states[..., self.rotary_emb.dim :],
547
+ )
548
+ key_rot, key_pass = (
549
+ key_states[..., : self.rotary_emb.dim],
550
+ key_states[..., self.rotary_emb.dim :],
551
+ )
552
+ # [batch_size, seq_length, num_heads, head_dim // config.partial_rotary_factor]
553
+ query_rot, key_rot = apply_rotary_pos_emb(
554
+ query_rot, key_rot, cos, sin, position_ids
555
+ )
556
+
557
+ # [batch_size, seq_length, num_heads, head_dim]
558
+ query_states = torch.cat((query_rot, query_pass), dim=-1)
559
+ key_states = torch.cat((key_rot, key_pass), dim=-1)
560
+
561
+ if past_key_value is not None:
562
+ # Specific to RoPE models with partial rotation
563
+ cache_kwargs = {
564
+ "sin": sin,
565
+ "cos": cos,
566
+ "partial_rotation_size": self.rotary_emb.dim,
567
+ }
568
+ key_states, value_states = past_key_value.update(
569
+ key_states, value_states, self.layer_idx, cache_kwargs
570
+ )
571
+
572
+ # Repeat k/v heads if n_kv_heads < n_heads
573
+ key_states = repeat_kv(key_states, self.num_key_value_groups)
574
+ value_states = repeat_kv(value_states, self.num_key_value_groups)
575
+
576
+ # SDPA with memory-efficient backend is currently (torch==2.1.2) bugged with non-contiguous inputs with custom attn_mask,
577
+ # Reference: https://github.com/pytorch/pytorch/issues/112577.
578
+ if query_states.device.type == "cuda" and attention_mask is not None:
579
+ query_states = query_states.contiguous()
580
+ key_states = key_states.contiguous()
581
+ value_states = value_states.contiguous()
582
+
583
+ attn_output = torch.nn.functional.scaled_dot_product_attention(
584
+ query_states,
585
+ key_states,
586
+ value_states,
587
+ attn_mask=attention_mask,
588
+ dropout_p=self.attention_dropout.p if self.training else 0.0,
589
+ # 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.
590
+ is_causal=self.is_causal and attention_mask is None and q_len > 1,
591
+ )
592
+
593
+ attn_output = attn_output.transpose(1, 2).contiguous()
594
+ attn_output = attn_output.view(bsz, q_len, self.hidden_size)
595
+
596
+ attn_output = self.o_proj(attn_output)
597
+
598
+ return attn_output, None, past_key_value
599
+
600
+
601
+ class StableLmFlashAttention2(StableLmAttention):
602
  """
603
+ StableLM flash attention module. This module inherits from `StableLmAttention` as the weights of the module stays
604
+ untouched. The only required change would be on the forward pass where it needs to correctly call the public API of
605
+ flash attention and deal with padding tokens in case the input contains any of them.
606
  """
607
 
608
+ # Copied from transformers.models.llama.modeling_llama.LlamaFlashAttention2.__init__
609
  def __init__(self, *args, **kwargs):
610
  super().__init__(*args, **kwargs)
611
 
612
  # TODO: Should be removed once Flash Attention for RoCm is bumped to 2.1.
613
+ # flash_attn<2.1 generates top-left aligned causal mask, while what is needed here is bottom-right alignment, 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.
614
  # 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).
615
  self._flash_attn_uses_top_left_mask = not is_flash_attn_greater_or_equal_2_10()
616
 
 
624
  use_cache: bool = False,
625
  **kwargs,
626
  ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
627
+ # StableLmFlashAttention2 attention does not support output_attentions
 
 
 
 
 
 
 
628
 
629
  output_attentions = False
630
 
 
647
  bsz, q_len, self.num_key_value_heads, self.head_dim
648
  ).transpose(1, 2)
649
 
650
+ if self.qk_layernorm:
651
+ query_states = self.q_layernorm(query_states)
652
+ key_states = self.k_layernorm(key_states)
 
 
 
 
 
653
 
654
  kv_seq_len = key_states.shape[-2]
655
  if past_key_value is not None:
656
+ if self.layer_idx is None:
657
+ raise ValueError(
658
+ f"The cache structure has changed since version v4.36. If you are using {self.__class__.__name__} "
659
+ "for auto-regressive decoding with k/v caching, please make sure to initialize the attention class "
660
+ "with a layer index."
661
+ )
662
+ kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx)
663
  cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
664
+
665
+ # Partial rotary embedding
666
+ query_rot, query_pass = (
667
+ query_states[..., : self.rotary_emb.dim],
668
+ query_states[..., self.rotary_emb.dim :],
669
+ )
670
+ key_rot, key_pass = (
671
+ key_states[..., : self.rotary_emb.dim],
672
+ key_states[..., self.rotary_emb.dim :],
673
+ )
674
+ query_rot, key_rot = apply_rotary_pos_emb(
675
  query_rot, key_rot, cos, sin, position_ids
676
  )
677
 
678
+ # [batch_size, seq_length, num_heads, head_dim]
679
+ query_states = torch.cat((query_rot, query_pass), dim=-1)
680
+ key_states = torch.cat((key_rot, key_pass), dim=-1)
681
 
682
  if past_key_value is not None:
683
+ cache_kwargs = {
684
+ "sin": sin,
685
+ "cos": cos,
686
+ "partial_rotation_size": self.rotary_emb.dim,
687
+ }
688
+ key_states, value_states = past_key_value.update(
689
+ key_states, value_states, self.layer_idx, cache_kwargs
690
+ )
691
 
692
  # TODO: These transpose are quite inefficient but Flash Attention requires the layout [batch_size, sequence_length, num_heads, head_dim]. We would need to refactor the KV cache
693
  # to be able to avoid many of these transpose/reshape/view.
 
695
  key_states = key_states.transpose(1, 2)
696
  value_states = value_states.transpose(1, 2)
697
 
698
+ dropout_rate = self.attention_dropout.p if self.training else 0.0
699
 
700
  attn_output = self._flash_attention_forward(
701
  query_states,
 
705
  q_len,
706
  dropout=dropout_rate,
707
  )
708
+
709
  attn_output = attn_output.reshape(bsz, q_len, self.hidden_size).contiguous()
710
  attn_output = self.o_proj(attn_output)
711
 
 
714
 
715
  return attn_output, attn_weights, past_key_value
716
 
717
+ # Copied from transformers.models.llama.modeling_llama.LlamaFlashAttention2._flash_attention_forward
718
  def _flash_attention_forward(
719
  self,
720
  query_states,
 
739
  attention_mask (`torch.Tensor`):
740
  The padding mask - corresponds to a tensor of size `(batch_size, seq_len)` where 0 stands for the
741
  position of padding tokens and 1 for the position of non-padding tokens.
742
+ dropout (`float`):
743
  Attention dropout
744
  softmax_scale (`float`, *optional*):
745
  The scaling of QK^T before applying softmax. Default to 1 / sqrt(head_dim)
 
747
  if not self._flash_attn_uses_top_left_mask:
748
  causal = self.is_causal
749
  else:
750
+ # 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__.
751
  causal = self.is_causal and query_length != 1
752
 
753
  # Contains at least one padding token in the sequence
 
795
 
796
  return attn_output
797
 
798
+ # Copied from transformers.models.llama.modeling_llama.LlamaFlashAttention2._upad_input
799
  def _upad_input(
800
  self, query_layer, key_layer, value_layer, attention_mask, query_length
801
  ):
 
843
 
844
 
845
  ATTENTION_CLASSES = {
846
+ "eager": StableLmAttention,
847
+ "sdpa": StableLmSdpaAttention,
848
+ "flash_attention_2": StableLmFlashAttention2,
849
  }
850
 
851
 
852
+ class StableLmDecoderLayer(nn.Module):
853
+ def __init__(self, config: StableLmConfig, layer_idx: int):
854
  super().__init__()
855
+ self.use_parallel_residual = config.use_parallel_residual
856
+ self.hidden_size = config.hidden_size
857
+ self.self_attn = ATTENTION_CLASSES[config._attn_implementation](
858
+ config, layer_idx=layer_idx
859
+ )
860
+ self.mlp = StableLmMLP(config)
861
  self.input_layernorm = nn.LayerNorm(
862
+ config.hidden_size, eps=config.layer_norm_eps
863
  )
864
+ self.post_attention_layernorm = None
865
+ if not self.use_parallel_residual:
866
+ self.post_attention_layernorm = nn.LayerNorm(
867
+ config.hidden_size, eps=config.layer_norm_eps
868
+ )
869
+ self.dropout = nn.Dropout(config.hidden_dropout)
870
 
871
  def forward(
872
  self,
873
+ hidden_states: torch.Tensor,
874
+ attention_mask: Optional[torch.Tensor] = None,
875
  position_ids: Optional[torch.LongTensor] = None,
876
  past_key_value: Optional[Tuple[torch.Tensor]] = None,
877
  output_attentions: Optional[bool] = False,
878
  use_cache: Optional[bool] = False,
879
+ ) -> Tuple[
880
+ torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]
 
881
  ]:
882
+ """
883
+ Args:
884
+ hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`
885
+ attention_mask (`torch.FloatTensor`, *optional*): attention mask of size
886
+ `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values.
887
+ position_ids (`torch.LongTensor` of shape `({0})`, *optional*):
888
+ Indices of positions of each input sequence tokens in the position embeddings. Selected in the range
889
+ `[0, config.n_positions - 1]`.
890
+
891
+ [What are position IDs?](../glossary#position-ids)
892
+ past_key_value (`Tuple(torch.FloatTensor)`, *optional*):
893
+ cached past key and value projection states
894
+ output_attentions (`bool`, *optional*):
895
+ Whether or not to return the attentions tensors of all attention layers. See `attentions` under
896
+ returned tensors for more detail.
897
+ use_cache (`bool`, *optional*):
898
+ If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding
899
+ (see `past_key_values`).
900
+ """
901
+
902
  residual = hidden_states
903
 
904
  hidden_states = self.input_layernorm(hidden_states)
 
913
  use_cache=use_cache,
914
  )
915
 
916
+ # copied from transformers.models.gpt_neox.modeling_gpt_neox.GPTNeoXLayer.forward
917
+ if self.use_parallel_residual:
918
+ # x = x + attn(ln1(x)) + mlp(ln1(x))
919
+ # Fully Connected
920
+ mlp_output = self.mlp(hidden_states)
921
+ mlp_output = self.dropout(mlp_output)
922
+ hidden_states = residual + self_attn_output + mlp_output
923
+ else:
924
+ # x = x + attn(ln1(x))
925
+ # x = x + mlp(ln2(x))
926
+ residual = residual + self_attn_output
927
+ # Fully Connected
928
+ mlp_output = self.mlp(self.post_attention_layernorm(residual))
929
+ mlp_output = self.dropout(mlp_output)
930
+ hidden_states = residual + mlp_output
931
 
932
  outputs = (hidden_states,)
933
 
 
940
  return outputs
941
 
942
 
943
+ STABLELM_START_DOCSTRING = r"""
944
+ This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the
945
+ library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
946
+ etc.)
947
+
948
+ This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.
949
+ Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage
950
+ and behavior.
951
+
952
+ Parameters:
953
+ config ([`StableLmConfig`]):
954
+ Model configuration class with all the parameters of the model. Initializing with a config file does not
955
+ load the weights associated with the model, only the configuration. Check out the
956
+ [`~PreTrainedModel.from_pretrained`] method to load the model weights.
957
+ """
958
 
959
+
960
+ @add_start_docstrings(
961
+ "The bare StableLm Model outputting raw hidden-states without any specific head on top.",
962
+ STABLELM_START_DOCSTRING,
963
+ )
964
+ class StableLmPreTrainedModel(PreTrainedModel):
965
+ config_class = StableLmConfig
966
+ base_model_prefix = "model"
967
  supports_gradient_checkpointing = True
968
+ _no_split_modules = ["StableLmDecoderLayer"]
969
  _skip_keys_device_placement = "past_key_values"
970
  _supports_flash_attn_2 = True
971
+ _supports_cache_class = True
972
+ _supports_sdpa = True
973
 
974
+ def _init_weights(self, module):
975
+ std = self.config.initializer_range
976
  if isinstance(module, nn.Linear):
977
+ module.weight.data.normal_(mean=0.0, std=std)
978
  if module.bias is not None:
979
  module.bias.data.zero_()
980
  elif isinstance(module, nn.Embedding):
981
+ module.weight.data.normal_(mean=0.0, std=std)
982
  if module.padding_idx is not None:
983
  module.weight.data[module.padding_idx].zero_()
 
 
 
 
984
 
 
 
 
985
 
986
+ STABLELM_INPUTS_DOCSTRING = r"""
987
+ Args:
988
+ input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
989
+ Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide
990
+ it.
991
+
992
+ Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
993
+ [`PreTrainedTokenizer.__call__`] for details.
994
+
995
+ [What are input IDs?](../glossary#input-ids)
996
+ attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
997
+ Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
998
+
999
+ - 1 for tokens that are **not masked**,
1000
+ - 0 for tokens that are **masked**.
1001
+
1002
+ [What are attention masks?](../glossary#attention-mask)
1003
+
1004
+ Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
1005
+ [`PreTrainedTokenizer.__call__`] for details.
1006
+
1007
+ If `past_key_values` is used, optionally only the last `decoder_input_ids` have to be input (see
1008
+ `past_key_values`).
1009
+
1010
+ If you want to change padding behavior, you should read [`modeling_opt._prepare_decoder_attention_mask`]
1011
+ and modify to your needs. See diagram 1 in [the paper](https://arxiv.org/abs/1910.13461) for more
1012
+ information on the default strategy.
1013
+
1014
+ - 1 indicates the head is **not masked**,
1015
+ - 0 indicates the head is **masked**.
1016
+ position_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
1017
+ Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,
1018
+ config.n_positions - 1]`.
1019
+
1020
+ [What are position IDs?](../glossary#position-ids)
1021
+ past_key_values (`Cache` or `tuple(tuple(torch.FloatTensor))`, *optional*):
1022
+ Pre-computed hidden-states (key and values in the self-attention blocks and in the cross-attention
1023
+ blocks) that can be used to speed up sequential decoding. This typically consists in the `past_key_values`
1024
+ returned by the model at a previous stage of decoding, when `use_cache=True` or `config.use_cache=True`.
1025
+
1026
+ Two formats are allowed:
1027
+ - a [`~cache_utils.Cache`] instance;
1028
+ - Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of
1029
+ shape `(batch_size, num_heads, sequence_length, embed_size_per_head)`). This is also known as the legacy
1030
+ cache format.
1031
+
1032
+ The model will output the same cache format that is fed as input. If no `past_key_values` are passed, the
1033
+ legacy cache format will be returned.
1034
+
1035
+ If `past_key_values` are used, the user can optionally input only the last `input_ids` (those that don't
1036
+ have their past key value states given to this model) of shape `(batch_size, 1)` instead of all `input_ids`
1037
+ of shape `(batch_size, sequence_length)`.
1038
+ inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
1039
+ Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This
1040
+ is useful if you want more control over how to convert `input_ids` indices into associated vectors than the
1041
+ model's internal embedding lookup matrix.
1042
+ use_cache (`bool`, *optional*):
1043
+ If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see
1044
+ `past_key_values`).
1045
+ output_attentions (`bool`, *optional*):
1046
+ Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
1047
+ tensors for more detail.
1048
+ output_hidden_states (`bool`, *optional*):
1049
+ Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
1050
+ more detail.
1051
+ return_dict (`bool`, *optional*):
1052
+ Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
1053
+ """
1054
+
1055
+
1056
+ @add_start_docstrings(
1057
+ "The bare StableLm Model outputting raw hidden-states without any specific head on top.",
1058
+ STABLELM_START_DOCSTRING,
1059
+ )
1060
+ class StableLmModel(StableLmPreTrainedModel):
1061
+ """
1062
+ Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`StableLmDecoderLayer`]
1063
 
1064
+ Args:
1065
+ config: StableLmConfig
1066
+ """
1067
+
1068
+ def __init__(self, config: StableLmConfig):
1069
  super().__init__(config)
1070
+ self.padding_idx = config.pad_token_id
1071
+ self.vocab_size = config.vocab_size
1072
+
1073
  self.embed_tokens = nn.Embedding(
1074
+ config.vocab_size, config.hidden_size, self.padding_idx
1075
  )
1076
  self.layers = nn.ModuleList(
1077
+ [
1078
+ StableLmDecoderLayer(config, layer_idx)
1079
+ for layer_idx in range(config.num_hidden_layers)
1080
+ ]
1081
  )
1082
+ self.norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
1083
 
1084
+ self._attn_implementation = config._attn_implementation
1085
  self.gradient_checkpointing = False
1086
  # Initialize weights and apply final processing
1087
  self.post_init()
 
1089
  def get_input_embeddings(self):
1090
  return self.embed_tokens
1091
 
1092
+ def set_input_embeddings(self, value):
1093
  self.embed_tokens = value
1094
 
1095
+ @add_start_docstrings_to_model_forward(STABLELM_INPUTS_DOCSTRING)
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1096
  def forward(
1097
  self,
1098
+ input_ids: torch.LongTensor = None,
1099
+ attention_mask: Optional[torch.Tensor] = None,
1100
  position_ids: Optional[torch.LongTensor] = None,
1101
+ past_key_values: Optional[List[torch.FloatTensor]] = None,
1102
  inputs_embeds: Optional[torch.FloatTensor] = None,
1103
  use_cache: Optional[bool] = None,
1104
  output_attentions: Optional[bool] = None,
 
1121
  return_dict if return_dict is not None else self.config.use_return_dict
1122
  )
1123
 
1124
+ # retrieve input_ids and inputs_embeds
1125
  if input_ids is not None and inputs_embeds is not None:
1126
  raise ValueError(
1127
  "You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time"
 
1138
  seq_length_with_past = seq_length
1139
  past_key_values_length = 0
1140
 
1141
+ if self.gradient_checkpointing and self.training:
1142
+ if use_cache:
1143
+ logger.warning_once(
1144
+ "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
1145
+ )
1146
+ use_cache = False
1147
+
1148
+ if use_cache:
1149
+ use_legacy_cache = not isinstance(past_key_values, Cache)
1150
+ if use_legacy_cache:
1151
+ past_key_values = DynamicCache.from_legacy_cache(past_key_values)
1152
+ past_key_values_length = past_key_values.get_usable_length(seq_length)
1153
+ seq_length_with_past = seq_length_with_past + past_key_values_length
1154
+
1155
  if position_ids is None:
1156
  device = input_ids.device if input_ids is not None else inputs_embeds.device
1157
  position_ids = torch.arange(
 
1160
  dtype=torch.long,
1161
  device=device,
1162
  )
1163
+ position_ids = position_ids.unsqueeze(0)
 
 
1164
 
1165
  if inputs_embeds is None:
1166
  inputs_embeds = self.embed_tokens(input_ids)
1167
+ # embed positions
1168
+ if self._attn_implementation == "flash_attention_2":
1169
  # 2d mask is passed through the layers
1170
  attention_mask = (
1171
  attention_mask
1172
  if (attention_mask is not None and 0 in attention_mask)
1173
  else None
1174
  )
1175
+ # for output_attentions case used fallback to eager attention realization
1176
+ elif self._attn_implementation == "sdpa" and not output_attentions:
1177
+ attention_mask = _prepare_4d_causal_attention_mask_for_sdpa(
1178
+ attention_mask,
1179
+ (batch_size, seq_length),
1180
+ inputs_embeds,
1181
+ past_key_values_length,
1182
+ )
1183
  else:
1184
+ # 4d mask is passed through the layers
1185
+ attention_mask = _prepare_4d_causal_attention_mask(
 
 
 
 
 
1186
  attention_mask,
1187
  (batch_size, seq_length),
1188
  inputs_embeds,
 
1191
 
1192
  hidden_states = inputs_embeds
1193
 
1194
+ # decoder layers
 
 
 
 
 
 
 
1195
  all_hidden_states = () if output_hidden_states else None
1196
  all_self_attns = () if output_attentions else None
1197
+ next_decoder_cache = None
1198
 
1199
+ for decoder_layer in self.layers:
1200
  if output_hidden_states:
1201
  all_hidden_states += (hidden_states,)
1202
 
 
 
 
 
1203
  if self.gradient_checkpointing and self.training:
1204
+ layer_outputs = self._gradient_checkpointing_func(
1205
+ decoder_layer.__call__,
 
 
 
 
 
 
 
 
1206
  hidden_states,
1207
  attention_mask,
1208
  position_ids,
1209
+ past_key_values,
1210
+ output_attentions,
1211
  )
1212
  else:
1213
  layer_outputs = decoder_layer(
1214
  hidden_states,
1215
  attention_mask=attention_mask,
1216
  position_ids=position_ids,
1217
+ past_key_value=past_key_values,
1218
  output_attentions=output_attentions,
1219
  use_cache=use_cache,
1220
  )
 
1222
  hidden_states = layer_outputs[0]
1223
 
1224
  if use_cache:
1225
+ next_decoder_cache = layer_outputs[2 if output_attentions else 1]
1226
 
1227
  if output_attentions:
1228
  all_self_attns += (layer_outputs[1],)
1229
 
1230
  hidden_states = self.norm(hidden_states)
1231
 
1232
+ # add hidden states from the last decoder layer
1233
  if output_hidden_states:
1234
  all_hidden_states += (hidden_states,)
1235
 
1236
+ next_cache = None
1237
+ if use_cache:
1238
+ next_cache = (
1239
+ next_decoder_cache.to_legacy_cache()
1240
+ if use_legacy_cache
1241
+ else next_decoder_cache
1242
+ )
1243
+
1244
  if not return_dict:
1245
  return tuple(
1246
  v
 
1255
  )
1256
 
1257
 
1258
+ # Copied from transformers.models.persimmon.modeling_persimmon.PersimmonForCausalLM with PERSIMMON->STABLELM,Persimmon->StableLm
1259
+ class StableLmForCausalLM(StableLmPreTrainedModel):
1260
  _tied_weights_keys = ["lm_head.weight"]
1261
 
1262
+ # Copied from transformers.models.llama.modeling_llama.LlamaForCausalLM.__init__ with LLAMA->STABLELM,Llama->StableLm
1263
+ def __init__(self, config):
1264
  super().__init__(config)
1265
+ self.model = StableLmModel(config)
1266
+ self.vocab_size = config.vocab_size
1267
  self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
1268
 
1269
  # Initialize weights and apply final processing
1270
  self.post_init()
1271
 
1272
+ # Copied from transformers.models.llama.modeling_llama.LlamaForCausalLM.get_input_embeddings
1273
  def get_input_embeddings(self):
1274
  return self.model.embed_tokens
1275
 
1276
+ # Copied from transformers.models.llama.modeling_llama.LlamaForCausalLM.set_input_embeddings
1277
  def set_input_embeddings(self, value):
1278
  self.model.embed_tokens = value
1279
 
1280
+ # Copied from transformers.models.llama.modeling_llama.LlamaForCausalLM.get_output_embeddings
1281
  def get_output_embeddings(self):
1282
  return self.lm_head
1283
 
1284
+ # Copied from transformers.models.llama.modeling_llama.LlamaForCausalLM.set_output_embeddings
1285
+ def set_output_embeddings(self, new_embeddings):
1286
  self.lm_head = new_embeddings
1287
 
1288
+ # Copied from transformers.models.llama.modeling_llama.LlamaForCausalLM.set_decoder
 
 
1289
  def set_decoder(self, decoder):
1290
  self.model = decoder
1291
 
1292
+ # Copied from transformers.models.llama.modeling_llama.LlamaForCausalLM.get_decoder
1293
+ def get_decoder(self):
1294
+ return self.model
1295
+
1296
+ @add_start_docstrings_to_model_forward(STABLELM_INPUTS_DOCSTRING)
1297
+ @replace_return_docstrings(
1298
+ output_type=CausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC
1299
+ )
1300
+ # Ignore copy
1301
  def forward(
1302
  self,
1303
+ input_ids: torch.LongTensor = None,
1304
+ attention_mask: Optional[torch.Tensor] = None,
1305
  position_ids: Optional[torch.LongTensor] = None,
1306
+ past_key_values: Optional[List[torch.FloatTensor]] = None,
1307
  inputs_embeds: Optional[torch.FloatTensor] = None,
1308
  labels: Optional[torch.LongTensor] = None,
1309
  use_cache: Optional[bool] = None,
 
1311
  output_hidden_states: Optional[bool] = None,
1312
  return_dict: Optional[bool] = None,
1313
  ) -> Union[Tuple, CausalLMOutputWithPast]:
1314
+ r"""
1315
+ Args:
1316
+ labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
1317
+ Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
1318
+ config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
1319
+ (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
1320
+
1321
+ Returns:
1322
+
1323
+ Example:
1324
+
1325
+ ```python
1326
+ >>> from transformers import AutoTokenizer, StableLmForCausalLM
1327
+
1328
+ >>> model = StableLmForCausalLM.from_pretrained("stabilityai/stablelm-3b-4e1t")
1329
+ >>> tokenizer = AutoTokenizer.from_pretrained("stabilityai/stablelm-3b-4e1t")
1330
+
1331
+ >>> prompt = "The weather is always wonderful in"
1332
+ >>> inputs = tokenizer(prompt, return_tensors="pt")
1333
+
1334
+ >>> # Generate
1335
+ >>> generate_ids = model.generate(inputs.input_ids, max_length=30)
1336
+ >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
1337
+ 'The weather is always wonderful in the summer in the city of San Diego. The city is located on the coast of the Pacific Ocean and is surrounded by'
1338
+ ```"""
1339
+
1340
  output_attentions = (
1341
  output_attentions
1342
  if output_attentions is not None
 
1351
  return_dict if return_dict is not None else self.config.use_return_dict
1352
  )
1353
 
 
1354
  outputs = self.model(
1355
+ input_ids=input_ids,
1356
  attention_mask=attention_mask,
1357
  position_ids=position_ids,
1358
  past_key_values=past_key_values,
 
1364
  )
1365
 
1366
  hidden_states = outputs[0]
1367
+ logits = self.lm_head(hidden_states)
1368
 
1369
  loss = None
1370
  if labels is not None:
 
1394
  def prepare_inputs_for_generation(
1395
  self,
1396
  input_ids,
1397
+ past_key_values=None,
1398
+ attention_mask=None,
1399
+ inputs_embeds=None,
1400
  **kwargs,
1401
  ):
 
1402
  if past_key_values is not None:
1403
+ if isinstance(past_key_values, Cache):
1404
+ cache_length = past_key_values.get_seq_length()
1405
+ past_length = past_key_values.seen_tokens
1406
+ max_cache_length = past_key_values.get_max_length()
 
1407
  else:
1408
+ cache_length = past_length = past_key_values[0][0].shape[2]
1409
+ max_cache_length = None
1410
+
1411
+ # Keep only the unprocessed tokens:
1412
+ # 1 - If the length of the attention_mask exceeds the length of input_ids, then we are in a setting where
1413
+ # some of the inputs are exclusively passed as part of the cache (e.g. when passing input_embeds as
1414
+ # input)
1415
+ if (
1416
+ attention_mask is not None
1417
+ and attention_mask.shape[1] > input_ids.shape[1]
1418
+ ):
1419
+ input_ids = input_ids[:, -(attention_mask.shape[1] - past_length) :]
1420
+ # 2 - If the past_length is smaller than input_ids', then input_ids holds all input tokens. We can discard
1421
+ # input_ids based on the past_length.
1422
+ elif past_length < input_ids.shape[1]:
1423
+ input_ids = input_ids[:, past_length:]
1424
+ # 3 - Otherwise (past_length >= input_ids.shape[1]), let's assume input_ids only has unprocessed tokens.
1425
+
1426
+ # If we are about to go beyond the maximum cache length, we need to crop the input attention mask.
1427
+ if (
1428
+ max_cache_length is not None
1429
+ and attention_mask is not None
1430
+ and cache_length + input_ids.shape[1] > max_cache_length
1431
+ ):
1432
+ attention_mask = attention_mask[:, -max_cache_length:]
1433
 
1434
  position_ids = kwargs.get("position_ids", None)
1435
  if attention_mask is not None and position_ids is None:
1436
+ # create position_ids on the fly for batch generation
1437
  position_ids = attention_mask.long().cumsum(-1) - 1
1438
  position_ids.masked_fill_(attention_mask == 0, 1)
1439
  if past_key_values:
1440
+ position_ids = position_ids[:, -input_ids.shape[1] :]
1441
 
1442
+ # if `inputs_embeds` are passed, we only want to use them in the 1st generation step
1443
  if inputs_embeds is not None and past_key_values is None:
1444
  model_inputs = {"inputs_embeds": inputs_embeds}
1445
  else:
 
1447
 
1448
  model_inputs.update(
1449
  {
1450
+ "position_ids": position_ids,
1451
  "past_key_values": past_key_values,
1452
  "use_cache": kwargs.get("use_cache"),
1453
+ "attention_mask": attention_mask,
1454
  }
1455
  )
1456
  return model_inputs
 
1468
  return reordered_past
1469
 
1470
 
1471
+ @add_start_docstrings(
1472
+ """
1473
+ The StableLm transformer with a sequence classification head on top (linear layer).
1474
+
1475
+ [`StableLmForSequenceClassification`] uses the last token in order to do the classification, as other causal
1476
+ models (e.g. GPT-2) do.
1477
+
1478
+ Since it does classification on the last token, it requires to know the position of the last token. If a
1479
+ `pad_token_id` is defined in the configuration, it finds the last token that is not a padding token in each row. If
1480
+ no `pad_token_id` is defined, it simply takes the last value in each row of the batch. Since it cannot guess the
1481
+ padding tokens when `inputs_embeds` are passed instead of `input_ids`, it does the same (take the last value in
1482
+ each row of the batch).
1483
+ """,
1484
+ STABLELM_START_DOCSTRING,
1485
+ )
1486
+ # Copied from transformers.models.llama.modeling_llama.LlamaForSequenceClassification with LLAMA->STABLELM,Llama->StableLm
1487
+ class StableLmForSequenceClassification(StableLmPreTrainedModel):
1488
+ def __init__(self, config):
1489
+ super().__init__(config)
1490
+ self.num_labels = config.num_labels
1491
+ self.model = StableLmModel(config)
1492
+ self.score = nn.Linear(config.hidden_size, self.num_labels, bias=False)
1493
+
1494
+ # Initialize weights and apply final processing
1495
+ self.post_init()
1496
+
1497
+ def get_input_embeddings(self):
1498
+ return self.model.embed_tokens
1499
+
1500
+ def set_input_embeddings(self, value):
1501
+ self.model.embed_tokens = value
1502
+
1503
+ @add_start_docstrings_to_model_forward(STABLELM_INPUTS_DOCSTRING)
1504
+ def forward(
1505
+ self,
1506
+ input_ids: torch.LongTensor = None,
1507
+ attention_mask: Optional[torch.Tensor] = None,
1508
+ position_ids: Optional[torch.LongTensor] = None,
1509
+ past_key_values: Optional[List[torch.FloatTensor]] = None,
1510
+ inputs_embeds: Optional[torch.FloatTensor] = None,
1511
+ labels: Optional[torch.LongTensor] = None,
1512
+ use_cache: Optional[bool] = None,
1513
+ output_attentions: Optional[bool] = None,
1514
+ output_hidden_states: Optional[bool] = None,
1515
+ return_dict: Optional[bool] = None,
1516
+ ) -> Union[Tuple, SequenceClassifierOutputWithPast]:
1517
+ r"""
1518
+ labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
1519
+ Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
1520
+ config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
1521
+ `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
1522
+ """
1523
+ return_dict = (
1524
+ return_dict if return_dict is not None else self.config.use_return_dict
1525
+ )
1526
+
1527
+ transformer_outputs = self.model(
1528
+ input_ids,
1529
+ attention_mask=attention_mask,
1530
+ position_ids=position_ids,
1531
+ past_key_values=past_key_values,
1532
+ inputs_embeds=inputs_embeds,
1533
+ use_cache=use_cache,
1534
+ output_attentions=output_attentions,
1535
+ output_hidden_states=output_hidden_states,
1536
+ return_dict=return_dict,
1537
+ )
1538
+ hidden_states = transformer_outputs[0]
1539
+ logits = self.score(hidden_states)
1540
+
1541
+ if input_ids is not None:
1542
+ batch_size = input_ids.shape[0]
1543
+ else:
1544
+ batch_size = inputs_embeds.shape[0]
1545
+
1546
+ if self.config.pad_token_id is None and batch_size != 1:
1547
+ raise ValueError(
1548
+ "Cannot handle batch sizes > 1 if no padding token is defined."
1549
+ )
1550
+ if self.config.pad_token_id is None:
1551
+ sequence_lengths = -1
1552
+ else:
1553
+ if input_ids is not None:
1554
+ # if no pad token found, use modulo instead of reverse indexing for ONNX compatibility
1555
+ sequence_lengths = (
1556
+ torch.eq(input_ids, self.config.pad_token_id).int().argmax(-1) - 1
1557
+ )
1558
+ sequence_lengths = sequence_lengths % input_ids.shape[-1]
1559
+ sequence_lengths = sequence_lengths.to(logits.device)
1560
+ else:
1561
+ sequence_lengths = -1
1562
+
1563
+ pooled_logits = logits[
1564
+ torch.arange(batch_size, device=logits.device), sequence_lengths
1565
+ ]
1566
+
1567
+ loss = None
1568
+ if labels is not None:
1569
+ labels = labels.to(logits.device)
1570
+ if self.config.problem_type is None:
1571
+ if self.num_labels == 1:
1572
+ self.config.problem_type = "regression"
1573
+ elif self.num_labels > 1 and (
1574
+ labels.dtype == torch.long or labels.dtype == torch.int
1575
+ ):
1576
+ self.config.problem_type = "single_label_classification"
1577
+ else:
1578
+ self.config.problem_type = "multi_label_classification"
1579
+
1580
+ if self.config.problem_type == "regression":
1581
+ loss_fct = MSELoss()
1582
+ if self.num_labels == 1:
1583
+ loss = loss_fct(pooled_logits.squeeze(), labels.squeeze())
1584
+ else:
1585
+ loss = loss_fct(pooled_logits, labels)
1586
+ elif self.config.problem_type == "single_label_classification":
1587
+ loss_fct = CrossEntropyLoss()
1588
+ loss = loss_fct(
1589
+ pooled_logits.view(-1, self.num_labels), labels.view(-1)
1590
+ )
1591
+ elif self.config.problem_type == "multi_label_classification":
1592
+ loss_fct = BCEWithLogitsLoss()
1593
+ loss = loss_fct(pooled_logits, labels)
1594
+ if not return_dict:
1595
+ output = (pooled_logits,) + transformer_outputs[1:]
1596
+ return ((loss,) + output) if loss is not None else output
1597
+
1598
+ return SequenceClassifierOutputWithPast(
1599
+ loss=loss,
1600
+ logits=pooled_logits,
1601
+ past_key_values=transformer_outputs.past_key_values,
1602
+ hidden_states=transformer_outputs.hidden_states,
1603
+ attentions=transformer_outputs.attentions,
1604
+ )
1605
+
1606
+
1607
+ StableLmConfig.register_for_auto_class()
1608
+ StableLmForCausalLM.register_for_auto_class("AutoModelForCausalLM")
special_tokens_map.json CHANGED
@@ -1,5 +1,65 @@
1
  {
2
- "bos_token": "<|endoftext|>",
3
- "eos_token": "<|endoftext|>",
4
- "pad_token": "<|endoftext|>"
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
5
  }
 
1
  {
2
+ "additional_special_tokens": [
3
+ "<|reg_extra|>",
4
+ "<|endoftext|>",
5
+ "<|fim_prefix|>",
6
+ "<|fim_middle|>",
7
+ "<|fim_suffix|>",
8
+ "<|fim_pad|>",
9
+ "<gh_stars>",
10
+ "<filename>",
11
+ "<issue_start>",
12
+ "<issue_comment>",
13
+ "<issue_closed>",
14
+ "<jupyter_start>",
15
+ "<jupyter_text>",
16
+ "<jupyter_code>",
17
+ "<jupyter_output>",
18
+ "<empty_output>",
19
+ "<commit_before>",
20
+ "<commit_msg>",
21
+ "<commit_after>",
22
+ "<reponame>",
23
+ "<|endofprompt|>",
24
+ "<|im_start|>",
25
+ "<|im_end|>",
26
+ "<|pause|>",
27
+ "<|reg0|>",
28
+ "<|reg1|>",
29
+ "<|reg2|>",
30
+ "<|reg3|>",
31
+ "<|reg4|>",
32
+ "<|reg5|>",
33
+ "<|reg6|>",
34
+ "<|reg7|>",
35
+ "<|extra0|>"
36
+ ],
37
+ "bos_token": {
38
+ "content": "<|endoftext|>",
39
+ "lstrip": false,
40
+ "normalized": false,
41
+ "rstrip": false,
42
+ "single_word": false
43
+ },
44
+ "eos_token": {
45
+ "content": "<|endoftext|>",
46
+ "lstrip": false,
47
+ "normalized": false,
48
+ "rstrip": false,
49
+ "single_word": false
50
+ },
51
+ "pad_token": {
52
+ "content": "<|endoftext|>",
53
+ "lstrip": false,
54
+ "normalized": false,
55
+ "rstrip": false,
56
+ "single_word": false
57
+ },
58
+ "unk_token": {
59
+ "content": "<|endoftext|>",
60
+ "lstrip": false,
61
+ "normalized": false,
62
+ "rstrip": false,
63
+ "single_word": false
64
+ }
65
  }
tokenization_arcade100k.py DELETED
@@ -1,292 +0,0 @@
1
- # coding=utf-8
2
- # Copyright (c) 2023 Alibaba Cloud & Stability AI.
3
- #
4
- # Tongyi Qianwen LICENSE AGREEMENT:
5
- # https://github.com/QwenLM/Qwen/blob/5aa84bdfd3237b37f01bc88cd49b3279b9a71d0b/Tongyi%20Qianwen%20LICENSE%20AGREEMENT
6
- """Tokenization classes for Arcade100k."""
7
-
8
- import base64
9
- import os
10
- import unicodedata
11
- from typing import Collection, Dict, List, Set, Tuple, Union
12
-
13
- import tiktoken
14
- from transformers.utils import logging
15
- from transformers import PreTrainedTokenizer, AddedToken
16
-
17
- logger = logging.get_logger(__name__)
18
-
19
- VOCAB_FILES_NAMES = {"vocab_file": "arcade100k.tiktoken"}
20
- NAME = "arcade100k"
21
-
22
-
23
- def _load_tiktoken_bpe(tiktoken_bpe_file: str) -> Dict[bytes, int]:
24
- with open(tiktoken_bpe_file, "rb") as f:
25
- contents = f.read()
26
- return {
27
- base64.b64decode(token): int(rank)
28
- for token, rank in (line.split() for line in contents.splitlines() if line)
29
- }
30
-
31
-
32
- ENDOFTEXT = "<|endoftext|>"
33
- FIM = [
34
- "<|fim_prefix|>",
35
- "<|fim_middle|>",
36
- "<|fim_suffix|>",
37
- "<|fim_pad|>",
38
- ]
39
- # `StarCoder` Tokens
40
- CODE = [
41
- "<gh_stars>",
42
- "<filename>",
43
- "<issue_start>",
44
- "<issue_comment>",
45
- "<issue_closed>",
46
- "<jupyter_start>",
47
- "<jupyter_text>",
48
- "<jupyter_code>",
49
- "<jupyter_output>",
50
- "<empty_output>",
51
- "<commit_before>",
52
- "<commit_msg>",
53
- "<commit_after>",
54
- "<reponame>",
55
- ]
56
- CHAT = [
57
- "<|im_start|>", # Chat: Input message start
58
- "<|im_end|>", # Chat: Input message end
59
- ]
60
- PAUSE = "<|pause|>" # Think before you speak (https://arxiv.org/abs/2310.02226)
61
- REGISTERS = [
62
- f"<|reg{i}|>" for i in range(0, 8)
63
- ] # Register 0 sink token (https://arxiv.org/abs/2309.17453)
64
- ENDOFPROMPT = "<|endofprompt|>"
65
- SPECIAL_TOKENS_NAMES = (
66
- [ENDOFTEXT]
67
- + FIM
68
- + CODE
69
- + [ENDOFPROMPT]
70
- + CHAT
71
- + [PAUSE]
72
- + REGISTERS
73
- + ["<|extra0|>"]
74
- )
75
- START_ID = 100257
76
- SPECIAL_TOKENS = {t: START_ID + i for i, t in enumerate(SPECIAL_TOKENS_NAMES)}
77
-
78
-
79
- def _arcade100k(vocab_file: str):
80
- mergeable_ranks = _load_tiktoken_bpe(vocab_file)
81
-
82
- return {
83
- "name": NAME,
84
- "pat_str": r"""(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+""",
85
- "mergeable_ranks": mergeable_ranks,
86
- "special_tokens": SPECIAL_TOKENS,
87
- }
88
-
89
-
90
- class Arcade100kTokenizer(PreTrainedTokenizer):
91
- """
92
- Construct a Arcade100k tokenizer backed by `tiktoken`.
93
-
94
- Args:
95
- vocab_file (`str`):
96
- Path to the vocabulary file.
97
- errors (`str`, *optional*, defaults to `"replace"`):
98
- How to handle errors in decoding UTF-8 byte sequences.
99
- WARNING: the default behaviour of this function is lossy, since decoded bytes are not
100
- guaranteed to be valid UTF-8. You can control this behaviour using the `errors` parameter,
101
- for instance, setting `errors=strict`.
102
- """
103
-
104
- vocab_files_names = VOCAB_FILES_NAMES
105
- model_input_names = ["input_ids", "attention_mask"]
106
-
107
- def __init__(
108
- self,
109
- vocab_file: str,
110
- errors: str = "replace",
111
- **kwargs,
112
- ):
113
- super().__init__(errors=errors, **kwargs)
114
- self.errors = errors
115
-
116
- self._tiktoken_config = _arcade100k(vocab_file)
117
- self.tokenizer = tiktoken.Encoding(**self._tiktoken_config)
118
-
119
- # TODO: Remove this assertion
120
- assert (
121
- len(self.tokenizer._mergeable_ranks)
122
- + len(self.tokenizer._special_tokens)
123
- + 1
124
- == self.tokenizer.n_vocab
125
- ), f"{len(self.tokenizer._mergeable_ranks) + len(self.tokenizer._special_tokens)} != {self.tokenizer.n_vocab} in encoding"
126
-
127
- self.decoder = {i: n for n, i in self.tokenizer._mergeable_ranks.items()}
128
- self.decoder.update({i: n for n, i in self.tokenizer._special_tokens.items()})
129
- # Provide default `eos_token` and `pad_token`
130
- if self.eos_token is None:
131
- self.eos_token = self.decoder[self.tokenizer.eot_token]
132
- if self.pad_token is None:
133
- self.pad_token = self.decoder[self.tokenizer.pad_token]
134
-
135
- # Expose for convenience
136
- self.mergeable_ranks = self.tokenizer._mergeable_ranks
137
- self.special_tokens = self.tokenizer._special_tokens
138
-
139
- def __len__(self):
140
- return self.tokenizer.n_vocab
141
-
142
- def __getstate__(self):
143
- # Required for `pickle` support
144
- state = self.__dict__.copy()
145
- del state["tokenizer"]
146
- return state
147
-
148
- def __setstate__(self, state):
149
- self.__dict__.update(state)
150
- self.tokenizer = tiktoken.Encoding(**self._tiktoken_config)
151
-
152
- @property
153
- def vocab_size(self):
154
- return self.tokenizer.n_vocab
155
-
156
- def get_vocab(self) -> Dict[bytes, int]:
157
- return self.tokenizer._mergeable_ranks
158
-
159
- def convert_tokens_to_ids(
160
- self, tokens: Union[bytes, str, List[Union[bytes, str]]]
161
- ) -> List[int]:
162
- ids = []
163
- if isinstance(tokens, (str, bytes)):
164
- if tokens in self.tokenizer._special_tokens:
165
- return self.tokenizer._special_tokens[tokens]
166
- else:
167
- return self.tokenizer._mergeable_ranks.get(tokens)
168
- for token in tokens:
169
- if token in self.tokenizer._special_tokens:
170
- ids.append(self.tokenizer._special_tokens[token])
171
- else:
172
- ids.append(self.tokenizer._mergeable_ranks.get(token))
173
- return ids
174
-
175
- def _add_tokens(
176
- self,
177
- new_tokens: Union[List[str], List[AddedToken]],
178
- special_tokens: bool = False,
179
- ) -> int:
180
- if not special_tokens and new_tokens:
181
- raise ValueError("Adding regular tokens is not supported")
182
- for token in new_tokens:
183
- surface_form = token.content if isinstance(token, AddedToken) else token
184
- if surface_form not in SPECIAL_TOKENS:
185
- raise ValueError("Adding unknown special tokens is not supported")
186
- return 0
187
-
188
- def save_vocabulary(self, save_directory: str, **kwargs) -> Tuple[str]:
189
- """
190
- Save only the vocabulary of the tokenizer (vocabulary).
191
-
192
- Returns:
193
- `Tuple(str)`: Paths to the files saved.
194
- """
195
- file_path = os.path.join(save_directory, "arcade100k.tiktoken")
196
- with open(file_path, "w", encoding="utf8") as w:
197
- for k, v in self.tokenizer._mergeable_ranks.items():
198
- line = base64.b64encode(k).decode("utf8") + " " + str(v) + "\n"
199
- w.write(line)
200
- return (file_path,)
201
-
202
- def tokenize(
203
- self,
204
- text: str,
205
- allowed_special: Union[Set, str] = "all",
206
- disallowed_special: Union[Collection, str] = (),
207
- **kwargs,
208
- ) -> List[Union[bytes, str]]:
209
- """
210
- Converts a string in a sequence of tokens.
211
-
212
- Args:
213
- text (`str`):
214
- The sequence to be encoded.
215
- allowed_special (`Literal["all"]` or `set`):
216
- The surface forms of the tokens to be encoded as special tokens in regular texts.
217
- Default to "all".
218
- disallowed_special (`Literal["all"]` or `Collection`):
219
- The surface forms of the tokens that should not be in regular texts and trigger errors.
220
- Default to an empty tuple.
221
-
222
- kwargs (additional keyword arguments, *optional*):
223
- Will be passed to the underlying model specific encode method.
224
-
225
- Returns:
226
- `List[bytes|str]`: The list of tokens.
227
- """
228
- tokens = []
229
- text = unicodedata.normalize("NFC", text)
230
-
231
- # this implementation takes a detour: text -> token id -> token surface forms
232
- for t in self.tokenizer.encode(
233
- text, allowed_special=allowed_special, disallowed_special=disallowed_special
234
- ):
235
- tokens.append(self.decoder[t])
236
- return tokens
237
-
238
- def convert_tokens_to_string(self, tokens: List[Union[bytes, str]]) -> str:
239
- """
240
- Converts a sequence of tokens in a single string.
241
- """
242
- text = ""
243
- temp = b""
244
- for t in tokens:
245
- if isinstance(t, str):
246
- if temp:
247
- text += temp.decode("utf-8", errors=self.errors)
248
- temp = b""
249
- text += t
250
- elif isinstance(t, bytes):
251
- temp += t
252
- else:
253
- raise TypeError("token should only be of type types or str")
254
- if temp:
255
- text += temp.decode("utf-8", errors=self.errors)
256
- return text
257
-
258
- def _convert_id_to_token(self, index: int) -> Union[bytes, str]:
259
- """Converts an id to a token, special tokens included"""
260
- if index in self.decoder:
261
- return self.decoder[index]
262
- raise ValueError("unknown ids")
263
-
264
- def _convert_token_to_id(self, token: Union[bytes, str]) -> int:
265
- """Converts a token to an id using the vocab, special tokens included"""
266
- if token in self.tokenizer._special_tokens:
267
- return self.tokenizer._special_tokens[token]
268
- if token in self.tokenizer._mergeable_ranks:
269
- return self.tokenizer._mergeable_ranks[token]
270
- raise ValueError("unknown token")
271
-
272
- def _tokenize(self, text: str, **kwargs):
273
- """
274
- Converts a string in a sequence of tokens (string), using the tokenizer. Split in words for word-based
275
- vocabulary or sub-words for sub-word-based vocabularies (BPE/SentencePieces/WordPieces).
276
-
277
- Do NOT take care of added tokens.
278
- """
279
- raise NotImplementedError
280
-
281
- def _decode(
282
- self,
283
- token_ids: Union[int, List[int]],
284
- skip_special_tokens: bool = False,
285
- errors: str = None,
286
- **kwargs,
287
- ) -> str:
288
- if isinstance(token_ids, int):
289
- token_ids = [token_ids]
290
- if skip_special_tokens:
291
- token_ids = [i for i in token_ids if i < self.tokenizer.eot_token]
292
- return self.tokenizer.decode(token_ids)
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
tokenizer.json ADDED
The diff for this file is too large to render. See raw diff
 
tokenizer_config.json CHANGED
@@ -1,17 +1,316 @@
1
  {
2
- "added_tokens_decoder": {},
3
- "auto_map": {
4
- "AutoTokenizer": [
5
- "tokenization_arcade100k.Arcade100kTokenizer",
6
- null
7
- ]
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
8
  },
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
9
  "bos_token": "<|endoftext|>",
10
  "chat_template": "{% if messages[0]['role'] == 'system' %}{% set loop_messages = messages[1:] %}{% set system_message = messages[0]['content'] %}{% else %}{% set loop_messages = messages %}{% set system_message = 'You are a helpful assistant.' %}{% endif %}{% if not add_generation_prompt is defined %}{% set add_generation_prompt = false %}{% endif %}{% for message in loop_messages %}{% if loop.index0 == 0 %}{{'<|im_start|>system\n' + system_message + '<|im_end|>\n'}}{% endif %}{{'<|im_start|>' + message['role'] + '\n' + message['content'] + '<|im_end|>' + '\n'}}{% endfor %}{% if add_generation_prompt %}{{ '<|im_start|>assistant\n' }}{% endif %}",
11
  "clean_up_tokenization_spaces": true,
12
  "eos_token": "<|endoftext|>",
13
- "errors": "replace",
14
  "model_max_length": 4096,
 
15
  "pad_token": "<|endoftext|>",
16
- "tokenizer_class": "Arcade100kTokenizer"
 
 
 
17
  }
 
1
  {
2
+ "add_prefix_space": false,
3
+ "added_tokens_decoder": {
4
+ "100256": {
5
+ "content": "<|reg_extra|>",
6
+ "lstrip": false,
7
+ "normalized": false,
8
+ "rstrip": false,
9
+ "single_word": false,
10
+ "special": true
11
+ },
12
+ "100257": {
13
+ "content": "<|endoftext|>",
14
+ "lstrip": false,
15
+ "normalized": false,
16
+ "rstrip": false,
17
+ "single_word": false,
18
+ "special": true
19
+ },
20
+ "100258": {
21
+ "content": "<|fim_prefix|>",
22
+ "lstrip": false,
23
+ "normalized": false,
24
+ "rstrip": false,
25
+ "single_word": false,
26
+ "special": true
27
+ },
28
+ "100259": {
29
+ "content": "<|fim_middle|>",
30
+ "lstrip": false,
31
+ "normalized": false,
32
+ "rstrip": false,
33
+ "single_word": false,
34
+ "special": true
35
+ },
36
+ "100260": {
37
+ "content": "<|fim_suffix|>",
38
+ "lstrip": false,
39
+ "normalized": false,
40
+ "rstrip": false,
41
+ "single_word": false,
42
+ "special": true
43
+ },
44
+ "100261": {
45
+ "content": "<|fim_pad|>",
46
+ "lstrip": false,
47
+ "normalized": false,
48
+ "rstrip": false,
49
+ "single_word": false,
50
+ "special": true
51
+ },
52
+ "100262": {
53
+ "content": "<gh_stars>",
54
+ "lstrip": false,
55
+ "normalized": false,
56
+ "rstrip": false,
57
+ "single_word": false,
58
+ "special": true
59
+ },
60
+ "100263": {
61
+ "content": "<filename>",
62
+ "lstrip": false,
63
+ "normalized": false,
64
+ "rstrip": false,
65
+ "single_word": false,
66
+ "special": true
67
+ },
68
+ "100264": {
69
+ "content": "<issue_start>",
70
+ "lstrip": false,
71
+ "normalized": false,
72
+ "rstrip": false,
73
+ "single_word": false,
74
+ "special": true
75
+ },
76
+ "100265": {
77
+ "content": "<issue_comment>",
78
+ "lstrip": false,
79
+ "normalized": false,
80
+ "rstrip": false,
81
+ "single_word": false,
82
+ "special": true
83
+ },
84
+ "100266": {
85
+ "content": "<issue_closed>",
86
+ "lstrip": false,
87
+ "normalized": false,
88
+ "rstrip": false,
89
+ "single_word": false,
90
+ "special": true
91
+ },
92
+ "100267": {
93
+ "content": "<jupyter_start>",
94
+ "lstrip": false,
95
+ "normalized": false,
96
+ "rstrip": false,
97
+ "single_word": false,
98
+ "special": true
99
+ },
100
+ "100268": {
101
+ "content": "<jupyter_text>",
102
+ "lstrip": false,
103
+ "normalized": false,
104
+ "rstrip": false,
105
+ "single_word": false,
106
+ "special": true
107
+ },
108
+ "100269": {
109
+ "content": "<jupyter_code>",
110
+ "lstrip": false,
111
+ "normalized": false,
112
+ "rstrip": false,
113
+ "single_word": false,
114
+ "special": true
115
+ },
116
+ "100270": {
117
+ "content": "<jupyter_output>",
118
+ "lstrip": false,
119
+ "normalized": false,
120
+ "rstrip": false,
121
+ "single_word": false,
122
+ "special": true
123
+ },
124
+ "100271": {
125
+ "content": "<empty_output>",
126
+ "lstrip": false,
127
+ "normalized": false,
128
+ "rstrip": false,
129
+ "single_word": false,
130
+ "special": true
131
+ },
132
+ "100272": {
133
+ "content": "<commit_before>",
134
+ "lstrip": false,
135
+ "normalized": false,
136
+ "rstrip": false,
137
+ "single_word": false,
138
+ "special": true
139
+ },
140
+ "100273": {
141
+ "content": "<commit_msg>",
142
+ "lstrip": false,
143
+ "normalized": false,
144
+ "rstrip": false,
145
+ "single_word": false,
146
+ "special": true
147
+ },
148
+ "100274": {
149
+ "content": "<commit_after>",
150
+ "lstrip": false,
151
+ "normalized": false,
152
+ "rstrip": false,
153
+ "single_word": false,
154
+ "special": true
155
+ },
156
+ "100275": {
157
+ "content": "<reponame>",
158
+ "lstrip": false,
159
+ "normalized": false,
160
+ "rstrip": false,
161
+ "single_word": false,
162
+ "special": true
163
+ },
164
+ "100276": {
165
+ "content": "<|endofprompt|>",
166
+ "lstrip": false,
167
+ "normalized": false,
168
+ "rstrip": false,
169
+ "single_word": false,
170
+ "special": true
171
+ },
172
+ "100277": {
173
+ "content": "<|im_start|>",
174
+ "lstrip": false,
175
+ "normalized": false,
176
+ "rstrip": false,
177
+ "single_word": false,
178
+ "special": true
179
+ },
180
+ "100278": {
181
+ "content": "<|im_end|>",
182
+ "lstrip": false,
183
+ "normalized": false,
184
+ "rstrip": false,
185
+ "single_word": false,
186
+ "special": true
187
+ },
188
+ "100279": {
189
+ "content": "<|pause|>",
190
+ "lstrip": false,
191
+ "normalized": false,
192
+ "rstrip": false,
193
+ "single_word": false,
194
+ "special": true
195
+ },
196
+ "100280": {
197
+ "content": "<|reg0|>",
198
+ "lstrip": false,
199
+ "normalized": false,
200
+ "rstrip": false,
201
+ "single_word": false,
202
+ "special": true
203
+ },
204
+ "100281": {
205
+ "content": "<|reg1|>",
206
+ "lstrip": false,
207
+ "normalized": false,
208
+ "rstrip": false,
209
+ "single_word": false,
210
+ "special": true
211
+ },
212
+ "100282": {
213
+ "content": "<|reg2|>",
214
+ "lstrip": false,
215
+ "normalized": false,
216
+ "rstrip": false,
217
+ "single_word": false,
218
+ "special": true
219
+ },
220
+ "100283": {
221
+ "content": "<|reg3|>",
222
+ "lstrip": false,
223
+ "normalized": false,
224
+ "rstrip": false,
225
+ "single_word": false,
226
+ "special": true
227
+ },
228
+ "100284": {
229
+ "content": "<|reg4|>",
230
+ "lstrip": false,
231
+ "normalized": false,
232
+ "rstrip": false,
233
+ "single_word": false,
234
+ "special": true
235
+ },
236
+ "100285": {
237
+ "content": "<|reg5|>",
238
+ "lstrip": false,
239
+ "normalized": false,
240
+ "rstrip": false,
241
+ "single_word": false,
242
+ "special": true
243
+ },
244
+ "100286": {
245
+ "content": "<|reg6|>",
246
+ "lstrip": false,
247
+ "normalized": false,
248
+ "rstrip": false,
249
+ "single_word": false,
250
+ "special": true
251
+ },
252
+ "100287": {
253
+ "content": "<|reg7|>",
254
+ "lstrip": false,
255
+ "normalized": false,
256
+ "rstrip": false,
257
+ "single_word": false,
258
+ "special": true
259
+ },
260
+ "100288": {
261
+ "content": "<|extra0|>",
262
+ "lstrip": false,
263
+ "normalized": false,
264
+ "rstrip": false,
265
+ "single_word": false,
266
+ "special": true
267
+ }
268
  },
269
+ "additional_special_tokens": [
270
+ "<|reg_extra|>",
271
+ "<|endoftext|>",
272
+ "<|fim_prefix|>",
273
+ "<|fim_middle|>",
274
+ "<|fim_suffix|>",
275
+ "<|fim_pad|>",
276
+ "<gh_stars>",
277
+ "<filename>",
278
+ "<issue_start>",
279
+ "<issue_comment>",
280
+ "<issue_closed>",
281
+ "<jupyter_start>",
282
+ "<jupyter_text>",
283
+ "<jupyter_code>",
284
+ "<jupyter_output>",
285
+ "<empty_output>",
286
+ "<commit_before>",
287
+ "<commit_msg>",
288
+ "<commit_after>",
289
+ "<reponame>",
290
+ "<|endofprompt|>",
291
+ "<|im_start|>",
292
+ "<|im_end|>",
293
+ "<|pause|>",
294
+ "<|reg0|>",
295
+ "<|reg1|>",
296
+ "<|reg2|>",
297
+ "<|reg3|>",
298
+ "<|reg4|>",
299
+ "<|reg5|>",
300
+ "<|reg6|>",
301
+ "<|reg7|>",
302
+ "<|extra0|>"
303
+ ],
304
  "bos_token": "<|endoftext|>",
305
  "chat_template": "{% if messages[0]['role'] == 'system' %}{% set loop_messages = messages[1:] %}{% set system_message = messages[0]['content'] %}{% else %}{% set loop_messages = messages %}{% set system_message = 'You are a helpful assistant.' %}{% endif %}{% if not add_generation_prompt is defined %}{% set add_generation_prompt = false %}{% endif %}{% for message in loop_messages %}{% if loop.index0 == 0 %}{{'<|im_start|>system\n' + system_message + '<|im_end|>\n'}}{% endif %}{{'<|im_start|>' + message['role'] + '\n' + message['content'] + '<|im_end|>' + '\n'}}{% endfor %}{% if add_generation_prompt %}{{ '<|im_start|>assistant\n' }}{% endif %}",
306
  "clean_up_tokenization_spaces": true,
307
  "eos_token": "<|endoftext|>",
308
+ "max_length": null,
309
  "model_max_length": 4096,
310
+ "pad_to_multiple_of": null,
311
  "pad_token": "<|endoftext|>",
312
+ "pad_token_type_id": 0,
313
+ "padding_side": "left",
314
+ "tokenizer_class": "GPT2Tokenizer",
315
+ "unk_token": "<|endoftext|>"
316
  }
vocab.json ADDED
The diff for this file is too large to render. See raw diff