radh25sh commited on
Commit
54bae31
1 Parent(s): a2a2efc

Upload folder using huggingface_hub

Browse files
config.json ADDED
@@ -0,0 +1,37 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "_name_or_path": "models/GPTNeoX-160m",
3
+ "architectures": [
4
+ "GPTNeoXForCausalLM"
5
+ ],
6
+ "attention_bias": true,
7
+ "attention_dropout": 0.0,
8
+ "auto_map": {
9
+ "AutoConfig": "config_custom.GPTNeoXConfig",
10
+ "AutoModel": "modeling_custom.GPTNeoXModel",
11
+ "AutoModelForCausalLM": "modeling_custom.GPTNeoXForCausalLM"
12
+ },
13
+ "bos_token_id": 0,
14
+ "classifier_dropout": 0.1,
15
+ "eos_token_id": 0,
16
+ "hidden_act": "gelu",
17
+ "hidden_dropout": 0.0,
18
+ "hidden_size": 768,
19
+ "initializer_range": 0.02,
20
+ "intermediate_size": 3072,
21
+ "layer_norm_eps": 1e-05,
22
+ "max_position_embeddings": 2048,
23
+ "model_type": "gpt_neox",
24
+ "num_attention_heads": 12,
25
+ "num_hidden_layers": 12,
26
+ "partial_rotary_factor": 0.25,
27
+ "rope_scaling": null,
28
+ "rope_theta": 10000,
29
+ "rotary_emb_base": 10000,
30
+ "rotary_pct": 0.25,
31
+ "tie_word_embeddings": false,
32
+ "torch_dtype": "bfloat16",
33
+ "transformers_version": "4.46.2",
34
+ "use_cache": true,
35
+ "use_parallel_residual": true,
36
+ "vocab_size": 50304
37
+ }
config_custom.py ADDED
@@ -0,0 +1,191 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # coding=utf-8
2
+ # Copyright 2022 EleutherAI and The HuggingFace Inc. team. All rights reserved.
3
+ #
4
+ # Licensed under the Apache License, Version 2.0 (the "License");
5
+ # you may not use this file except in compliance with the License.
6
+ # You may obtain a copy of the License at
7
+ #
8
+ # http://www.apache.org/licenses/LICENSE-2.0
9
+ #
10
+ # Unless required by applicable law or agreed to in writing, software
11
+ # distributed under the License is distributed on an "AS IS" BASIS,
12
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13
+ # See the License for the specific language governing permissions and
14
+ # limitations under the License.
15
+ """GPTNeoX model configuration"""
16
+
17
+ from transformers.configuration_utils import PretrainedConfig
18
+ from transformers.modeling_rope_utils import rope_config_validation
19
+ from transformers.utils import logging
20
+
21
+
22
+ logger = logging.get_logger(__name__)
23
+
24
+
25
+ class GPTNeoXConfig(PretrainedConfig):
26
+ r"""
27
+ This is the configuration class to store the configuration of a [`GPTNeoXModel`]. It is used to instantiate an
28
+ GPTNeoX model according to the specified arguments, defining the model architecture. Instantiating a configuration
29
+ with the defaults will yield a similar configuration to that of the GPTNeoX
30
+ [EleutherAI/gpt-neox-20b](https://huggingface.co/EleutherAI/gpt-neox-20b) architecture.
31
+
32
+ Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
33
+ documentation from [`PretrainedConfig`] for more information.
34
+
35
+
36
+ Args:
37
+ vocab_size (`int`, *optional*, defaults to 50432):
38
+ Vocabulary size of the GPTNeoX model. Defines the number of different tokens that can be represented by the
39
+ `inputs_ids` passed when calling [`GPTNeoXModel`].
40
+ hidden_size (`int`, *optional*, defaults to 6144):
41
+ Dimension of the encoder layers and the pooler layer.
42
+ num_hidden_layers (`int`, *optional*, defaults to 44):
43
+ Number of hidden layers in the Transformer encoder.
44
+ num_attention_heads (`int`, *optional*, defaults to 64):
45
+ Number of attention heads for each attention layer in the Transformer encoder.
46
+ intermediate_size (`int`, *optional*, defaults to 24576):
47
+ Dimension of the "intermediate" (i.e., feed-forward) layer in the Transformer encoder.
48
+ hidden_act (`str` or `function`, *optional*, defaults to `"gelu"`):
49
+ The non-linear activation function (function or string) in the encoder and pooler. If string, `"gelu"`,
50
+ `"relu"`, `"selu"` and `"gelu_new"` are supported.
51
+ rotary_pct (`float`, *optional*, defaults to 0.25):
52
+ percentage of hidden dimensions to allocate to rotary embeddings
53
+ rotary_emb_base (`int`, *optional*, defaults to 10000)
54
+ base for computing rotary embeddings frequency
55
+ attention_dropout (`float`, *optional*, defaults to 0.0):
56
+ The dropout ratio probability of the attention score.
57
+ hidden_dropout (`float`, *optional*, defaults to 0.0):
58
+ The dropout ratio of (1) the word embeddings, (2) the post-attention hidden states, and (3) the post-mlp
59
+ hidden states.
60
+ classifier_dropout (`float`, *optional*, defaults to 0.1):
61
+ Argument used when doing token classification, used in the model [`GPTNeoXForTokenClassification`].
62
+
63
+ The dropout ratio for the hidden layer.
64
+ max_position_embeddings (`int`, *optional*, defaults to 2048):
65
+ The maximum sequence length that this model might ever be used with. Typically set this to something large
66
+ just in case (e.g., 512 or 1024 or 2048).
67
+ initializer_range (`float`, *optional*, defaults to 1e-5):
68
+ The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
69
+ layer_norm_eps (`float`, *optional*, defaults to 1e-12):
70
+ The epsilon used by the layer normalization layers.
71
+ use_cache (`bool`, *optional*, defaults to `True`):
72
+ Whether or not the model should return the last key/values attentions (not used by all models). Only
73
+ relevant if `config.is_decoder=True`.
74
+ use_parallel_residual (`bool`, *optional*, defaults to `True`):
75
+ Whether to use a "parallel" formulation in each Transformer layer, which can provide a slight training
76
+ speedup at large scales (e.g. 20B).
77
+ rope_scaling (`Dict`, *optional*):
78
+ Dictionary containing the scaling configuration for the RoPE embeddings. NOTE: if you apply new rope type
79
+ and you expect the model to work on longer `max_position_embeddings`, we recommend you to update this value
80
+ accordingly.
81
+ Expected contents:
82
+ `rope_type` (`str`):
83
+ The sub-variant of RoPE to use. Can be one of ['default', 'linear', 'dynamic', 'yarn', 'longrope',
84
+ 'llama3'], with 'default' being the original RoPE implementation.
85
+ `factor` (`float`, *optional*):
86
+ Used with all rope types except 'default'. The scaling factor to apply to the RoPE embeddings. In
87
+ most scaling types, a `factor` of x will enable the model to handle sequences of length x *
88
+ original maximum pre-trained length.
89
+ `original_max_position_embeddings` (`int`, *optional*):
90
+ Used with 'dynamic', 'longrope' and 'llama3'. The original max position embeddings used during
91
+ pretraining.
92
+ `attention_factor` (`float`, *optional*):
93
+ Used with 'yarn' and 'longrope'. The scaling factor to be applied on the attention
94
+ computation. If unspecified, it defaults to value recommended by the implementation, using the
95
+ `factor` field to infer the suggested value.
96
+ `beta_fast` (`float`, *optional*):
97
+ Only used with 'yarn'. Parameter to set the boundary for extrapolation (only) in the linear
98
+ ramp function. If unspecified, it defaults to 32.
99
+ `beta_slow` (`float`, *optional*):
100
+ Only used with 'yarn'. Parameter to set the boundary for interpolation (only) in the linear
101
+ ramp function. If unspecified, it defaults to 1.
102
+ `short_factor` (`List[float]`, *optional*):
103
+ Only used with 'longrope'. The scaling factor to be applied to short contexts (<
104
+ `original_max_position_embeddings`). Must be a list of numbers with the same length as the hidden
105
+ size divided by the number of attention heads divided by 2
106
+ `long_factor` (`List[float]`, *optional*):
107
+ Only used with 'longrope'. The scaling factor to be applied to long contexts (<
108
+ `original_max_position_embeddings`). Must be a list of numbers with the same length as the hidden
109
+ size divided by the number of attention heads divided by 2
110
+ `low_freq_factor` (`float`, *optional*):
111
+ Only used with 'llama3'. Scaling factor applied to low frequency components of the RoPE
112
+ `high_freq_factor` (`float`, *optional*):
113
+ Only used with 'llama3'. Scaling factor applied to high frequency components of the RoPE
114
+ attention_bias (`bool`, *optional*, defaults to `True`):
115
+ Whether to use a bias in the query, key, value and output projection layers during self-attention.
116
+
117
+ Example:
118
+
119
+ ```python
120
+ >>> from transformers import GPTNeoXConfig, GPTNeoXModel
121
+
122
+ >>> # Initializing a GPTNeoX gpt-neox-20b style configuration
123
+ >>> configuration = GPTNeoXConfig()
124
+
125
+ >>> # Initializing a model (with random weights) from the gpt-neox-20b style configuration
126
+ >>> model = GPTNeoXModel(configuration) # doctest: +SKIP
127
+
128
+ >>> # Accessing the model configuration
129
+ >>> configuration = model.config # doctest: +SKIP
130
+ ```"""
131
+
132
+ model_type = "gpt_neox"
133
+ keys_to_ignore_at_inference = ["past_key_values"]
134
+
135
+ def __init__(
136
+ self,
137
+ vocab_size=50432,
138
+ hidden_size=6144,
139
+ num_hidden_layers=44,
140
+ num_attention_heads=64,
141
+ intermediate_size=24576,
142
+ hidden_act="gelu",
143
+ rotary_pct=0.25,
144
+ rotary_emb_base=10000,
145
+ attention_dropout=0.0,
146
+ hidden_dropout=0.0,
147
+ classifier_dropout=0.1,
148
+ max_position_embeddings=2048,
149
+ initializer_range=0.02,
150
+ layer_norm_eps=1e-5,
151
+ use_cache=True,
152
+ bos_token_id=0,
153
+ eos_token_id=2,
154
+ tie_word_embeddings=False,
155
+ use_parallel_residual=True,
156
+ rope_scaling=None,
157
+ attention_bias=True,
158
+ **kwargs,
159
+ ):
160
+ super().__init__(bos_token_id=bos_token_id, eos_token_id=eos_token_id, **kwargs)
161
+ self.vocab_size = vocab_size
162
+ self.max_position_embeddings = max_position_embeddings
163
+ self.hidden_size = hidden_size
164
+ self.num_hidden_layers = num_hidden_layers
165
+ self.num_attention_heads = num_attention_heads
166
+ self.intermediate_size = intermediate_size
167
+ self.hidden_act = hidden_act
168
+ self.rotary_pct = rotary_pct
169
+ self.partial_rotary_factor = rotary_pct
170
+ self.rotary_emb_base = rotary_emb_base
171
+ self.rope_theta = rotary_emb_base
172
+ self.attention_dropout = attention_dropout
173
+ self.hidden_dropout = hidden_dropout
174
+ self.classifier_dropout = classifier_dropout
175
+ self.initializer_range = initializer_range
176
+ self.layer_norm_eps = layer_norm_eps
177
+ self.use_cache = use_cache
178
+ self.tie_word_embeddings = tie_word_embeddings
179
+ self.use_parallel_residual = use_parallel_residual
180
+ self.rope_scaling = rope_scaling
181
+ self.attention_bias = attention_bias
182
+ # Validate the correctness of rotary position embeddings parameters
183
+ # BC: if there is a 'type' field, move it to 'rope_type'.
184
+ if self.rope_scaling is not None and "type" in self.rope_scaling:
185
+ self.rope_scaling["rope_type"] = self.rope_scaling["type"]
186
+ rope_config_validation(self)
187
+
188
+ if self.hidden_size % self.num_attention_heads != 0:
189
+ raise ValueError(
190
+ "The hidden size is not divisble by the number of attention heads! Make sure to update them!"
191
+ )
generation_config.json ADDED
@@ -0,0 +1,6 @@
 
 
 
 
 
 
 
1
+ {
2
+ "_from_model_config": true,
3
+ "bos_token_id": 0,
4
+ "eos_token_id": 0,
5
+ "transformers_version": "4.46.2"
6
+ }
model.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:460a8c0e6854140f1a81b53e7d0f949dd2b525c1d62484d1d1d84e38a560a918
3
+ size 324662984
modeling_custom.py ADDED
@@ -0,0 +1,1575 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # coding=utf-8
2
+ # Copyright 2022 EleutherAI The HuggingFace Inc. team. All rights reserved.
3
+ #
4
+ # Licensed under the Apache License, Version 2.0 (the "License");
5
+ # you may not use this file except in compliance with the License.
6
+ # You may obtain a copy of the License at
7
+ #
8
+ # http://www.apache.org/licenses/LICENSE-2.0
9
+ #
10
+ # Unless required by applicable law or agreed to in writing, software
11
+ # distributed under the License is distributed on an "AS IS" BASIS,
12
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13
+ # See the License for the specific language governing permissions and
14
+ # limitations under the License.
15
+ """PyTorch GPTNeoX model."""
16
+
17
+ from typing import Optional, Tuple, Union
18
+
19
+ import torch
20
+ import torch.utils.checkpoint
21
+ from packaging import version
22
+ from torch import nn
23
+ from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
24
+
25
+ from transformers.activations import ACT2FN
26
+ from transformers.cache_utils import Cache, DynamicCache, StaticCache
27
+ from transformers.file_utils import (
28
+ add_code_sample_docstrings,
29
+ add_start_docstrings,
30
+ add_start_docstrings_to_model_forward,
31
+ replace_return_docstrings,
32
+ )
33
+ from transformers.generation import GenerationMixin
34
+ from transformers.modeling_attn_mask_utils import AttentionMaskConverter
35
+ from transformers.modeling_outputs import (
36
+ BaseModelOutputWithPast,
37
+ CausalLMOutputWithPast,
38
+ QuestionAnsweringModelOutput,
39
+ SequenceClassifierOutputWithPast,
40
+ TokenClassifierOutput,
41
+ )
42
+ from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS
43
+ from transformers.modeling_utils import PreTrainedModel
44
+ from transformers.utils import (
45
+ get_torch_version,
46
+ is_flash_attn_2_available,
47
+ is_flash_attn_greater_or_equal_2_10,
48
+ logging,
49
+ )
50
+ from .config_custom import GPTNeoXConfig
51
+
52
+
53
+ if is_flash_attn_2_available():
54
+ from transformers.modeling_flash_attention_utils import _flash_attention_forward
55
+
56
+ logger = logging.get_logger(__name__)
57
+
58
+ _CHECKPOINT_FOR_DOC = "trl-internal-testing/tiny-random-GPTNeoXForCausalLM"
59
+ _REAL_CHECKPOINT_FOR_DOC = "EleutherAI/gpt-neox-20b"
60
+ _CONFIG_FOR_DOC = "GPTNeoXConfig"
61
+
62
+
63
+ # Copied from transformers.models.llama.modeling_llama._prepare_4d_causal_attention_mask_with_cache_position
64
+ def _prepare_4d_causal_attention_mask_with_cache_position(
65
+ attention_mask: torch.Tensor,
66
+ sequence_length: int,
67
+ target_length: int,
68
+ dtype: torch.dtype,
69
+ device: torch.device,
70
+ min_dtype: float,
71
+ cache_position: torch.Tensor,
72
+ batch_size: int,
73
+ ):
74
+ """
75
+ Creates a causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape
76
+ `(batch_size, key_value_length)`, or if the input `attention_mask` is already 4D, do nothing.
77
+
78
+ Args:
79
+ attention_mask (`torch.Tensor`):
80
+ A 2D attention mask of shape `(batch_size, key_value_length)` or a 4D attention mask of shape `(batch_size, 1, query_length, key_value_length)`.
81
+ sequence_length (`int`):
82
+ The sequence length being processed.
83
+ target_length (`int`):
84
+ The target length: when generating with static cache, the mask should be as long as the static cache, to account for the 0 padding, the part of the cache that is not filled yet.
85
+ dtype (`torch.dtype`):
86
+ The dtype to use for the 4D attention mask.
87
+ device (`torch.device`):
88
+ The device to plcae the 4D attention mask on.
89
+ min_dtype (`float`):
90
+ The minimum value representable with the dtype `dtype`.
91
+ cache_position (`torch.Tensor`):
92
+ Indices depicting the position of the input sequence tokens in the sequence.
93
+ batch_size (`torch.Tensor`):
94
+ Batch size.
95
+ """
96
+ if attention_mask is not None and attention_mask.dim() == 4:
97
+ # In this case we assume that the mask comes already in inverted form and requires no inversion or slicing.
98
+ causal_mask = attention_mask
99
+ else:
100
+ causal_mask = torch.full((sequence_length, target_length), fill_value=min_dtype, dtype=dtype, device=device)
101
+ if sequence_length != 1:
102
+ causal_mask = torch.triu(causal_mask, diagonal=1)
103
+ causal_mask *= torch.arange(target_length, device=device) > cache_position.reshape(-1, 1)
104
+ causal_mask = causal_mask[None, None, :, :].expand(batch_size, 1, -1, -1)
105
+ if attention_mask is not None:
106
+ causal_mask = causal_mask.clone() # copy to contiguous memory for in-place edit
107
+ mask_length = attention_mask.shape[-1]
108
+ padding_mask = causal_mask[:, :, :, :mask_length] + attention_mask[:, None, None, :]
109
+ padding_mask = padding_mask == 0
110
+ causal_mask[:, :, :, :mask_length] = causal_mask[:, :, :, :mask_length].masked_fill(
111
+ padding_mask, min_dtype
112
+ )
113
+
114
+ return causal_mask
115
+
116
+
117
+ class GPTNeoXPreTrainedModel(PreTrainedModel):
118
+ """
119
+ An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
120
+ models.
121
+ """
122
+
123
+ config_class = GPTNeoXConfig
124
+ base_model_prefix = "gpt_neox"
125
+ supports_gradient_checkpointing = True
126
+ _no_split_modules = ["GPTNeoXLayer"]
127
+ _skip_keys_device_placement = "past_key_values"
128
+ _supports_flash_attn_2 = True
129
+ _supports_cache_class = True
130
+ _supports_quantized_cache = True
131
+ _supports_static_cache = True
132
+ _supports_sdpa = True
133
+
134
+ def _init_weights(self, module):
135
+ """Initialize the weights"""
136
+ if isinstance(module, nn.Linear):
137
+ module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
138
+ if module.bias is not None:
139
+ module.bias.data.zero_()
140
+ elif isinstance(module, nn.Embedding):
141
+ module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
142
+ if module.padding_idx is not None:
143
+ module.weight.data[module.padding_idx].zero_()
144
+ elif isinstance(module, nn.LayerNorm):
145
+ module.bias.data.zero_()
146
+ module.weight.data.fill_(1.0)
147
+
148
+
149
+ class GPTNeoXAttention(nn.Module):
150
+ def __init__(self, config, layer_idx=None):
151
+ super().__init__()
152
+ self.config = config
153
+ self.num_attention_heads = config.num_attention_heads
154
+ self.hidden_size = config.hidden_size
155
+ if self.hidden_size % self.num_attention_heads != 0:
156
+ raise ValueError(
157
+ "The hidden size is not divisble by the number of attention heads! Make sure to update them"
158
+ )
159
+ self.head_size = self.hidden_size // self.num_attention_heads
160
+ self.rotary_ndims = int(self.head_size * config.rotary_pct)
161
+ self.rope_theta = config.rotary_emb_base
162
+ self._init_bias(config.max_position_embeddings)
163
+
164
+ self.register_buffer("masked_bias", torch.tensor(-1e9), persistent=False)
165
+ self.rotary_emb = GPTNeoXRotaryEmbedding(config=self.config)
166
+
167
+ if layer_idx is None:
168
+ logger.warning_once(
169
+ f"Instantiating {self.__class__.__name__} without passing a `layer_idx` is not recommended and will "
170
+ "lead to errors during the forward call if caching is used. Please make sure to provide a `layer_idx` "
171
+ "when creating this class."
172
+ )
173
+ self.norm_factor = self.head_size**-0.5
174
+ self.query_key_value = nn.Linear(config.hidden_size, 3 * config.hidden_size, bias=config.attention_bias)
175
+ self.dense = nn.Linear(config.hidden_size, config.hidden_size, bias=config.attention_bias)
176
+ self.attention_dropout = nn.Dropout(config.attention_dropout)
177
+ self.is_causal = True
178
+ self.layer_idx = layer_idx
179
+
180
+ def _init_bias(self, max_positions, device=None):
181
+ self.register_buffer(
182
+ "bias",
183
+ torch.tril(torch.ones((max_positions, max_positions), dtype=torch.bool)).view(
184
+ 1, 1, max_positions, max_positions
185
+ ),
186
+ persistent=False,
187
+ )
188
+ if device is not None:
189
+ self.bias = self.bias.to(device)
190
+
191
+ def forward(
192
+ self,
193
+ hidden_states: torch.FloatTensor,
194
+ attention_mask: torch.FloatTensor,
195
+ position_ids: torch.LongTensor,
196
+ head_mask: Optional[torch.FloatTensor] = None,
197
+ layer_past: Optional[Cache] = None,
198
+ use_cache: Optional[bool] = False,
199
+ output_attentions: Optional[bool] = False,
200
+ padding_mask: Optional[torch.Tensor] = None,
201
+ cache_position: Optional[torch.LongTensor] = None,
202
+ position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, # will become mandatory in v4.46
203
+ ):
204
+ # Apply attention-specific projections and rope
205
+ query, key, value, present = self._attn_projections_and_rope(
206
+ hidden_states=hidden_states,
207
+ position_ids=position_ids,
208
+ layer_past=layer_past,
209
+ use_cache=use_cache,
210
+ position_embeddings=position_embeddings,
211
+ )
212
+
213
+ # Compute attention
214
+ attn_output, attn_weights = self._attn(query, key, value, attention_mask, head_mask)
215
+
216
+ # Reshape outputs
217
+ attn_output = self._merge_heads(attn_output, self.num_attention_heads, self.head_size)
218
+ attn_output = self.dense(attn_output)
219
+
220
+ outputs = (attn_output, present)
221
+ if output_attentions:
222
+ outputs += (attn_weights,)
223
+
224
+ return outputs
225
+
226
+ @classmethod
227
+ def _split_heads(cls, tensor, num_attention_heads, attn_head_size):
228
+ """
229
+ Splits hidden dim into attn_head_size and num_attention_heads
230
+ """
231
+ # tensor: [bs, seq_len, hidden_size]
232
+ new_shape = tensor.size()[:-1] + (num_attention_heads, attn_head_size)
233
+ # -> [bs, seq_len, num_attention_heads, attn_head_size]
234
+ tensor = tensor.view(new_shape)
235
+ # -> [bs, num_attention_heads, seq_len, attn_head_size]
236
+ tensor = tensor.permute(0, 2, 1, 3)
237
+ return tensor
238
+
239
+ @classmethod
240
+ def _merge_heads(cls, tensor, num_attention_heads, attn_head_size):
241
+ """
242
+ Merges attn_head_size dim and num_attn_heads dim into hidden dim
243
+ """
244
+ # tensor [bs, num_attention_heads, seq_len, attn_head_size]
245
+ tensor = tensor.permute(0, 2, 1, 3).contiguous()
246
+ # -> [bs, seq_len, num_attention_heads, attn_head_size]
247
+ tensor = tensor.view(tensor.size(0), tensor.size(1), num_attention_heads * attn_head_size)
248
+ # -> [bs, seq_len, hidden_size]
249
+ return tensor
250
+
251
+ def _attn_projections_and_rope(
252
+ self,
253
+ hidden_states: torch.FloatTensor,
254
+ position_ids: torch.LongTensor,
255
+ layer_past: Optional[Tuple[torch.Tensor]] = None,
256
+ use_cache: Optional[bool] = False,
257
+ cache_position: Optional[torch.LongTensor] = None,
258
+ position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, # will become mandatory in v4.46
259
+ ):
260
+ # Compute QKV
261
+ # Attention heads [batch, seq_len, hidden_size]
262
+ # --> [batch, seq_len, (np * 3 * head_size)]
263
+ qkv = self.query_key_value(hidden_states)
264
+
265
+ # [batch, seq_len, (num_heads * 3 * head_size)]
266
+ # --> [batch, seq_len, num_heads, 3 * head_size]
267
+ new_qkv_shape = qkv.size()[:-1] + (self.num_attention_heads, 3 * self.head_size)
268
+ qkv = qkv.view(*new_qkv_shape)
269
+
270
+ # [batch, seq_len, num_attention_heads, 3 * head_size] --> 3 [batch, num_attention_heads, seq_len, head_size]
271
+ query = qkv[..., : self.head_size].permute(0, 2, 1, 3)
272
+ key = qkv[..., self.head_size : 2 * self.head_size].permute(0, 2, 1, 3)
273
+ value = qkv[..., 2 * self.head_size :].permute(0, 2, 1, 3)
274
+
275
+ # Compute rotary embeddings on rotary_ndims
276
+ query_rot = query[..., : self.rotary_ndims]
277
+ query_pass = query[..., self.rotary_ndims :]
278
+ key_rot = key[..., : self.rotary_ndims]
279
+ key_pass = key[..., self.rotary_ndims :]
280
+
281
+ if position_embeddings is None:
282
+ logger.warning_once(
283
+ "The attention layers in this model are transitioning from computing the RoPE embeddings internally "
284
+ "through `position_ids` (2D tensor with the indexes of the tokens), to using externally computed "
285
+ "`position_embeddings` (Tuple of tensors, containing cos and sin). In v4.46 `position_ids` will be "
286
+ "removed and `position_embeddings` will be mandatory."
287
+ )
288
+ cos, sin = self.rotary_emb(value, position_ids)
289
+ else:
290
+ cos, sin = position_embeddings
291
+ query, key = apply_rotary_pos_emb(query_rot, key_rot, cos, sin)
292
+ query = torch.cat((query, query_pass), dim=-1)
293
+ key = torch.cat((key, key_pass), dim=-1)
294
+
295
+ # Cache QKV values
296
+ if layer_past is not None:
297
+ cache_kwargs = {
298
+ "sin": sin,
299
+ "cos": cos,
300
+ "partial_rotation_size": self.rotary_ndims,
301
+ "cache_position": cache_position,
302
+ }
303
+ key, value = layer_past.update(key, value, self.layer_idx, cache_kwargs)
304
+
305
+ return query, key, value, layer_past
306
+
307
+ def _attn(self, query, key, value, attention_mask=None, head_mask=None):
308
+ # q, k, v: [bs, num_attention_heads, seq_len, attn_head_size]
309
+ # compute causal mask from causal mask buffer
310
+ batch_size, num_attention_heads, query_length, attn_head_size = query.size()
311
+ key_length = key.size(-2)
312
+
313
+ # dynamically increase the causal mask with the key length, if needed.
314
+ if key_length > self.bias.shape[-1]:
315
+ self._init_bias(key_length, device=key.device)
316
+ causal_mask = self.bias[:, :, key_length - query_length : key_length, :key_length]
317
+
318
+ query = query.view(batch_size * num_attention_heads, query_length, attn_head_size)
319
+ key = key.view(batch_size * num_attention_heads, key_length, attn_head_size)
320
+ attn_scores = torch.zeros(
321
+ batch_size * num_attention_heads,
322
+ query_length,
323
+ key_length,
324
+ dtype=query.dtype,
325
+ device=key.device,
326
+ )
327
+ attn_scores = torch.baddbmm(
328
+ attn_scores,
329
+ query,
330
+ key.transpose(1, 2),
331
+ beta=1.0,
332
+ alpha=self.norm_factor,
333
+ )
334
+ attn_scores = attn_scores.view(batch_size, num_attention_heads, query_length, key_length)
335
+
336
+ mask_value = torch.finfo(attn_scores.dtype).min
337
+ # Need to be a tensor, otherwise we get error: `RuntimeError: expected scalar type float but found double`.
338
+ # Need to be on the same device, otherwise `RuntimeError: ..., x and y to be on the same device`
339
+ mask_value = torch.tensor(mask_value, dtype=attn_scores.dtype).to(attn_scores.device)
340
+ attn_scores = torch.where(causal_mask, attn_scores, mask_value)
341
+
342
+ if attention_mask is not None: # no matter the length, we just slice it
343
+ causal_mask = attention_mask[:, :, :, : key.shape[-2]]
344
+ attn_scores = attn_scores + causal_mask
345
+
346
+ attn_weights = nn.functional.softmax(attn_scores, dim=-1)
347
+ attn_weights = attn_weights.to(value.dtype)
348
+
349
+ # Mask heads if we want to
350
+ if head_mask is not None:
351
+ attn_weights = attn_weights * head_mask
352
+
353
+ attn_weights = self.attention_dropout(attn_weights)
354
+
355
+ attn_output = torch.matmul(attn_weights, value)
356
+ return attn_output, attn_weights
357
+
358
+
359
+ class GPTNeoXFlashAttention2(GPTNeoXAttention):
360
+ """
361
+ GPTNeoX flash attention module. This module inherits from `GPTNeoXAttention` as the weights of the module stays
362
+ untouched. The only required change would be on the forward pass where it needs to correctly call the public API of
363
+ flash attention and deal with padding tokens in case the input contains any of them.
364
+ """
365
+
366
+ # Copied from transformers.models.llama.modeling_llama.LlamaFlashAttention2.__init__
367
+ def __init__(self, *args, **kwargs):
368
+ super().__init__(*args, **kwargs)
369
+
370
+ # TODO: Should be removed once Flash Attention for RoCm is bumped to 2.1.
371
+ # 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.
372
+ # 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).
373
+ self._flash_attn_uses_top_left_mask = not is_flash_attn_greater_or_equal_2_10()
374
+
375
+ def forward(
376
+ self,
377
+ hidden_states: torch.FloatTensor,
378
+ attention_mask: torch.FloatTensor,
379
+ position_ids: torch.LongTensor,
380
+ head_mask: Optional[torch.FloatTensor] = None,
381
+ layer_past: Optional[Cache] = None,
382
+ use_cache: Optional[bool] = False,
383
+ output_attentions: Optional[bool] = False,
384
+ cache_position: Optional[torch.LongTensor] = None,
385
+ position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, # will become mandatory in v4.46
386
+ ):
387
+ # Apply attention-specific projections and rope
388
+ query, key, value, present = self._attn_projections_and_rope(
389
+ hidden_states=hidden_states,
390
+ position_ids=position_ids,
391
+ layer_past=layer_past,
392
+ use_cache=use_cache,
393
+ cache_position=cache_position,
394
+ position_embeddings=position_embeddings,
395
+ )
396
+
397
+ query_length = query.shape[-2]
398
+
399
+ # GPT-neo-X casts query and key in fp32 to apply rotary embedding in full precision
400
+ target_dtype = value.dtype
401
+ if query.dtype != target_dtype:
402
+ query = query.to(target_dtype)
403
+ if key.dtype != target_dtype:
404
+ key = key.to(target_dtype)
405
+
406
+ #TODO: Permute to get the expected shape for Flash Attention
407
+ query = query.permute(0, 2, 1, 3).contiguous() # (batch_size, num_heads, seq_len, head_dim)
408
+ key = key.permute(0, 2, 1, 3).contiguous()
409
+ value = value.permute(0, 2, 1, 3).contiguous()
410
+ # query =
411
+ # key = ...
412
+ # value = ...
413
+
414
+ attention_dropout = self.config.attention_dropout if self.training else 0.0
415
+
416
+ #TODO: Compute attention with _flash_attention_forward
417
+ attn_weights = _flash_attention_forward(
418
+ query, key, value,
419
+ attention_mask,
420
+ query_length,
421
+ dropout=attention_dropout,
422
+ use_top_left_mask=self._flash_attn_uses_top_left_mask,
423
+ is_causal=self.is_causal
424
+ )
425
+ # attn_weights = ...
426
+
427
+ #TODO: Reshape outputs before projection
428
+ # attn_output = ...
429
+ attn_output = attn_weights.view(hidden_states.size(0), query_length, -1)
430
+
431
+ attn_output = self.dense(attn_output)
432
+
433
+ outputs = (attn_output, layer_past)
434
+ if output_attentions:
435
+ outputs += (attn_weights,)
436
+
437
+ return outputs
438
+
439
+
440
+ class GPTNeoXSdpaAttention(GPTNeoXAttention):
441
+ """
442
+ GPTNeoX attention module using torch.nn.functional.scaled_dot_product_attention. This module inherits from
443
+ `GPTNeoXAttention` as the weights of the module stays untouched. The only changes are on the forward pass
444
+ to adapt to the SDPA API.
445
+ """
446
+
447
+ def __init__(self, config, layer_idx=None):
448
+ super().__init__(config, layer_idx=layer_idx)
449
+
450
+ # SDPA with memory-efficient backend is broken in torch==2.1.2 when using non-contiguous inputs and a custom
451
+ # attn_mask, so we need to call `.contiguous()`. This was fixed in torch==2.2.0.
452
+ # Reference: https://github.com/pytorch/pytorch/issues/112577
453
+ self.require_contiguous_qkv = version.parse(get_torch_version()) < version.parse("2.2.0")
454
+
455
+ def forward(
456
+ self,
457
+ hidden_states: torch.FloatTensor,
458
+ attention_mask: torch.FloatTensor,
459
+ position_ids: torch.LongTensor,
460
+ head_mask: Optional[torch.FloatTensor] = None,
461
+ layer_past: Optional[Tuple[torch.Tensor]] = None,
462
+ use_cache: Optional[bool] = False,
463
+ output_attentions: Optional[bool] = False,
464
+ cache_position: Optional[torch.LongTensor] = None,
465
+ position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, # will become mandatory in v4.46
466
+ ):
467
+ if output_attentions or head_mask is not None:
468
+ logger.warning_once(
469
+ "`GPTNeoXSdpaAttention` is used but `torch.nn.functional.scaled_dot_product_attention` does not support "
470
+ "`output_attentions=True` or `head_mask`. Falling back to the manual attention implementation, but "
471
+ "specifying the manual implementation will be required from Transformers version v5.0.0 onwards. "
472
+ 'This warning can be removed using the argument `attn_implementation="eager"` when loading the model.'
473
+ )
474
+ return super().forward(
475
+ hidden_states=hidden_states,
476
+ attention_mask=attention_mask,
477
+ position_ids=position_ids,
478
+ head_mask=head_mask,
479
+ layer_past=layer_past,
480
+ use_cache=use_cache,
481
+ output_attentions=output_attentions,
482
+ cache_position=cache_position,
483
+ )
484
+
485
+ bsz, q_len, _ = hidden_states.size()
486
+
487
+ # Apply attention-specific projections and rope
488
+ query, key, value, present = self._attn_projections_and_rope(
489
+ hidden_states=hidden_states,
490
+ position_ids=position_ids,
491
+ layer_past=layer_past,
492
+ use_cache=use_cache,
493
+ cache_position=cache_position,
494
+ position_embeddings=position_embeddings,
495
+ )
496
+
497
+ causal_mask = attention_mask
498
+ if attention_mask is not None:
499
+ causal_mask = causal_mask[:, :, :, : key.shape[-2]]
500
+
501
+ # GPT-neo-X casts query and key in fp32 to apply rotary embedding in full precision
502
+ target_dtype = value.dtype
503
+ if query.dtype != target_dtype:
504
+ query = query.to(target_dtype)
505
+ if key.dtype != target_dtype:
506
+ key = key.to(target_dtype)
507
+
508
+ # Avoid torch==2.1.2 specific bug for the memory-efficient backend in SDPA
509
+ if self.require_contiguous_qkv and query.device.type == "cuda" and attention_mask is not None:
510
+ query = query.contiguous()
511
+ key = key.contiguous()
512
+ value = value.contiguous()
513
+
514
+ # We dispatch to SDPA's Flash Attention or Efficient kernels via this `is_causal` if statement instead of an inline conditional assignment
515
+ # in SDPA to support both torch.compile's dynamic shapes and full graph options. An inline conditional prevents dynamic shapes from compiling.
516
+ is_causal = True if causal_mask is None and q_len > 1 else False
517
+
518
+ attn_output = torch.nn.functional.scaled_dot_product_attention(
519
+ query=query,
520
+ key=key,
521
+ value=value,
522
+ attn_mask=causal_mask,
523
+ dropout_p=self.attention_dropout.p if self.training else 0.0,
524
+ is_causal=is_causal,
525
+ )
526
+
527
+ # Reshape outputs
528
+ attn_output = attn_output.transpose(1, 2).contiguous()
529
+ attn_output = attn_output.view(bsz, q_len, self.hidden_size)
530
+
531
+ attn_output = self.dense(attn_output)
532
+
533
+ return attn_output, present, None
534
+
535
+
536
+ def attention_mask_func(attention_scores, ltor_mask):
537
+ attention_scores.masked_fill_(~ltor_mask, torch.finfo(attention_scores.dtype).min)
538
+ return attention_scores
539
+
540
+
541
+ # Copied from transformers.models.llama.modeling_llama.LlamaRotaryEmbedding with Llama->GPTNeoX
542
+ class GPTNeoXRotaryEmbedding(nn.Module):
543
+ def __init__(
544
+ self,
545
+ dim=None,
546
+ max_position_embeddings=2048,
547
+ base=10000,
548
+ device=None,
549
+ scaling_factor=1.0,
550
+ rope_type="default",
551
+ config: Optional[GPTNeoXConfig] = None,
552
+ ):
553
+ super().__init__()
554
+ # TODO (joao): remove the `if` below, only used for BC
555
+ self.rope_kwargs = {}
556
+ if config is None:
557
+ logger.warning_once(
558
+ "`GPTNeoXRotaryEmbedding` can now be fully parameterized by passing the model config through the "
559
+ "`config` argument. All other arguments will be removed in v4.46"
560
+ )
561
+ self.rope_kwargs = {
562
+ "rope_type": rope_type,
563
+ "factor": scaling_factor,
564
+ "dim": dim,
565
+ "base": base,
566
+ "max_position_embeddings": max_position_embeddings,
567
+ }
568
+ self.rope_type = rope_type
569
+ self.max_seq_len_cached = max_position_embeddings
570
+ self.original_max_seq_len = max_position_embeddings
571
+ else:
572
+ # BC: "rope_type" was originally "type"
573
+ if config.rope_scaling is not None:
574
+ self.rope_type = config.rope_scaling.get("rope_type", config.rope_scaling.get("type"))
575
+ else:
576
+ self.rope_type = "default"
577
+ self.max_seq_len_cached = config.max_position_embeddings
578
+ self.original_max_seq_len = config.max_position_embeddings
579
+
580
+ self.config = config
581
+ self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
582
+
583
+ inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device, **self.rope_kwargs)
584
+ self.register_buffer("inv_freq", inv_freq, persistent=False)
585
+ self.original_inv_freq = self.inv_freq
586
+
587
+ def _dynamic_frequency_update(self, position_ids, device):
588
+ """
589
+ dynamic RoPE layers should recompute `inv_freq` in the following situations:
590
+ 1 - growing beyond the cached sequence length (allow scaling)
591
+ 2 - the current sequence length is in the original scale (avoid losing precision with small sequences)
592
+ """
593
+ seq_len = torch.max(position_ids) + 1
594
+ if seq_len > self.max_seq_len_cached: # growth
595
+ inv_freq, self.attention_scaling = self.rope_init_fn(
596
+ self.config, device, seq_len=seq_len, **self.rope_kwargs
597
+ )
598
+ self.register_buffer("inv_freq", inv_freq, persistent=False) # TODO joao: may break with compilation
599
+ self.max_seq_len_cached = seq_len
600
+
601
+ if seq_len < self.original_max_seq_len and self.max_seq_len_cached > self.original_max_seq_len: # reset
602
+ self.register_buffer("inv_freq", self.original_inv_freq, persistent=False)
603
+ self.max_seq_len_cached = self.original_max_seq_len
604
+
605
+ @torch.no_grad()
606
+ def forward(self, x, position_ids):
607
+ if "dynamic" in self.rope_type:
608
+ self._dynamic_frequency_update(position_ids, device=x.device)
609
+
610
+ # Core RoPE block
611
+ inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1)
612
+ position_ids_expanded = position_ids[:, None, :].float()
613
+ # Force float32 (see https://github.com/huggingface/transformers/pull/29285)
614
+ device_type = x.device.type
615
+ device_type = device_type if isinstance(device_type, str) and device_type != "mps" else "cpu"
616
+ with torch.autocast(device_type=device_type, enabled=False):
617
+ freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
618
+ emb = torch.cat((freqs, freqs), dim=-1)
619
+ cos = emb.cos()
620
+ sin = emb.sin()
621
+
622
+ # Advanced RoPE types (e.g. yarn) apply a post-processing scaling factor, equivalent to scaling attention
623
+ cos = cos * self.attention_scaling
624
+ sin = sin * self.attention_scaling
625
+
626
+ return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
627
+
628
+
629
+ # Copied from transformers.models.llama.modeling_llama.LlamaLinearScalingRotaryEmbedding with Llama->GPTNeoX
630
+ class GPTNeoXLinearScalingRotaryEmbedding(GPTNeoXRotaryEmbedding):
631
+ """GPTNeoXRotaryEmbedding extended with linear scaling. Credits to the Reddit user /u/kaiokendev"""
632
+
633
+ def __init__(self, *args, **kwargs):
634
+ logger.warning_once(
635
+ "`GPTNeoXLinearScalingRotaryEmbedding` is deprecated an will be removed in v4.46. Please use "
636
+ "`GPTNeoXRotaryEmbedding`, which now also does linear scaling (simply pass the model config to __init__)."
637
+ )
638
+ kwargs["rope_type"] = "linear"
639
+ super().__init__(*args, **kwargs)
640
+
641
+
642
+ # Copied from transformers.models.llama.modeling_llama.LlamaDynamicNTKScalingRotaryEmbedding with Llama->GPTNeoX
643
+ class GPTNeoXDynamicNTKScalingRotaryEmbedding(GPTNeoXRotaryEmbedding):
644
+ """GPTNeoXRotaryEmbedding extended with Dynamic NTK scaling. Credits to the Reddit users /u/bloc97 and /u/emozilla"""
645
+
646
+ def __init__(self, *args, **kwargs):
647
+ logger.warning_once(
648
+ "`GPTNeoXDynamicNTKScalingRotaryEmbedding` is deprecated an will be removed in v4.46. Please use "
649
+ "`GPTNeoXRotaryEmbedding`, which now also does dynamic ntk scaling (simply pass the model config to "
650
+ "__init__)."
651
+ )
652
+ kwargs["rope_type"] = "dynamic"
653
+ super().__init__(*args, **kwargs)
654
+
655
+
656
+ def rotate_half(x):
657
+ """Rotates half the hidden dims of the input."""
658
+ x1 = x[..., : x.shape[-1] // 2]
659
+ x2 = x[..., x.shape[-1] // 2 :]
660
+ return torch.cat((-x2, x1), dim=-1)
661
+
662
+
663
+ # Copied from transformers.models.llama.modeling_llama.apply_rotary_pos_emb
664
+ def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
665
+ """Applies Rotary Position Embedding to the query and key tensors.
666
+
667
+ Args:
668
+ q (`torch.Tensor`): The query tensor.
669
+ k (`torch.Tensor`): The key tensor.
670
+ cos (`torch.Tensor`): The cosine part of the rotary embedding.
671
+ sin (`torch.Tensor`): The sine part of the rotary embedding.
672
+ position_ids (`torch.Tensor`, *optional*):
673
+ Deprecated and unused.
674
+ unsqueeze_dim (`int`, *optional*, defaults to 1):
675
+ The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
676
+ sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
677
+ that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
678
+ k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
679
+ cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
680
+ the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
681
+ Returns:
682
+ `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
683
+ """
684
+ cos = cos.unsqueeze(unsqueeze_dim)
685
+ sin = sin.unsqueeze(unsqueeze_dim)
686
+ q_embed = (q * cos) + (rotate_half(q) * sin)
687
+ k_embed = (k * cos) + (rotate_half(k) * sin)
688
+ return q_embed, k_embed
689
+
690
+
691
+ class GPTNeoXMLP(nn.Module):
692
+ def __init__(self, config):
693
+ super().__init__()
694
+ self.dense_h_to_4h = nn.Linear(config.hidden_size, config.intermediate_size)
695
+ self.dense_4h_to_h = nn.Linear(config.intermediate_size, config.hidden_size)
696
+ self.act = ACT2FN[config.hidden_act]
697
+
698
+ def forward(self, hidden_states):
699
+ hidden_states = self.dense_h_to_4h(hidden_states)
700
+ hidden_states = self.act(hidden_states)
701
+ hidden_states = self.dense_4h_to_h(hidden_states)
702
+ return hidden_states
703
+
704
+
705
+ GPT_NEOX_ATTENTION_CLASSES = {
706
+ "eager": GPTNeoXAttention,
707
+ "flash_attention_2": GPTNeoXFlashAttention2,
708
+ "sdpa": GPTNeoXSdpaAttention,
709
+ }
710
+
711
+
712
+ class GPTNeoXLayer(nn.Module):
713
+ def __init__(self, config, layer_idx):
714
+ super().__init__()
715
+ self.use_parallel_residual = config.use_parallel_residual
716
+ self.input_layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
717
+ self.post_attention_layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
718
+ self.post_attention_dropout = nn.Dropout(config.hidden_dropout)
719
+ self.post_mlp_dropout = nn.Dropout(config.hidden_dropout)
720
+ self.attention = GPT_NEOX_ATTENTION_CLASSES[config._attn_implementation](config, layer_idx)
721
+ self.mlp = GPTNeoXMLP(config)
722
+
723
+ def forward(
724
+ self,
725
+ hidden_states: Optional[torch.FloatTensor],
726
+ attention_mask: Optional[torch.FloatTensor] = None,
727
+ position_ids: Optional[torch.LongTensor] = None,
728
+ head_mask: Optional[torch.FloatTensor] = None,
729
+ use_cache: Optional[bool] = False,
730
+ layer_past: Optional[Cache] = None,
731
+ output_attentions: Optional[bool] = False,
732
+ cache_position: Optional[torch.LongTensor] = None,
733
+ position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, # will become mandatory in v4.46
734
+ ):
735
+ attention_layer_outputs = self.attention(
736
+ self.input_layernorm(hidden_states),
737
+ attention_mask=attention_mask,
738
+ position_ids=position_ids,
739
+ layer_past=layer_past,
740
+ head_mask=head_mask,
741
+ use_cache=use_cache,
742
+ output_attentions=output_attentions,
743
+ cache_position=cache_position,
744
+ position_embeddings=position_embeddings,
745
+ )
746
+ attn_output = attention_layer_outputs[0] # output_attn: attn_output, present, (attn_weights)
747
+ attn_output = self.post_attention_dropout(attn_output)
748
+ outputs = attention_layer_outputs[1:]
749
+
750
+ if self.use_parallel_residual:
751
+ # pseudocode:
752
+ # x = x + attn(ln1(x)) + mlp(ln2(x))
753
+ mlp_output = self.mlp(self.post_attention_layernorm(hidden_states))
754
+ mlp_output = self.post_mlp_dropout(mlp_output)
755
+ hidden_states = mlp_output + attn_output + hidden_states
756
+ else:
757
+ # pseudocode:
758
+ # x = x + attn(ln1(x))
759
+ # x = x + mlp(ln2(x))
760
+ attn_output = attn_output + hidden_states
761
+ mlp_output = self.mlp(self.post_attention_layernorm(attn_output))
762
+ mlp_output = self.post_mlp_dropout(mlp_output)
763
+ hidden_states = mlp_output + attn_output
764
+
765
+ if use_cache:
766
+ outputs = (hidden_states,) + outputs # hidden_states, present, (attn_weights)
767
+ else:
768
+ outputs = (hidden_states,) + outputs[1:] # hidden_states, (attn_weights)
769
+
770
+ return outputs
771
+
772
+
773
+ GPT_NEOX_START_DOCSTRING = r"""
774
+ This model is a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) sub-class. Use
775
+ it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage and
776
+ behavior.
777
+
778
+ Parameters:
779
+ config ([`~GPTNeoXConfig`]): Model configuration class with all the parameters of the model.
780
+ Initializing with a config file does not load the weights associated with the model, only the
781
+ configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
782
+ """
783
+
784
+ GPT_NEOX_INPUTS_DOCSTRING = r"""
785
+ Args:
786
+ input_ids (`torch.LongTensor` of shape `({0})`):
787
+ Indices of input sequence tokens in the vocabulary.
788
+
789
+ Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
790
+ [`PreTrainedTokenizer.__call__`] for details.
791
+
792
+ [What are input IDs?](../glossary#input-ids)
793
+ attention_mask (`torch.FloatTensor` of shape `({0})`, *optional*):
794
+ Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
795
+
796
+ - 1 for tokens that are **not masked**,
797
+ - 0 for tokens that are **masked**.
798
+
799
+ [What are attention masks?](../glossary#attention-mask)
800
+ position_ids (`torch.LongTensor` of shape `({0})`, *optional*):
801
+ Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,
802
+ config.n_positions - 1]`.
803
+
804
+ [What are position IDs?](../glossary#position-ids)
805
+ head_mask (`torch.FloatTensor` of shape `(num_heads,)` or `(num_layers, num_heads)`, *optional*):
806
+ Mask to nullify selected heads of the self-attention modules. Mask values selected in `[0, 1]`:
807
+
808
+ - 1 indicates the head is **not masked**,
809
+ - 0 indicates the head is **masked**.
810
+
811
+ inputs_embeds (`torch.FloatTensor` of shape `({0}, hidden_size)`, *optional*):
812
+ Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This
813
+ is useful if you want more control over how to convert *input_ids* indices into associated vectors than the
814
+ model's internal embedding lookup matrix.
815
+ past_key_values (`Cache` or `tuple(tuple(torch.FloatTensor))`, *optional*):
816
+ Pre-computed hidden-states (key and values in the self-attention blocks and in the cross-attention
817
+ blocks) that can be used to speed up sequential decoding. This typically consists in the `past_key_values`
818
+ returned by the model at a previous stage of decoding, when `use_cache=True` or `config.use_cache=True`.
819
+
820
+ Two formats are allowed:
821
+ - a [`~cache_utils.Cache`] instance, see our
822
+ [kv cache guide](https://huggingface.co/docs/transformers/en/kv_cache);
823
+ - Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of
824
+ shape `(batch_size, num_heads, sequence_length, embed_size_per_head)`). This is also known as the legacy
825
+ cache format.
826
+
827
+ The model will output the same cache format that is fed as input. If no `past_key_values` are passed, the
828
+ legacy cache format will be returned.
829
+
830
+ If `past_key_values` are used, the user can optionally input only the last `input_ids` (those that don't
831
+ have their past key value states given to this model) of shape `(batch_size, 1)` instead of all `input_ids`
832
+ of shape `(batch_size, sequence_length)`.
833
+ output_attentions (`bool`, *optional*):
834
+ Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
835
+ tensors for more detail.
836
+ output_hidden_states (`bool`, *optional*):
837
+ Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
838
+ more detail.
839
+ return_dict (`bool`, *optional*):
840
+ Whether or not to return a [`~file_utils.ModelOutput`] instead of a plain tuple.
841
+ cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*):
842
+ Indices depicting the position of the input sequence tokens in the sequence. Contrarily to `position_ids`,
843
+ this tensor is not affected by padding. It is used to update the cache in the correct position and to infer
844
+ the complete sequence length.
845
+ """
846
+
847
+
848
+ @add_start_docstrings(
849
+ "The bare GPTNeoX Model transformer outputting raw hidden-states without any specific head on top.",
850
+ GPT_NEOX_START_DOCSTRING,
851
+ )
852
+ class GPTNeoXModel(GPTNeoXPreTrainedModel):
853
+ def __init__(self, config):
854
+ super().__init__(config)
855
+ self.config = config
856
+
857
+ self.embed_in = nn.Embedding(config.vocab_size, config.hidden_size)
858
+ self.emb_dropout = nn.Dropout(config.hidden_dropout)
859
+ self.layers = nn.ModuleList([GPTNeoXLayer(config, i) for i in range(config.num_hidden_layers)])
860
+ self.final_layer_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
861
+ self.rotary_emb = GPTNeoXRotaryEmbedding(config=config)
862
+
863
+ self._attn_implementation = config._attn_implementation
864
+
865
+ self.gradient_checkpointing = False
866
+
867
+ # Initialize weights and apply final processing
868
+ self.post_init()
869
+
870
+ def get_input_embeddings(self):
871
+ return self.embed_in
872
+
873
+ def set_input_embeddings(self, value):
874
+ self.embed_in = value
875
+
876
+ @add_start_docstrings_to_model_forward(GPT_NEOX_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
877
+ @add_code_sample_docstrings(
878
+ checkpoint=_CHECKPOINT_FOR_DOC,
879
+ real_checkpoint=_REAL_CHECKPOINT_FOR_DOC,
880
+ output_type=BaseModelOutputWithPast,
881
+ config_class=_CONFIG_FOR_DOC,
882
+ )
883
+ def forward(
884
+ self,
885
+ input_ids: Optional[torch.LongTensor] = None,
886
+ attention_mask: Optional[torch.FloatTensor] = None,
887
+ position_ids: Optional[torch.LongTensor] = None,
888
+ head_mask: Optional[torch.FloatTensor] = None,
889
+ inputs_embeds: Optional[torch.FloatTensor] = None,
890
+ past_key_values: Optional[Union[Cache, Tuple[Tuple[torch.FloatTensor]]]] = None,
891
+ use_cache: Optional[bool] = None,
892
+ output_attentions: Optional[bool] = None,
893
+ output_hidden_states: Optional[bool] = None,
894
+ return_dict: Optional[bool] = None,
895
+ cache_position: Optional[torch.LongTensor] = None,
896
+ ) -> Union[Tuple, BaseModelOutputWithPast]:
897
+ r"""
898
+ use_cache (`bool`, *optional*):
899
+ If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see
900
+ `past_key_values`).
901
+ """
902
+ output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
903
+ output_hidden_states = (
904
+ output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
905
+ )
906
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
907
+ use_cache = use_cache if use_cache is not None else self.config.use_cache
908
+
909
+ if (input_ids is None) ^ (inputs_embeds is not None):
910
+ raise ValueError(
911
+ "You cannot specify both input_ids and inputs_embeds at the same time, and must specify either one"
912
+ )
913
+
914
+ if self.gradient_checkpointing and self.training:
915
+ if use_cache:
916
+ logger.warning_once(
917
+ "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
918
+ )
919
+ use_cache = False
920
+
921
+ if inputs_embeds is None:
922
+ inputs_embeds = self.embed_in(input_ids)
923
+
924
+ # kept for BC (non `Cache` `past_key_values` inputs)
925
+ return_legacy_cache = False
926
+ if use_cache and not isinstance(past_key_values, Cache):
927
+ return_legacy_cache = True
928
+ if past_key_values is None:
929
+ past_key_values = DynamicCache()
930
+ else:
931
+ past_key_values = DynamicCache.from_legacy_cache(past_key_values)
932
+ logger.warning_once(
933
+ "We detected that you are passing `past_key_values` as a tuple of tuples. This is deprecated and "
934
+ "will be removed in v4.47. Please convert your cache or use an appropriate `Cache` class "
935
+ "(https://huggingface.co/docs/transformers/kv_cache#legacy-cache-format)"
936
+ )
937
+
938
+ seq_length = inputs_embeds.shape[1]
939
+ if cache_position is None:
940
+ past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
941
+ cache_position = torch.arange(past_seen_tokens, past_seen_tokens + seq_length, device=inputs_embeds.device)
942
+
943
+ if position_ids is None:
944
+ position_ids = cache_position.unsqueeze(0)
945
+
946
+ causal_mask = self._update_causal_mask(
947
+ attention_mask, inputs_embeds, cache_position, past_key_values, output_attentions
948
+ )
949
+
950
+ # Prepare head mask if needed
951
+ # 1.0 in head_mask indicate we keep the head
952
+ # attention_probs has shape bsz x n_heads x N x N
953
+ # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads]
954
+ # and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length]
955
+ head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers)
956
+ hidden_states = self.emb_dropout(inputs_embeds)
957
+
958
+ # create position embeddings to be shared across the decoder layers
959
+ position_embeddings = self.rotary_emb(hidden_states, position_ids)
960
+
961
+ next_decoder_cache = None
962
+ all_attentions = () if output_attentions else None
963
+ all_hidden_states = () if output_hidden_states else None
964
+ for i, layer in enumerate(
965
+ self.layers,
966
+ ):
967
+ if output_hidden_states:
968
+ all_hidden_states = all_hidden_states + (hidden_states,)
969
+
970
+ if self.gradient_checkpointing and self.training:
971
+ outputs = self._gradient_checkpointing_func(
972
+ layer.__call__,
973
+ hidden_states,
974
+ causal_mask,
975
+ position_ids,
976
+ head_mask[i],
977
+ use_cache,
978
+ None,
979
+ output_attentions,
980
+ cache_position,
981
+ position_embeddings,
982
+ )
983
+ else:
984
+ outputs = layer(
985
+ hidden_states,
986
+ attention_mask=causal_mask,
987
+ position_ids=position_ids,
988
+ head_mask=head_mask[i],
989
+ layer_past=past_key_values,
990
+ use_cache=use_cache,
991
+ output_attentions=output_attentions,
992
+ cache_position=cache_position,
993
+ position_embeddings=position_embeddings,
994
+ )
995
+ hidden_states = outputs[0]
996
+ if use_cache is True:
997
+ next_decoder_cache = outputs[1]
998
+ if output_attentions:
999
+ all_attentions = all_attentions + (outputs[2 if use_cache else 1],)
1000
+
1001
+ hidden_states = self.final_layer_norm(hidden_states)
1002
+ # Add last hidden state
1003
+ if output_hidden_states:
1004
+ all_hidden_states = all_hidden_states + (hidden_states,)
1005
+
1006
+ next_cache = next_decoder_cache if use_cache else None
1007
+ if return_legacy_cache:
1008
+ next_cache = next_cache.to_legacy_cache()
1009
+
1010
+ if not return_dict:
1011
+ return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_attentions] if v is not None)
1012
+
1013
+ return BaseModelOutputWithPast(
1014
+ last_hidden_state=hidden_states,
1015
+ past_key_values=next_cache,
1016
+ hidden_states=all_hidden_states,
1017
+ attentions=all_attentions,
1018
+ )
1019
+
1020
+ # Copied from transformers.models.llama.modeling_llama.LlamaModel._update_causal_mask
1021
+ def _update_causal_mask(
1022
+ self,
1023
+ attention_mask: torch.Tensor,
1024
+ input_tensor: torch.Tensor,
1025
+ cache_position: torch.Tensor,
1026
+ past_key_values: Cache,
1027
+ output_attentions: bool,
1028
+ ):
1029
+ if self.config._attn_implementation == "flash_attention_2":
1030
+ if attention_mask is not None and 0.0 in attention_mask:
1031
+ return attention_mask
1032
+ return None
1033
+
1034
+ # For SDPA, when possible, we will rely on its `is_causal` argument instead of its `attn_mask` argument, in
1035
+ # order to dispatch on Flash Attention 2. This feature is not compatible with static cache, as SDPA will fail
1036
+ # to infer the attention mask.
1037
+ past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
1038
+ using_static_cache = isinstance(past_key_values, StaticCache)
1039
+
1040
+ # When output attentions is True, sdpa implementation's forward method calls the eager implementation's forward
1041
+ if self.config._attn_implementation == "sdpa" and not using_static_cache and not output_attentions:
1042
+ if AttentionMaskConverter._ignore_causal_mask_sdpa(
1043
+ attention_mask,
1044
+ inputs_embeds=input_tensor,
1045
+ past_key_values_length=past_seen_tokens,
1046
+ is_training=self.training,
1047
+ ):
1048
+ return None
1049
+
1050
+ dtype, device = input_tensor.dtype, input_tensor.device
1051
+ min_dtype = torch.finfo(dtype).min
1052
+ sequence_length = input_tensor.shape[1]
1053
+ if using_static_cache:
1054
+ target_length = past_key_values.get_max_length()
1055
+ else:
1056
+ target_length = (
1057
+ attention_mask.shape[-1]
1058
+ if isinstance(attention_mask, torch.Tensor)
1059
+ else past_seen_tokens + sequence_length + 1
1060
+ )
1061
+
1062
+ # In case the provided `attention` mask is 2D, we generate a causal mask here (4D).
1063
+ causal_mask = _prepare_4d_causal_attention_mask_with_cache_position(
1064
+ attention_mask,
1065
+ sequence_length=sequence_length,
1066
+ target_length=target_length,
1067
+ dtype=dtype,
1068
+ device=device,
1069
+ min_dtype=min_dtype,
1070
+ cache_position=cache_position,
1071
+ batch_size=input_tensor.shape[0],
1072
+ )
1073
+
1074
+ if (
1075
+ self.config._attn_implementation == "sdpa"
1076
+ and attention_mask is not None
1077
+ and attention_mask.device.type == "cuda"
1078
+ and not output_attentions
1079
+ ):
1080
+ # Attend to all tokens in fully masked rows in the causal_mask, for example the relevant first rows when
1081
+ # using left padding. This is required by F.scaled_dot_product_attention memory-efficient attention path.
1082
+ # Details: https://github.com/pytorch/pytorch/issues/110213
1083
+ causal_mask = AttentionMaskConverter._unmask_unattended(causal_mask, min_dtype)
1084
+
1085
+ return causal_mask
1086
+
1087
+
1088
+ @add_start_docstrings(
1089
+ """GPTNeoX Model with a `language modeling` head on top for CLM fine-tuning.""", GPT_NEOX_START_DOCSTRING
1090
+ )
1091
+ class GPTNeoXForCausalLM(GPTNeoXPreTrainedModel, GenerationMixin):
1092
+ _tied_weights_keys = ["embed_out.weight"]
1093
+
1094
+ def __init__(self, config):
1095
+ super().__init__(config)
1096
+
1097
+ self.gpt_neox = GPTNeoXModel(config)
1098
+ self.embed_out = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
1099
+
1100
+ # Initialize weights and apply final processing
1101
+ self.post_init()
1102
+
1103
+ def get_output_embeddings(self):
1104
+ return self.embed_out
1105
+
1106
+ def set_output_embeddings(self, new_embeddings):
1107
+ self.embed_out = new_embeddings
1108
+
1109
+ @add_start_docstrings_to_model_forward(GPT_NEOX_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
1110
+ @replace_return_docstrings(output_type=CausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC)
1111
+ def forward(
1112
+ self,
1113
+ input_ids: Optional[torch.LongTensor] = None,
1114
+ attention_mask: Optional[torch.FloatTensor] = None,
1115
+ position_ids: Optional[torch.LongTensor] = None,
1116
+ inputs_embeds: Optional[torch.FloatTensor] = None,
1117
+ head_mask: Optional[torch.FloatTensor] = None,
1118
+ past_key_values: Optional[Union[Cache, Tuple[Tuple[torch.FloatTensor]]]] = None,
1119
+ labels: Optional[torch.LongTensor] = None,
1120
+ use_cache: Optional[bool] = None,
1121
+ output_attentions: Optional[bool] = None,
1122
+ output_hidden_states: Optional[bool] = None,
1123
+ return_dict: Optional[bool] = None,
1124
+ cache_position: Optional[torch.LongTensor] = None,
1125
+ ) -> Union[Tuple, CausalLMOutputWithPast]:
1126
+ r"""
1127
+ labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
1128
+ Labels for computing the left-to-right language modeling loss (next word prediction). Indices should be in
1129
+ `[-100, 0, ..., config.vocab_size]` (see `input_ids` docstring) Tokens with indices set to `-100` are
1130
+ ignored (masked), the loss is only computed for the tokens with labels n `[0, ..., config.vocab_size]`.
1131
+ use_cache (`bool`, *optional*):
1132
+ If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see
1133
+ `past_key_values`).
1134
+
1135
+ Returns:
1136
+
1137
+ Example:
1138
+
1139
+ ```python
1140
+ >>> from transformers import AutoTokenizer, GPTNeoXForCausalLM, GPTNeoXConfig
1141
+ >>> import torch
1142
+
1143
+ >>> tokenizer = AutoTokenizer.from_pretrained("EleutherAI/gpt-neox-20b")
1144
+ >>> config = GPTNeoXConfig.from_pretrained("EleutherAI/gpt-neox-20b")
1145
+ >>> config.is_decoder = True
1146
+ >>> model = GPTNeoXForCausalLM.from_pretrained("EleutherAI/gpt-neox-20b", config=config)
1147
+
1148
+ >>> inputs = tokenizer("Hello, my dog is cute", return_tensors="pt")
1149
+ >>> outputs = model(**inputs)
1150
+
1151
+ >>> prediction_logits = outputs.logits
1152
+ ```"""
1153
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
1154
+
1155
+ outputs = self.gpt_neox(
1156
+ input_ids,
1157
+ attention_mask=attention_mask,
1158
+ position_ids=position_ids,
1159
+ head_mask=head_mask,
1160
+ inputs_embeds=inputs_embeds,
1161
+ past_key_values=past_key_values,
1162
+ use_cache=use_cache,
1163
+ output_attentions=output_attentions,
1164
+ output_hidden_states=output_hidden_states,
1165
+ return_dict=return_dict,
1166
+ cache_position=cache_position,
1167
+ )
1168
+
1169
+ hidden_states = outputs[0]
1170
+ lm_logits = self.embed_out(hidden_states)
1171
+
1172
+ lm_loss = None
1173
+ if labels is not None:
1174
+ # move labels to correct device to enable model parallelism
1175
+ labels = labels.to(lm_logits.device)
1176
+ # we are doing next-token prediction; shift prediction scores and input ids by one
1177
+ shift_logits = lm_logits[:, :-1, :].contiguous()
1178
+ labels = labels[:, 1:].contiguous()
1179
+ loss_fct = CrossEntropyLoss()
1180
+ lm_loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), labels.view(-1))
1181
+
1182
+ if not return_dict:
1183
+ output = (lm_logits,) + outputs[1:]
1184
+ return ((lm_loss,) + output) if lm_loss is not None else output
1185
+
1186
+ return CausalLMOutputWithPast(
1187
+ loss=lm_loss,
1188
+ logits=lm_logits,
1189
+ past_key_values=outputs.past_key_values,
1190
+ hidden_states=outputs.hidden_states,
1191
+ attentions=outputs.attentions,
1192
+ )
1193
+
1194
+ # can't be copied from llama, gpt-neox has embed_out and not lm_head
1195
+ def prepare_inputs_for_generation(
1196
+ self,
1197
+ input_ids,
1198
+ past_key_values=None,
1199
+ attention_mask=None,
1200
+ inputs_embeds=None,
1201
+ cache_position=None,
1202
+ position_ids=None,
1203
+ use_cache=True,
1204
+ **kwargs,
1205
+ ):
1206
+ # If we have cache: let's slice `input_ids` through `cache_position`, to keep only the unprocessed tokens
1207
+ # Exception 1: when passing input_embeds, input_ids may be missing entries
1208
+ # Exception 2: some generation methods do special slicing of input_ids, so we don't need to do it here
1209
+ if past_key_values is not None:
1210
+ if inputs_embeds is not None: # Exception 1
1211
+ input_ids = input_ids[:, -cache_position.shape[0] :]
1212
+ elif input_ids.shape[1] != cache_position.shape[0]: # Default case (the "else", a no op, is Exception 2)
1213
+ input_ids = input_ids[:, cache_position]
1214
+
1215
+ if attention_mask is not None and position_ids is None:
1216
+ # create position_ids on the fly for batch generation
1217
+ position_ids = attention_mask.long().cumsum(-1) - 1
1218
+ position_ids.masked_fill_(attention_mask == 0, 1)
1219
+ if past_key_values:
1220
+ position_ids = position_ids[:, -input_ids.shape[1] :]
1221
+
1222
+ # This `clone` call is needed to avoid recapturing cuda graphs with `torch.compile`'s `mode="reduce-overhead`, as otherwise the input `position_ids` would have various stride during the decoding. Here, simply using `.contiguous()` is not sufficient as in the batch size = 1 case, `position_ids` is already contiguous but with varying stride which retriggers a capture.
1223
+ position_ids = position_ids.clone(memory_format=torch.contiguous_format)
1224
+
1225
+ # if `inputs_embeds` are passed, we only want to use them in the 1st generation step
1226
+ if inputs_embeds is not None and cache_position[0] == 0:
1227
+ model_inputs = {"inputs_embeds": inputs_embeds, "input_ids": None}
1228
+ else:
1229
+ # The clone here is for the same reason as for `position_ids`.
1230
+ model_inputs = {"input_ids": input_ids.clone(memory_format=torch.contiguous_format), "inputs_embeds": None}
1231
+
1232
+ if isinstance(past_key_values, StaticCache) and attention_mask.ndim == 2:
1233
+ if model_inputs["inputs_embeds"] is not None:
1234
+ batch_size, sequence_length, _ = model_inputs["inputs_embeds"].shape
1235
+ device = model_inputs["inputs_embeds"].device
1236
+ else:
1237
+ batch_size, sequence_length = model_inputs["input_ids"].shape
1238
+ device = model_inputs["input_ids"].device
1239
+
1240
+ dtype = self.embed_out.weight.dtype
1241
+ min_dtype = torch.finfo(dtype).min
1242
+
1243
+ attention_mask = _prepare_4d_causal_attention_mask_with_cache_position(
1244
+ attention_mask,
1245
+ sequence_length=sequence_length,
1246
+ target_length=past_key_values.get_max_length(),
1247
+ dtype=dtype,
1248
+ device=device,
1249
+ min_dtype=min_dtype,
1250
+ cache_position=cache_position,
1251
+ batch_size=batch_size,
1252
+ )
1253
+
1254
+ model_inputs.update(
1255
+ {
1256
+ "position_ids": position_ids,
1257
+ "cache_position": cache_position,
1258
+ "past_key_values": past_key_values,
1259
+ "use_cache": use_cache,
1260
+ "attention_mask": attention_mask,
1261
+ }
1262
+ )
1263
+ return model_inputs
1264
+
1265
+ def _reorder_cache(self, past_key_values, beam_idx):
1266
+ reordered_past = ()
1267
+ for layer_past in past_key_values:
1268
+ reordered_past += (
1269
+ tuple(past_state.index_select(0, beam_idx.to(past_state.device)) for past_state in layer_past[:2])
1270
+ + layer_past[2:],
1271
+ )
1272
+ return reordered_past
1273
+
1274
+
1275
+ @add_start_docstrings(
1276
+ """
1277
+ The GPTNeoX Model transformer with a sequence classification head on top (linear layer).
1278
+
1279
+ [`GPTNeoXForSequenceClassification`] uses the last token in order to do the classification, as other causal models
1280
+ (e.g. GPT-1) do.
1281
+
1282
+ Since it does classification on the last token, it requires to know the position of the last token. If a
1283
+ `pad_token_id` is defined in the configuration, it finds the last token that is not a padding token in each row. If
1284
+ no `pad_token_id` is defined, it simply takes the last value in each row of the batch. Since it cannot guess the
1285
+ padding tokens when `inputs_embeds` are passed instead of `input_ids`, it does the same (take the last value in
1286
+ each row of the batch).
1287
+ """,
1288
+ GPT_NEOX_START_DOCSTRING,
1289
+ )
1290
+ class GPTNeoXForSequenceClassification(GPTNeoXPreTrainedModel):
1291
+ def __init__(self, config):
1292
+ super().__init__(config)
1293
+ self.num_labels = config.num_labels
1294
+ self.gpt_neox = GPTNeoXModel(config)
1295
+ self.score = nn.Linear(config.hidden_size, self.num_labels, bias=False)
1296
+
1297
+ # Initialize weights and apply final processing
1298
+ self.post_init()
1299
+
1300
+ @add_start_docstrings_to_model_forward(GPT_NEOX_INPUTS_DOCSTRING)
1301
+ @add_code_sample_docstrings(
1302
+ checkpoint=_CHECKPOINT_FOR_DOC,
1303
+ output_type=SequenceClassifierOutputWithPast,
1304
+ config_class=_CONFIG_FOR_DOC,
1305
+ )
1306
+ def forward(
1307
+ self,
1308
+ input_ids: Optional[torch.LongTensor] = None,
1309
+ attention_mask: Optional[torch.FloatTensor] = None,
1310
+ position_ids: Optional[torch.LongTensor] = None,
1311
+ inputs_embeds: Optional[torch.FloatTensor] = None,
1312
+ head_mask: Optional[torch.FloatTensor] = None,
1313
+ past_key_values: Optional[Union[Cache, Tuple[Tuple[torch.FloatTensor]]]] = None,
1314
+ labels: Optional[torch.LongTensor] = None,
1315
+ use_cache: Optional[bool] = None,
1316
+ output_attentions: Optional[bool] = None,
1317
+ output_hidden_states: Optional[bool] = None,
1318
+ return_dict: Optional[bool] = None,
1319
+ ) -> Union[Tuple[torch.Tensor], SequenceClassifierOutputWithPast]:
1320
+ r"""
1321
+ labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
1322
+ Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
1323
+ config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
1324
+ `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
1325
+ """
1326
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
1327
+
1328
+ outputs = self.gpt_neox(
1329
+ input_ids,
1330
+ attention_mask=attention_mask,
1331
+ position_ids=position_ids,
1332
+ head_mask=head_mask,
1333
+ inputs_embeds=inputs_embeds,
1334
+ past_key_values=past_key_values,
1335
+ use_cache=use_cache,
1336
+ output_attentions=output_attentions,
1337
+ output_hidden_states=output_hidden_states,
1338
+ return_dict=return_dict,
1339
+ )
1340
+ hidden_states = outputs[0]
1341
+ logits = self.score(hidden_states)
1342
+
1343
+ if input_ids is not None:
1344
+ batch_size, sequence_length = input_ids.shape[:2]
1345
+ else:
1346
+ batch_size, sequence_length = inputs_embeds.shape[:2]
1347
+
1348
+ if self.config.pad_token_id is None and batch_size != 1:
1349
+ raise ValueError("Cannot handle batch sizes > 1 if no padding token is defined.")
1350
+ if self.config.pad_token_id is None:
1351
+ sequence_lengths = -1
1352
+ else:
1353
+ if input_ids is not None:
1354
+ # if no pad token found, use modulo instead of reverse indexing for ONNX compatibility
1355
+ sequence_lengths = torch.eq(input_ids, self.config.pad_token_id).int().argmax(-1) - 1
1356
+ sequence_lengths = sequence_lengths % input_ids.shape[-1]
1357
+ sequence_lengths = sequence_lengths.to(logits.device)
1358
+ else:
1359
+ sequence_lengths = -1
1360
+ logger.warning_once(
1361
+ f"{self.__class__.__name__} will not detect padding tokens in `inputs_embeds`. Results may be "
1362
+ "unexpected if using padding tokens in conjunction with `inputs_embeds.`"
1363
+ )
1364
+
1365
+ pooled_logits = logits[torch.arange(batch_size, device=logits.device), sequence_lengths]
1366
+
1367
+ loss = None
1368
+ if labels is not None:
1369
+ labels = labels.to(logits.device)
1370
+ if self.config.problem_type is None:
1371
+ if self.num_labels == 1:
1372
+ self.config.problem_type = "regression"
1373
+ elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
1374
+ self.config.problem_type = "single_label_classification"
1375
+ else:
1376
+ self.config.problem_type = "multi_label_classification"
1377
+
1378
+ if self.config.problem_type == "regression":
1379
+ loss_fct = MSELoss()
1380
+ if self.num_labels == 1:
1381
+ loss = loss_fct(pooled_logits.squeeze(), labels.squeeze())
1382
+ else:
1383
+ loss = loss_fct(pooled_logits, labels)
1384
+ elif self.config.problem_type == "single_label_classification":
1385
+ loss_fct = CrossEntropyLoss()
1386
+ loss = loss_fct(pooled_logits.view(-1, self.num_labels), labels.view(-1))
1387
+ elif self.config.problem_type == "multi_label_classification":
1388
+ loss_fct = BCEWithLogitsLoss()
1389
+ loss = loss_fct(pooled_logits, labels)
1390
+ if not return_dict:
1391
+ output = (pooled_logits,) + outputs[1:]
1392
+ return ((loss,) + output) if loss is not None else output
1393
+
1394
+ return SequenceClassifierOutputWithPast(
1395
+ loss=loss,
1396
+ logits=pooled_logits,
1397
+ past_key_values=outputs.past_key_values,
1398
+ hidden_states=outputs.hidden_states,
1399
+ attentions=outputs.attentions,
1400
+ )
1401
+
1402
+
1403
+ class GPTNeoXForTokenClassification(GPTNeoXPreTrainedModel):
1404
+ def __init__(self, config):
1405
+ super().__init__(config)
1406
+ self.num_labels = config.num_labels
1407
+
1408
+ self.gpt_neox = GPTNeoXModel(config)
1409
+ self.dropout = nn.Dropout(config.classifier_dropout)
1410
+ self.classifier = nn.Linear(config.hidden_size, config.num_labels)
1411
+
1412
+ # Initialize weights and apply final processing
1413
+ self.post_init()
1414
+
1415
+ @add_start_docstrings_to_model_forward(GPT_NEOX_INPUTS_DOCSTRING)
1416
+ @add_code_sample_docstrings(
1417
+ checkpoint="LarsJonasson/pythia-410m-deduped-sft-swedish",
1418
+ output_type=TokenClassifierOutput,
1419
+ config_class=_CONFIG_FOR_DOC,
1420
+ expected_loss=0.25,
1421
+ )
1422
+ def forward(
1423
+ self,
1424
+ input_ids: Optional[torch.LongTensor] = None,
1425
+ past_key_values: Optional[Union[Cache, Tuple[Tuple[torch.Tensor]]]] = None,
1426
+ attention_mask: Optional[torch.FloatTensor] = None,
1427
+ token_type_ids: Optional[torch.LongTensor] = None,
1428
+ position_ids: Optional[torch.LongTensor] = None,
1429
+ head_mask: Optional[torch.FloatTensor] = None,
1430
+ inputs_embeds: Optional[torch.FloatTensor] = None,
1431
+ labels: Optional[torch.LongTensor] = None,
1432
+ use_cache: Optional[bool] = None,
1433
+ output_attentions: Optional[bool] = None,
1434
+ output_hidden_states: Optional[bool] = None,
1435
+ return_dict: Optional[bool] = None,
1436
+ ) -> Union[Tuple, TokenClassifierOutput]:
1437
+ r"""
1438
+ labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
1439
+ Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
1440
+ config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
1441
+ `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
1442
+ """
1443
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
1444
+
1445
+ outputs = self.gpt_neox(
1446
+ input_ids,
1447
+ past_key_values=past_key_values,
1448
+ attention_mask=attention_mask,
1449
+ position_ids=position_ids,
1450
+ head_mask=head_mask,
1451
+ inputs_embeds=inputs_embeds,
1452
+ use_cache=use_cache,
1453
+ output_attentions=output_attentions,
1454
+ output_hidden_states=output_hidden_states,
1455
+ return_dict=return_dict,
1456
+ )
1457
+
1458
+ hidden_states = outputs[0]
1459
+ hidden_states = self.dropout(hidden_states)
1460
+ logits = self.classifier(hidden_states)
1461
+
1462
+ loss = None
1463
+ if labels is not None:
1464
+ labels = labels.to(logits.device)
1465
+ loss_fct = CrossEntropyLoss()
1466
+ loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
1467
+
1468
+ if not return_dict:
1469
+ output = (logits,) + outputs[2:]
1470
+ return ((loss,) + output) if loss is not None else output
1471
+
1472
+ return TokenClassifierOutput(
1473
+ loss=loss,
1474
+ logits=logits,
1475
+ hidden_states=outputs.hidden_states,
1476
+ attentions=outputs.attentions,
1477
+ )
1478
+
1479
+
1480
+ @add_start_docstrings(
1481
+ """
1482
+ The GPT-NeoX Model transformer with a span classification head on top for extractive question-answering tasks like
1483
+ SQuAD (a linear layer on top of the hidden-states output to compute `span start logits` and `span end logits`).
1484
+ """,
1485
+ GPT_NEOX_START_DOCSTRING,
1486
+ )
1487
+ class GPTNeoXForQuestionAnswering(GPTNeoXPreTrainedModel):
1488
+ def __init__(self, config):
1489
+ super().__init__(config)
1490
+ self.num_labels = config.num_labels
1491
+ self.gpt_neox = GPTNeoXModel(config)
1492
+ self.qa_outputs = nn.Linear(config.hidden_size, 2)
1493
+
1494
+ # Initialize weights and apply final processing
1495
+ self.post_init()
1496
+
1497
+ @add_start_docstrings_to_model_forward(GPT_NEOX_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
1498
+ @add_code_sample_docstrings(
1499
+ checkpoint=_CHECKPOINT_FOR_DOC,
1500
+ output_type=QuestionAnsweringModelOutput,
1501
+ config_class=_CONFIG_FOR_DOC,
1502
+ real_checkpoint=_REAL_CHECKPOINT_FOR_DOC,
1503
+ )
1504
+ def forward(
1505
+ self,
1506
+ input_ids: Optional[torch.LongTensor] = None,
1507
+ attention_mask: Optional[torch.FloatTensor] = None,
1508
+ token_type_ids: Optional[torch.LongTensor] = None,
1509
+ position_ids: Optional[torch.LongTensor] = None,
1510
+ head_mask: Optional[torch.FloatTensor] = None,
1511
+ inputs_embeds: Optional[torch.FloatTensor] = None,
1512
+ start_positions: Optional[torch.LongTensor] = None,
1513
+ end_positions: Optional[torch.LongTensor] = None,
1514
+ output_attentions: Optional[bool] = None,
1515
+ output_hidden_states: Optional[bool] = None,
1516
+ return_dict: Optional[bool] = None,
1517
+ ) -> Union[Tuple, QuestionAnsweringModelOutput]:
1518
+ r"""
1519
+ start_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
1520
+ Labels for position (index) of the start of the labelled span for computing the token classification loss.
1521
+ Positions are clamped to the length of the sequence (`sequence_length`). Position outside of the sequence
1522
+ are not taken into account for computing the loss.
1523
+ end_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
1524
+ Labels for position (index) of the end of the labelled span for computing the token classification loss.
1525
+ Positions are clamped to the length of the sequence (`sequence_length`). Position outside of the sequence
1526
+ are not taken into account for computing the loss.
1527
+ """
1528
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
1529
+
1530
+ outputs = self.gpt_neox(
1531
+ input_ids,
1532
+ attention_mask=attention_mask,
1533
+ position_ids=position_ids,
1534
+ head_mask=head_mask,
1535
+ inputs_embeds=inputs_embeds,
1536
+ output_attentions=output_attentions,
1537
+ output_hidden_states=output_hidden_states,
1538
+ return_dict=return_dict,
1539
+ )
1540
+
1541
+ sequence_output = outputs[0]
1542
+
1543
+ logits = self.qa_outputs(sequence_output)
1544
+ start_logits, end_logits = logits.split(1, dim=-1)
1545
+ start_logits = start_logits.squeeze(-1).contiguous()
1546
+ end_logits = end_logits.squeeze(-1).contiguous()
1547
+
1548
+ total_loss = None
1549
+ if start_positions is not None and end_positions is not None:
1550
+ # If we are on multi-GPU, split add a dimension
1551
+ if len(start_positions.size()) > 1:
1552
+ start_positions = start_positions.squeeze(-1).to(start_logits.device)
1553
+ if len(end_positions.size()) > 1:
1554
+ end_positions = end_positions.squeeze(-1).to(end_logits.device)
1555
+ # sometimes the start/end positions are outside our model inputs, we ignore these terms
1556
+ ignored_index = start_logits.size(1)
1557
+ start_positions = start_positions.clamp(0, ignored_index)
1558
+ end_positions = end_positions.clamp(0, ignored_index)
1559
+
1560
+ loss_fct = CrossEntropyLoss(ignore_index=ignored_index)
1561
+ start_loss = loss_fct(start_logits, start_positions)
1562
+ end_loss = loss_fct(end_logits, end_positions)
1563
+ total_loss = (start_loss + end_loss) / 2
1564
+
1565
+ if not return_dict:
1566
+ output = (start_logits, end_logits) + outputs[2:]
1567
+ return ((total_loss,) + output) if total_loss is not None else output
1568
+
1569
+ return QuestionAnsweringModelOutput(
1570
+ loss=total_loss,
1571
+ start_logits=start_logits,
1572
+ end_logits=end_logits,
1573
+ hidden_states=outputs.hidden_states,
1574
+ attentions=outputs.attentions,
1575
+ )
special_tokens_map.json ADDED
@@ -0,0 +1,24 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "bos_token": {
3
+ "content": "<|endoftext|>",
4
+ "lstrip": false,
5
+ "normalized": false,
6
+ "rstrip": false,
7
+ "single_word": false
8
+ },
9
+ "eos_token": {
10
+ "content": "<|endoftext|>",
11
+ "lstrip": false,
12
+ "normalized": false,
13
+ "rstrip": false,
14
+ "single_word": false
15
+ },
16
+ "pad_token": "<|endoftext|>",
17
+ "unk_token": {
18
+ "content": "<|endoftext|>",
19
+ "lstrip": false,
20
+ "normalized": false,
21
+ "rstrip": false,
22
+ "single_word": false
23
+ }
24
+ }
tokenizer.json ADDED
The diff for this file is too large to render. See raw diff
 
tokenizer_config.json ADDED
@@ -0,0 +1,214 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "add_bos_token": false,
3
+ "add_eos_token": false,
4
+ "add_prefix_space": false,
5
+ "added_tokens_decoder": {
6
+ "0": {
7
+ "content": "<|endoftext|>",
8
+ "lstrip": false,
9
+ "normalized": false,
10
+ "rstrip": false,
11
+ "single_word": false,
12
+ "special": true
13
+ },
14
+ "1": {
15
+ "content": "<|padding|>",
16
+ "lstrip": false,
17
+ "normalized": false,
18
+ "rstrip": false,
19
+ "single_word": false,
20
+ "special": true
21
+ },
22
+ "50254": {
23
+ "content": " ",
24
+ "lstrip": false,
25
+ "normalized": true,
26
+ "rstrip": false,
27
+ "single_word": false,
28
+ "special": false
29
+ },
30
+ "50255": {
31
+ "content": " ",
32
+ "lstrip": false,
33
+ "normalized": true,
34
+ "rstrip": false,
35
+ "single_word": false,
36
+ "special": false
37
+ },
38
+ "50256": {
39
+ "content": " ",
40
+ "lstrip": false,
41
+ "normalized": true,
42
+ "rstrip": false,
43
+ "single_word": false,
44
+ "special": false
45
+ },
46
+ "50257": {
47
+ "content": " ",
48
+ "lstrip": false,
49
+ "normalized": true,
50
+ "rstrip": false,
51
+ "single_word": false,
52
+ "special": false
53
+ },
54
+ "50258": {
55
+ "content": " ",
56
+ "lstrip": false,
57
+ "normalized": true,
58
+ "rstrip": false,
59
+ "single_word": false,
60
+ "special": false
61
+ },
62
+ "50259": {
63
+ "content": " ",
64
+ "lstrip": false,
65
+ "normalized": true,
66
+ "rstrip": false,
67
+ "single_word": false,
68
+ "special": false
69
+ },
70
+ "50260": {
71
+ "content": " ",
72
+ "lstrip": false,
73
+ "normalized": true,
74
+ "rstrip": false,
75
+ "single_word": false,
76
+ "special": false
77
+ },
78
+ "50261": {
79
+ "content": " ",
80
+ "lstrip": false,
81
+ "normalized": true,
82
+ "rstrip": false,
83
+ "single_word": false,
84
+ "special": false
85
+ },
86
+ "50262": {
87
+ "content": " ",
88
+ "lstrip": false,
89
+ "normalized": true,
90
+ "rstrip": false,
91
+ "single_word": false,
92
+ "special": false
93
+ },
94
+ "50263": {
95
+ "content": " ",
96
+ "lstrip": false,
97
+ "normalized": true,
98
+ "rstrip": false,
99
+ "single_word": false,
100
+ "special": false
101
+ },
102
+ "50264": {
103
+ "content": " ",
104
+ "lstrip": false,
105
+ "normalized": true,
106
+ "rstrip": false,
107
+ "single_word": false,
108
+ "special": false
109
+ },
110
+ "50265": {
111
+ "content": " ",
112
+ "lstrip": false,
113
+ "normalized": true,
114
+ "rstrip": false,
115
+ "single_word": false,
116
+ "special": false
117
+ },
118
+ "50266": {
119
+ "content": " ",
120
+ "lstrip": false,
121
+ "normalized": true,
122
+ "rstrip": false,
123
+ "single_word": false,
124
+ "special": false
125
+ },
126
+ "50267": {
127
+ "content": " ",
128
+ "lstrip": false,
129
+ "normalized": true,
130
+ "rstrip": false,
131
+ "single_word": false,
132
+ "special": false
133
+ },
134
+ "50268": {
135
+ "content": " ",
136
+ "lstrip": false,
137
+ "normalized": true,
138
+ "rstrip": false,
139
+ "single_word": false,
140
+ "special": false
141
+ },
142
+ "50269": {
143
+ "content": " ",
144
+ "lstrip": false,
145
+ "normalized": true,
146
+ "rstrip": false,
147
+ "single_word": false,
148
+ "special": false
149
+ },
150
+ "50270": {
151
+ "content": " ",
152
+ "lstrip": false,
153
+ "normalized": true,
154
+ "rstrip": false,
155
+ "single_word": false,
156
+ "special": false
157
+ },
158
+ "50271": {
159
+ "content": " ",
160
+ "lstrip": false,
161
+ "normalized": true,
162
+ "rstrip": false,
163
+ "single_word": false,
164
+ "special": false
165
+ },
166
+ "50272": {
167
+ "content": " ",
168
+ "lstrip": false,
169
+ "normalized": true,
170
+ "rstrip": false,
171
+ "single_word": false,
172
+ "special": false
173
+ },
174
+ "50273": {
175
+ "content": " ",
176
+ "lstrip": false,
177
+ "normalized": true,
178
+ "rstrip": false,
179
+ "single_word": false,
180
+ "special": false
181
+ },
182
+ "50274": {
183
+ "content": " ",
184
+ "lstrip": false,
185
+ "normalized": true,
186
+ "rstrip": false,
187
+ "single_word": false,
188
+ "special": false
189
+ },
190
+ "50275": {
191
+ "content": " ",
192
+ "lstrip": false,
193
+ "normalized": true,
194
+ "rstrip": false,
195
+ "single_word": false,
196
+ "special": false
197
+ },
198
+ "50276": {
199
+ "content": " ",
200
+ "lstrip": false,
201
+ "normalized": true,
202
+ "rstrip": false,
203
+ "single_word": false,
204
+ "special": false
205
+ }
206
+ },
207
+ "bos_token": "<|endoftext|>",
208
+ "clean_up_tokenization_spaces": true,
209
+ "eos_token": "<|endoftext|>",
210
+ "model_max_length": 1000000000000000019884624838656,
211
+ "pad_token": "<|endoftext|>",
212
+ "tokenizer_class": "GPTNeoXTokenizer",
213
+ "unk_token": "<|endoftext|>"
214
+ }
training_args.bin ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:446228e93e6bcc929afde9fbd97866c60e49ecdadb13996b0c6cc76828fe253c
3
+ size 5368