H4nwei commited on
Commit
716c816
1 Parent(s): 7933a89

Upload 5 files

Browse files
configuration_mplug_owl2.py ADDED
@@ -0,0 +1,334 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright (c) Alibaba.
2
+ #
3
+ # This source code is licensed under the license found in the
4
+ # LICENSE file in the root directory of this source tree.
5
+ import copy
6
+ import os
7
+ from typing import Union
8
+
9
+ from transformers.configuration_utils import PretrainedConfig
10
+ from transformers.models.auto.modeling_auto import MODEL_FOR_CAUSAL_LM_MAPPING_NAMES
11
+ from transformers.utils import logging
12
+ from transformers.models.auto import CONFIG_MAPPING
13
+
14
+
15
+ class LlamaConfig(PretrainedConfig):
16
+ r"""
17
+ This is the configuration class to store the configuration of a [`LlamaModel`]. It is used to instantiate an LLaMA
18
+ model according to the specified arguments, defining the model architecture. Instantiating a configuration with the
19
+ defaults will yield a similar configuration to that of the LLaMA-7B.
20
+
21
+ Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
22
+ documentation from [`PretrainedConfig`] for more information.
23
+
24
+
25
+ Args:
26
+ vocab_size (`int`, *optional*, defaults to 32000):
27
+ Vocabulary size of the LLaMA model. Defines the number of different tokens that can be represented by the
28
+ `inputs_ids` passed when calling [`LlamaModel`]
29
+ hidden_size (`int`, *optional*, defaults to 4096):
30
+ Dimension of the hidden representations.
31
+ intermediate_size (`int`, *optional*, defaults to 11008):
32
+ Dimension of the MLP representations.
33
+ num_hidden_layers (`int`, *optional*, defaults to 32):
34
+ Number of hidden layers in the Transformer decoder.
35
+ num_attention_heads (`int`, *optional*, defaults to 32):
36
+ Number of attention heads for each attention layer in the Transformer decoder.
37
+ num_key_value_heads (`int`, *optional*):
38
+ This is the number of key_value heads that should be used to implement Grouped Query Attention. If
39
+ `num_key_value_heads=num_attention_heads`, the model will use Multi Head Attention (MHA), if
40
+ `num_key_value_heads=1 the model will use Multi Query Attention (MQA) otherwise GQA is used. When
41
+ converting a multi-head checkpoint to a GQA checkpoint, each group key and value head should be constructed
42
+ by meanpooling all the original heads within that group. For more details checkout [this
43
+ paper](https://arxiv.org/pdf/2305.13245.pdf). If it is not specified, will default to
44
+ `num_attention_heads`.
45
+ hidden_act (`str` or `function`, *optional*, defaults to `"silu"`):
46
+ The non-linear activation function (function or string) in the decoder.
47
+ max_position_embeddings (`int`, *optional*, defaults to 2048):
48
+ The maximum sequence length that this model might ever be used with. Llama 1 supports up to 2048 tokens,
49
+ Llama 2 up to 4096, CodeLlama up to 16384.
50
+ initializer_range (`float`, *optional*, defaults to 0.02):
51
+ The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
52
+ rms_norm_eps (`float`, *optional*, defaults to 1e-06):
53
+ The epsilon used by the rms normalization layers.
54
+ use_cache (`bool`, *optional*, defaults to `True`):
55
+ Whether or not the model should return the last key/values attentions (not used by all models). Only
56
+ relevant if `config.is_decoder=True`.
57
+ pad_token_id (`int`, *optional*):
58
+ Padding token id.
59
+ bos_token_id (`int`, *optional*, defaults to 1):
60
+ Beginning of stream token id.
61
+ eos_token_id (`int`, *optional*, defaults to 2):
62
+ End of stream token id.
63
+ pretraining_tp (`int`, *optional*, defaults to 1):
64
+ Experimental feature. Tensor parallelism rank used during pretraining. Please refer to [this
65
+ document](https://huggingface.co/docs/transformers/parallelism) to understand more about it. This value is
66
+ necessary to ensure exact reproducibility of the pretraining results. Please refer to [this
67
+ issue](https://github.com/pytorch/pytorch/issues/76232).
68
+ tie_word_embeddings (`bool`, *optional*, defaults to `False`):
69
+ Whether to tie weight embeddings
70
+ rope_theta (`float`, *optional*, defaults to 10000.0):
71
+ The base period of the RoPE embeddings.
72
+ rope_scaling (`Dict`, *optional*):
73
+ Dictionary containing the scaling configuration for the RoPE embeddings. Currently supports two scaling
74
+ strategies: linear and dynamic. Their scaling factor must be a float greater than 1. The expected format is
75
+ `{"type": strategy name, "factor": scaling factor}`. When using this flag, don't update
76
+ `max_position_embeddings` to the expected new maximum. See the following thread for more information on how
77
+ these scaling strategies behave:
78
+ https://www.reddit.com/r/LocalLLaMA/comments/14mrgpr/dynamically_scaled_rope_further_increases/. This is an
79
+ experimental feature, subject to breaking API changes in future versions.
80
+ attention_bias (`bool`, defaults to `False`, *optional*, defaults to `False`):
81
+ Whether to use a bias in the query, key, value and output projection layers during self-attention.
82
+
83
+
84
+ ```python
85
+ >>> from transformers import LlamaModel, LlamaConfig
86
+
87
+ >>> # Initializing a LLaMA llama-7b style configuration
88
+ >>> configuration = LlamaConfig()
89
+
90
+ >>> # Initializing a model from the llama-7b style configuration
91
+ >>> model = LlamaModel(configuration)
92
+
93
+ >>> # Accessing the model configuration
94
+ >>> configuration = model.config
95
+ ```"""
96
+ model_type = "llama"
97
+ keys_to_ignore_at_inference = ["past_key_values"]
98
+
99
+ def __init__(
100
+ self,
101
+ vocab_size=32000,
102
+ hidden_size=4096,
103
+ intermediate_size=11008,
104
+ num_hidden_layers=32,
105
+ num_attention_heads=32,
106
+ num_key_value_heads=None,
107
+ hidden_act="silu",
108
+ max_position_embeddings=2048,
109
+ initializer_range=0.02,
110
+ rms_norm_eps=1e-6,
111
+ use_cache=True,
112
+ pad_token_id=None,
113
+ bos_token_id=1,
114
+ eos_token_id=2,
115
+ pretraining_tp=1,
116
+ tie_word_embeddings=False,
117
+ rope_theta=10000.0,
118
+ rope_scaling=None,
119
+ attention_bias=False,
120
+ attention_dropout=0.0,
121
+ **kwargs,
122
+ ):
123
+ self.vocab_size = vocab_size
124
+ self.max_position_embeddings = max_position_embeddings
125
+ self.hidden_size = hidden_size
126
+ self.intermediate_size = intermediate_size
127
+ self.num_hidden_layers = num_hidden_layers
128
+ self.num_attention_heads = num_attention_heads
129
+
130
+ # for backward compatibility
131
+ if num_key_value_heads is None:
132
+ num_key_value_heads = num_attention_heads
133
+
134
+ self.num_key_value_heads = num_key_value_heads
135
+ self.hidden_act = hidden_act
136
+ self.initializer_range = initializer_range
137
+ self.rms_norm_eps = rms_norm_eps
138
+ self.pretraining_tp = pretraining_tp
139
+ self.use_cache = use_cache
140
+ self.rope_theta = rope_theta
141
+ self.rope_scaling = rope_scaling
142
+ self._rope_scaling_validation()
143
+ self.attention_bias = attention_bias
144
+ self.attention_dropout = attention_dropout
145
+
146
+ super().__init__(
147
+ pad_token_id=pad_token_id,
148
+ bos_token_id=bos_token_id,
149
+ eos_token_id=eos_token_id,
150
+ tie_word_embeddings=tie_word_embeddings,
151
+ **kwargs,
152
+ )
153
+
154
+ def _rope_scaling_validation(self):
155
+ """
156
+ Validate the `rope_scaling` configuration.
157
+ """
158
+ if self.rope_scaling is None:
159
+ return
160
+
161
+ if not isinstance(self.rope_scaling, dict) or len(self.rope_scaling) != 2:
162
+ raise ValueError(
163
+ "`rope_scaling` must be a dictionary with with two fields, `type` and `factor`, "
164
+ f"got {self.rope_scaling}"
165
+ )
166
+ rope_scaling_type = self.rope_scaling.get("type", None)
167
+ rope_scaling_factor = self.rope_scaling.get("factor", None)
168
+ if rope_scaling_type is None or rope_scaling_type not in ["linear", "dynamic"]:
169
+ raise ValueError(
170
+ f"`rope_scaling`'s type field must be one of ['linear', 'dynamic'], got {rope_scaling_type}"
171
+ )
172
+ if rope_scaling_factor is None or not isinstance(rope_scaling_factor, float) or rope_scaling_factor <= 1.0:
173
+ raise ValueError(f"`rope_scaling`'s factor field must be a float > 1, got {rope_scaling_factor}")
174
+
175
+
176
+ class MplugOwlVisionConfig(PretrainedConfig):
177
+ r"""
178
+ This is the configuration class to store the configuration of a [`MplugOwlVisionModel`]. It is used to instantiate
179
+ a
180
+ mPLUG-Owl vision encoder according to the specified arguments, defining the model architecture. Instantiating a
181
+ configuration defaults will yield a similar configuration to that of the mPLUG-Owl
182
+ [x-plug/x_plug-llama-7b](https://huggingface.co/x-plug/x_plug-llama-7b) architecture.
183
+
184
+ Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
185
+ documentation from [`PretrainedConfig`] for more information.
186
+
187
+ Args:
188
+ hidden_size (`int`, *optional*, defaults to 768):
189
+ Dimensionality of the encoder layers and the pooler layer.
190
+ intermediate_size (`int`, *optional*, defaults to 3072):
191
+ Dimensionality of the "intermediate" (i.e., feed-forward) layer in the Transformer encoder.
192
+ num_hidden_layers (`int`, *optional*, defaults to 12):
193
+ Number of hidden layers in the Transformer encoder.
194
+ num_attention_heads (`int`, *optional*, defaults to 12):
195
+ Number of attention heads for each attention layer in the Transformer encoder.
196
+ image_size (`int`, *optional*, defaults to 224):
197
+ The size (resolution) of each image.
198
+ patch_size (`int`, *optional*, defaults to 32):
199
+ The size (resolution) of each patch.
200
+ hidden_act (`str` or `function`, *optional*, defaults to `"quick_gelu"`):
201
+ The non-linear activation function (function or string) in the encoder and pooler. If string, `"gelu"`,
202
+ `"relu"`, `"selu"` and `"gelu_new"` ``"quick_gelu"` are supported.
203
+ layer_norm_eps (`float`, *optional*, defaults to 1e-5):
204
+ The epsilon used by the layer normalization layers.
205
+ attention_dropout (`float`, *optional*, defaults to 0.0):
206
+ The dropout ratio for the attention probabilities.
207
+ initializer_range (`float`, *optional*, defaults to 0.02):
208
+ The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
209
+ initializer_factor (`float`, *optional*, defaults to 1):
210
+ A factor for initializing all weight matrices (should be kept to 1, used internally for initialization
211
+ testing).
212
+
213
+
214
+ ```"""
215
+
216
+ model_type = "mplug_owl_vision_model"
217
+
218
+ def __init__(
219
+ self,
220
+ hidden_size=1024,
221
+ intermediate_size=4096,
222
+ projection_dim=768,
223
+ num_hidden_layers=24,
224
+ num_attention_heads=16,
225
+ num_channels=3,
226
+ image_size=448,
227
+ patch_size=14,
228
+ hidden_act="quick_gelu",
229
+ layer_norm_eps=1e-6,
230
+ attention_dropout=0.0,
231
+ initializer_range=0.02,
232
+ initializer_factor=1.0,
233
+ use_flash_attn=False,
234
+ **kwargs,
235
+ ):
236
+ super().__init__(**kwargs)
237
+ self.hidden_size = hidden_size
238
+ self.intermediate_size = intermediate_size
239
+ self.projection_dim = projection_dim
240
+ self.num_hidden_layers = num_hidden_layers
241
+ self.num_attention_heads = num_attention_heads
242
+ self.num_channels = num_channels
243
+ self.patch_size = patch_size
244
+ self.image_size = image_size
245
+ self.initializer_range = initializer_range
246
+ self.initializer_factor = initializer_factor
247
+ self.attention_dropout = attention_dropout
248
+ self.layer_norm_eps = layer_norm_eps
249
+ self.hidden_act = hidden_act
250
+ self.use_flash_attn = use_flash_attn
251
+
252
+ @classmethod
253
+ def from_pretrained(cls, pretrained_model_name_or_path: Union[str, os.PathLike], **kwargs) -> "PretrainedConfig":
254
+ config_dict, kwargs = cls.get_config_dict(pretrained_model_name_or_path, **kwargs)
255
+
256
+ # get the vision config dict if we are loading from MplugOwlConfig
257
+ if config_dict.get("model_type") == "mplug-owl":
258
+ config_dict = config_dict["vision_config"]
259
+
260
+ if "model_type" in config_dict and hasattr(cls, "model_type") and config_dict["model_type"] != cls.model_type:
261
+ logger.warning(
262
+ f"You are using a model of type {config_dict['model_type']} to instantiate a model of type "
263
+ f"{cls.model_type}. This is not supported for all configurations of models and can yield errors."
264
+ )
265
+
266
+ return cls.from_dict(config_dict, **kwargs)
267
+
268
+
269
+ class MplugOwlVisualAbstractorConfig(PretrainedConfig):
270
+ model_type = "mplug_owl_visual_abstract"
271
+
272
+ def __init__(
273
+ self,
274
+ num_learnable_queries=64,
275
+ hidden_size=1024,
276
+ num_hidden_layers=6,
277
+ num_attention_heads=16,
278
+ intermediate_size=2816,
279
+ attention_probs_dropout_prob=0.,
280
+ initializer_range=0.02,
281
+ layer_norm_eps=1e-6,
282
+ encoder_hidden_size=1024,
283
+ grid_size=None,
284
+ **kwargs,
285
+ ):
286
+ super().__init__(**kwargs)
287
+ self.hidden_size = hidden_size
288
+ self.num_learnable_queries = num_learnable_queries
289
+ self.num_hidden_layers = num_hidden_layers
290
+ self.num_attention_heads = num_attention_heads
291
+ self.intermediate_size = intermediate_size
292
+ self.attention_probs_dropout_prob = attention_probs_dropout_prob
293
+ self.initializer_range = initializer_range
294
+ self.layer_norm_eps = layer_norm_eps
295
+ self.encoder_hidden_size = encoder_hidden_size
296
+ self.grid_size = grid_size if grid_size else 32
297
+
298
+ @classmethod
299
+ def from_pretrained(cls, pretrained_model_name_or_path: Union[str, os.PathLike], **kwargs) -> "PretrainedConfig":
300
+ config_dict, kwargs = cls.get_config_dict(pretrained_model_name_or_path, **kwargs)
301
+
302
+ # get the visual_abstractor config dict if we are loading from MplugOwlConfig
303
+ if config_dict.get("model_type") == "mplug-owl":
304
+ config_dict = config_dict["abstractor_config"]
305
+
306
+ if "model_type" in config_dict and hasattr(cls, "model_type") and config_dict["model_type"] != cls.model_type:
307
+ logger.warning(
308
+ f"You are using a model of type {config_dict['model_type']} to instantiate a model of type "
309
+ f"{cls.model_type}. This is not supported for all configurations of models and can yield errors."
310
+ )
311
+
312
+ return cls.from_dict(config_dict, **kwargs)
313
+
314
+
315
+
316
+ DEFAULT_VISUAL_CONFIG = {
317
+ "visual_model": MplugOwlVisionConfig().to_dict(),
318
+ "visual_abstractor": MplugOwlVisualAbstractorConfig().to_dict()
319
+ }
320
+
321
+ class MPLUGOwl2Config(LlamaConfig):
322
+ model_type = "mplug_owl2"
323
+ def __init__(self, visual_config=None, **kwargs):
324
+ if visual_config is None:
325
+ self.visual_config = DEFAULT_VISUAL_CONFIG
326
+ else:
327
+ self.visual_config = visual_config
328
+
329
+ super().__init__(
330
+ **kwargs,
331
+ )
332
+
333
+ if __name__ == "__main__":
334
+ print(MplugOwlVisionConfig().to_dict())
modeling_attn_mask_utils.py ADDED
@@ -0,0 +1,247 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright 2023 The HuggingFace 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.
5
+ # You may obtain a copy of the License at
6
+ #
7
+ # http://www.apache.org/licenses/LICENSE-2.0
8
+ #
9
+ # Unless required by applicable law or agreed to in writing, software
10
+ # distributed under the License is distributed on an "AS IS" BASIS,
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
+ from typing import List, Optional, Tuple, Union
15
+
16
+ import torch
17
+
18
+
19
+ class AttentionMaskConverter:
20
+ """
21
+ A utility attention mask class that allows one to:
22
+ - Create a causal 4d mask
23
+ - Create a causal 4d mask with slided window
24
+ - Convert a 2d attention mask (batch_size, query_length) to a 4d attention mask (batch_size, 1, query_length,
25
+ key_value_length) that can be multiplied with attention scores
26
+
27
+ Parameters:
28
+ is_causal (`bool`):
29
+ Whether the attention mask should be a uni-directional (causal) or bi-directional mask.
30
+
31
+ sliding_window (`int`, *optional*):
32
+ Optionally, the sliding window masks can be created if `sliding_window` is defined to a positive integer.
33
+ """
34
+
35
+ def __init__(self, is_causal: bool, sliding_window: Optional[int] = None):
36
+ self.is_causal = is_causal
37
+ self.sliding_window = sliding_window
38
+
39
+ if self.sliding_window is not None and self.sliding_window <= 0:
40
+ raise ValueError(
41
+ f"Make sure that when passing `sliding_window` that its value is a strictly positive integer, not `{self.sliding_window}`"
42
+ )
43
+
44
+ def to_causal_4d(
45
+ self,
46
+ batch_size: int,
47
+ query_length: int,
48
+ key_value_length: int,
49
+ dtype: torch.dtype = torch.float32,
50
+ device: Union[torch.device, "str"] = "cpu",
51
+ ) -> torch.Tensor:
52
+ """
53
+ Creates a causal 4D mask of (bsz, head_dim=1, query_length, key_value_length) shape and adds large negative
54
+ bias to upper right hand triangular matrix (causal mask).
55
+ """
56
+ if not self.is_causal:
57
+ raise ValueError(f"Please use `to_causal_4d` only if {self.__class__} has `is_causal` set to True.")
58
+
59
+ # If shape is not cached, create a new causal mask and cache it
60
+ input_shape = (batch_size, query_length)
61
+ past_key_values_length = key_value_length - query_length
62
+
63
+ # create causal mask
64
+ # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
65
+ causal_4d_mask = None
66
+ if input_shape[-1] > 1 or self.sliding_window is not None:
67
+ causal_4d_mask = self._make_causal_mask(
68
+ input_shape,
69
+ dtype,
70
+ device=device,
71
+ past_key_values_length=past_key_values_length,
72
+ sliding_window=self.sliding_window,
73
+ )
74
+
75
+ return causal_4d_mask
76
+
77
+ def to_4d(
78
+ self,
79
+ attention_mask_2d: torch.Tensor,
80
+ query_length: int,
81
+ key_value_length: Optional[int] = None,
82
+ dtype: torch.dtype = torch.float32,
83
+ ) -> torch.Tensor:
84
+ """
85
+ Converts 2D attention mask to 4D attention mask by expanding mask to (bsz, head_dim=1, query_length,
86
+ key_value_length) shape and by adding a large negative bias to not-attended positions. If attention_mask is
87
+ causal, a causal mask will be added.
88
+ """
89
+ input_shape = (attention_mask_2d.shape[0], query_length)
90
+
91
+ # create causal mask
92
+ # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
93
+ causal_4d_mask = None
94
+ if (input_shape[-1] > 1 or self.sliding_window is not None) and self.is_causal:
95
+ if key_value_length is None:
96
+ raise ValueError(
97
+ "This attention mask converter is causal. Make sure to pass `key_value_length` to correctly create a causal mask."
98
+ )
99
+
100
+ past_key_values_length = key_value_length - query_length
101
+ causal_4d_mask = self._make_causal_mask(
102
+ input_shape,
103
+ dtype,
104
+ device=attention_mask_2d.device,
105
+ past_key_values_length=past_key_values_length,
106
+ sliding_window=self.sliding_window,
107
+ )
108
+ elif self.sliding_window is not None:
109
+ raise NotImplementedError("Sliding window is currently only implemented for causal masking")
110
+
111
+ # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
112
+ expanded_attn_mask = self._expand_mask(attention_mask_2d, dtype, tgt_len=input_shape[-1]).to(
113
+ attention_mask_2d.device
114
+ )
115
+ expanded_4d_mask = expanded_attn_mask if causal_4d_mask is None else expanded_attn_mask + causal_4d_mask
116
+
117
+ return expanded_4d_mask
118
+
119
+ @staticmethod
120
+ def _make_causal_mask(
121
+ input_ids_shape: torch.Size,
122
+ dtype: torch.dtype,
123
+ device: torch.device,
124
+ past_key_values_length: int = 0,
125
+ sliding_window: Optional[int] = None,
126
+ ):
127
+ """
128
+ Make causal mask used for bi-directional self-attention.
129
+ """
130
+ bsz, tgt_len = input_ids_shape
131
+ mask = torch.full((tgt_len, tgt_len), torch.finfo(dtype).min, device=device)
132
+ mask_cond = torch.arange(mask.size(-1), device=device)
133
+ mask.masked_fill_(mask_cond < (mask_cond + 1).view(mask.size(-1), 1), 0)
134
+
135
+ mask = mask.to(dtype)
136
+
137
+ if past_key_values_length > 0:
138
+ mask = torch.cat([torch.zeros(tgt_len, past_key_values_length, dtype=dtype, device=device), mask], dim=-1)
139
+
140
+ # add lower triangular sliding window mask if necessary
141
+ if sliding_window is not None:
142
+ diagonal = past_key_values_length - sliding_window + 1
143
+
144
+ context_mask = 1 - torch.triu(torch.ones_like(mask, dtype=torch.int), diagonal=diagonal)
145
+ mask.masked_fill_(context_mask.bool(), torch.finfo(dtype).min)
146
+
147
+ return mask[None, None, :, :].expand(bsz, 1, tgt_len, tgt_len + past_key_values_length)
148
+
149
+ @staticmethod
150
+ def _expand_mask(mask: torch.Tensor, dtype: torch.dtype, tgt_len: Optional[int] = None):
151
+ """
152
+ Expands attention_mask from `[bsz, seq_len]` to `[bsz, 1, tgt_seq_len, src_seq_len]`.
153
+ """
154
+ bsz, src_len = mask.size()
155
+ tgt_len = tgt_len if tgt_len is not None else src_len
156
+
157
+ expanded_mask = mask[:, None, None, :].expand(bsz, 1, tgt_len, src_len).to(dtype)
158
+
159
+ inverted_mask = 1.0 - expanded_mask
160
+
161
+ return inverted_mask.masked_fill(inverted_mask.to(torch.bool), torch.finfo(dtype).min)
162
+
163
+
164
+ def _prepare_4d_causal_attention_mask(
165
+ attention_mask: Optional[torch.Tensor],
166
+ input_shape: Union[torch.Size, Tuple, List],
167
+ inputs_embeds: torch.Tensor,
168
+ past_key_values_length: int,
169
+ sliding_window: Optional[int] = None,
170
+ ):
171
+ """
172
+ Creates a causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape
173
+ `(batch_size, key_value_length)`
174
+
175
+ Args:
176
+ attention_mask (`torch.Tensor` or `None`):
177
+ A 2D attention mask of shape `(batch_size, key_value_length)`
178
+ input_shape (`tuple(int)` or `list(int)` or `torch.Size`):
179
+ The input shape should be a tuple that defines `(batch_size, query_length)`.
180
+ inputs_embeds (`torch.Tensor`):
181
+ The embedded inputs as a torch Tensor.
182
+ past_key_values_length (`int`):
183
+ The length of the key value cache.
184
+ sliding_window (`int`, *optional*):
185
+ If the model uses windowed attention, a sliding window should be passed.
186
+ """
187
+ attn_mask_converter = AttentionMaskConverter(is_causal=True, sliding_window=sliding_window)
188
+
189
+ key_value_length = input_shape[-1] + past_key_values_length
190
+
191
+ # 4d mask is passed through the layers
192
+ if attention_mask is not None:
193
+ attention_mask = attn_mask_converter.to_4d(
194
+ attention_mask, input_shape[-1], key_value_length, dtype=inputs_embeds.dtype
195
+ )
196
+ else:
197
+ attention_mask = attn_mask_converter.to_causal_4d(
198
+ input_shape[0], input_shape[-1], key_value_length, dtype=inputs_embeds.dtype, device=inputs_embeds.device
199
+ )
200
+
201
+ return attention_mask
202
+
203
+
204
+ def _prepare_4d_attention_mask(mask: torch.Tensor, dtype: torch.dtype, tgt_len: Optional[int] = None):
205
+ """
206
+ Creates a non-causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape
207
+ `(batch_size, key_value_length)`
208
+
209
+ Args:
210
+ mask (`torch.Tensor` or `None`):
211
+ A 2D attention mask of shape `(batch_size, key_value_length)`
212
+ dtype (`torch.dtype`):
213
+ The torch dtype the created mask shall have.
214
+ tgt_len (`int`):
215
+ The target length or query length the created mask shall have.
216
+ """
217
+ return AttentionMaskConverter._expand_mask(mask=mask, dtype=dtype, tgt_len=tgt_len)
218
+
219
+
220
+ def _create_4d_causal_attention_mask(
221
+ input_shape: Union[torch.Size, Tuple, List],
222
+ dtype: torch.dtype,
223
+ device: torch.device,
224
+ past_key_values_length: int = 0,
225
+ sliding_window: Optional[int] = None,
226
+ ):
227
+ """
228
+ Creates a causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)`
229
+
230
+ Args:
231
+ input_shape (`tuple(int)` or `list(int)` or `torch.Size`):
232
+ The input shape should be a tuple that defines `(batch_size, query_length)`.
233
+ dtype (`torch.dtype`):
234
+ The torch dtype the created mask shall have.
235
+ device (`int`):
236
+ The torch device the created mask shall have.
237
+ sliding_window (`int`, *optional*):
238
+ If the model uses windowed attention, a sliding window should be passed.
239
+ """
240
+ attn_mask_converter = AttentionMaskConverter(is_causal=True, sliding_window=sliding_window)
241
+
242
+ key_value_length = past_key_values_length + input_shape[-1]
243
+ attention_mask = attn_mask_converter.to_causal_4d(
244
+ input_shape[0], input_shape[-1], key_value_length, dtype=dtype, device=device
245
+ )
246
+
247
+ return attention_mask
modeling_llama2.py ADDED
@@ -0,0 +1,837 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import math
2
+ import warnings
3
+ from functools import partial
4
+ from typing import List, Optional, Tuple, Union
5
+
6
+ import torch
7
+ import torch.nn.functional as F
8
+ import torch.utils.checkpoint
9
+ from torch import nn
10
+ from transformers.modeling_attn_mask_utils import _prepare_4d_causal_attention_mask, _prepare_4d_causal_attention_mask_for_sdpa
11
+
12
+
13
+ import copy
14
+ import os
15
+ import sys
16
+
17
+ dir_path = os.path.dirname(os.path.realpath(__file__))
18
+ sys.path.insert(0, dir_path)
19
+
20
+ import transformers
21
+ from transformers.models.llama.modeling_llama import *
22
+
23
+ def _get_unpad_data(attention_mask):
24
+ seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32)
25
+ indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()
26
+ max_seqlen_in_batch = seqlens_in_batch.max().item()
27
+ cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.torch.int32), (1, 0))
28
+ return (
29
+ indices,
30
+ cu_seqlens,
31
+ max_seqlen_in_batch,
32
+ )
33
+
34
+
35
+ from transformers.configuration_utils import PretrainedConfig
36
+ from transformers.utils import logging
37
+
38
+ from .modeling_attn_mask_utils import _prepare_4d_causal_attention_mask
39
+ from .configuration_mplug_owl2 import LlamaConfig
40
+ # from modeling_attn_mask_utils import _prepare_4d_causal_attention_mask
41
+ # from configuration_mplug_owl2 import LlamaConfig
42
+
43
+ class MultiwayNetwork(nn.Module):
44
+
45
+ def __init__(self, module_provider, num_multiway=2):
46
+ super(MultiwayNetwork, self).__init__()
47
+
48
+ self.multiway = torch.nn.ModuleList([module_provider() for _ in range(num_multiway)])
49
+
50
+ def forward(self, hidden_states, multiway_indices):
51
+
52
+ if len(self.multiway) == 1:
53
+ return self.multiway[0](hidden_states)
54
+
55
+ output_hidden_states = torch.empty_like(hidden_states)
56
+
57
+ for idx, subway in enumerate(self.multiway):
58
+ local_indices = multiway_indices.eq(idx).nonzero(as_tuple=True)
59
+ hidden = hidden_states[local_indices].unsqueeze(1).contiguous()
60
+ if hidden.numel():
61
+ output = subway(hidden)
62
+ if isinstance(output, tuple):
63
+ output = output[0]
64
+ output = output.squeeze(1)
65
+ output_hidden_states[local_indices] = output
66
+
67
+ return output_hidden_states.contiguous()
68
+
69
+
70
+ class LlamaAttention(nn.Module):
71
+ """Multi-headed attention from 'Attention Is All You Need' paper"""
72
+
73
+ def __init__(self, config: LlamaConfig, layer_idx: Optional[int] = None):
74
+ super().__init__()
75
+ self.config = config
76
+ self.layer_idx = layer_idx
77
+ if layer_idx is None:
78
+ logger.warning_once(
79
+ f"Instantiating {self.__class__.__name__} without passing `layer_idx` is not recommended and will "
80
+ "to errors during the forward call, if caching is used. Please make sure to provide a `layer_idx` "
81
+ "when creating this class."
82
+ )
83
+
84
+ self.attention_dropout = config.attention_dropout
85
+ self.hidden_size = config.hidden_size
86
+ self.num_heads = config.num_attention_heads
87
+ self.head_dim = self.hidden_size // self.num_heads
88
+ self.num_key_value_heads = config.num_key_value_heads
89
+ self.num_key_value_groups = self.num_heads // self.num_key_value_heads
90
+ self.max_position_embeddings = config.max_position_embeddings
91
+ self.rope_theta = config.rope_theta
92
+ self.is_causal = True
93
+
94
+ if (self.head_dim * self.num_heads) != self.hidden_size:
95
+ raise ValueError(
96
+ f"hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size}"
97
+ f" and `num_heads`: {self.num_heads})."
98
+ )
99
+ self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=config.attention_bias)
100
+ self.k_proj = MultiwayNetwork(module_provider=partial(
101
+ nn.Linear, in_features=self.hidden_size, out_features=self.num_key_value_heads * self.head_dim, bias=config.attention_bias)
102
+ )
103
+ self.v_proj = MultiwayNetwork(module_provider=partial(
104
+ nn.Linear, in_features=self.hidden_size, out_features=self.num_key_value_heads * self.head_dim, bias=config.attention_bias)
105
+ )
106
+ self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=config.attention_bias)
107
+ self._init_rope()
108
+
109
+ def _init_rope(self):
110
+ if self.config.rope_scaling is None:
111
+ self.rotary_emb = LlamaRotaryEmbedding(
112
+ self.head_dim,
113
+ max_position_embeddings=self.max_position_embeddings,
114
+ base=self.rope_theta,
115
+ )
116
+ else:
117
+ scaling_type = self.config.rope_scaling["type"]
118
+ scaling_factor = self.config.rope_scaling["factor"]
119
+ if scaling_type == "linear":
120
+ self.rotary_emb = LlamaLinearScalingRotaryEmbedding(
121
+ self.head_dim,
122
+ max_position_embeddings=self.max_position_embeddings,
123
+ scaling_factor=scaling_factor,
124
+ base=self.rope_theta,
125
+ )
126
+ elif scaling_type == "dynamic":
127
+ self.rotary_emb = LlamaDynamicNTKScalingRotaryEmbedding(
128
+ self.head_dim,
129
+ max_position_embeddings=self.max_position_embeddings,
130
+ scaling_factor=scaling_factor,
131
+ base=self.rope_theta,
132
+ )
133
+ else:
134
+ raise ValueError(f"Unknown RoPE scaling type {scaling_type}")
135
+
136
+ def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int):
137
+ return tensor.view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2).contiguous()
138
+
139
+ def forward(
140
+ self,
141
+ hidden_states: torch.Tensor,
142
+ modality_indicators: torch.Tensor,
143
+ attention_mask: Optional[torch.Tensor] = None,
144
+ position_ids: Optional[torch.LongTensor] = None,
145
+ past_key_value: Optional[Tuple[torch.Tensor]] = None,
146
+ output_attentions: bool = False,
147
+ use_cache: bool = False,
148
+ padding_mask: Optional[torch.LongTensor] = None,
149
+ ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
150
+ bsz, q_len, _ = hidden_states.size()
151
+
152
+ query_states = self.q_proj(hidden_states, )
153
+ key_states = self.k_proj(hidden_states, modality_indicators)
154
+ value_states = self.v_proj(hidden_states, modality_indicators)
155
+
156
+ query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
157
+ key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
158
+ value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
159
+
160
+ kv_seq_len = key_states.shape[-2]
161
+ if past_key_value is not None:
162
+ kv_seq_len += past_key_value[0].shape[-2]
163
+ cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
164
+ query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)
165
+
166
+ if past_key_value is not None:
167
+ # reuse k, v, self_attention
168
+ key_states = torch.cat([past_key_value[0], key_states], dim=2)
169
+ value_states = torch.cat([past_key_value[1], value_states], dim=2)
170
+
171
+ past_key_value = (key_states, value_states) if use_cache else None
172
+
173
+ key_states = repeat_kv(key_states, self.num_key_value_groups)
174
+ value_states = repeat_kv(value_states, self.num_key_value_groups)
175
+
176
+ attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim)
177
+
178
+ if attn_weights.size() != (bsz, self.num_heads, q_len, kv_seq_len):
179
+ raise ValueError(
180
+ f"Attention weights should be of size {(bsz, self.num_heads, q_len, kv_seq_len)}, but is"
181
+ f" {attn_weights.size()}"
182
+ )
183
+
184
+ if attention_mask is not None:
185
+ if attention_mask.size() != (bsz, 1, q_len, kv_seq_len):
186
+ raise ValueError(
187
+ f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}"
188
+ )
189
+ attn_weights = attn_weights + attention_mask
190
+
191
+ # upcast attention to fp32
192
+ attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype)
193
+ attn_output = torch.matmul(attn_weights, value_states)
194
+
195
+ if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim):
196
+ raise ValueError(
197
+ f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.head_dim)}, but is"
198
+ f" {attn_output.size()}"
199
+ )
200
+
201
+ attn_output = attn_output.transpose(1, 2).contiguous()
202
+
203
+ attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)
204
+
205
+ attn_output = self.o_proj(attn_output)
206
+
207
+ if not output_attentions:
208
+ attn_weights = None
209
+
210
+ return attn_output, attn_weights, past_key_value
211
+
212
+
213
+ class LlamaFlashAttention2(LlamaAttention):
214
+ """
215
+ Llama flash attention module. This module inherits from `LlamaAttention` as the weights of the module stays
216
+ untouched. The only required change would be on the forward pass where it needs to correctly call the public API of
217
+ flash attention and deal with padding tokens in case the input contains any of them.
218
+ """
219
+
220
+ def __init__(self, *args, **kwargs):
221
+ super().__init__(*args, **kwargs)
222
+
223
+ # TODO: Should be removed once Flash Attention for RoCm is bumped to 2.1.
224
+ # 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.
225
+ # 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).
226
+ self._flash_attn_uses_top_left_mask = not is_flash_attn_greater_or_equal_2_10()
227
+
228
+ def forward(
229
+ self,
230
+ hidden_states: torch.Tensor,
231
+ modality_indicators: torch.Tensor,
232
+ attention_mask: Optional[torch.LongTensor] = None,
233
+ position_ids: Optional[torch.LongTensor] = None,
234
+ past_key_value: Optional[Cache] = None,
235
+ output_attentions: bool = False,
236
+ use_cache: bool = False,
237
+ **kwargs,
238
+ ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
239
+ # LlamaFlashAttention2 attention does not support output_attentions
240
+ if "padding_mask" in kwargs:
241
+ warnings.warn(
242
+ "Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use `attention_mask` instead.`"
243
+ )
244
+
245
+ # overwrite attention_mask with padding_mask
246
+ attention_mask = kwargs.pop("padding_mask")
247
+
248
+ output_attentions = False
249
+
250
+ bsz, q_len, _ = hidden_states.size()
251
+
252
+ query_states = self.q_proj(hidden_states)
253
+ key_states = self.k_proj(hidden_states, modality_indicators)
254
+ value_states = self.v_proj(hidden_states, modality_indicators)
255
+
256
+ # Flash attention requires the input to have the shape
257
+ # batch_size x seq_length x head_dim x hidden_dim
258
+ # therefore we just need to keep the original shape
259
+ query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
260
+ key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
261
+ value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
262
+
263
+ kv_seq_len = key_states.shape[-2]
264
+ if past_key_value is not None:
265
+ kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx)
266
+ cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
267
+ query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)
268
+
269
+ if past_key_value is not None:
270
+ cache_kwargs = {"sin": sin, "cos": cos} # Specific to RoPE models
271
+ key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
272
+
273
+ # 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
274
+ # to be able to avoid many of these transpose/reshape/view.
275
+ query_states = query_states.transpose(1, 2)
276
+ key_states = key_states.transpose(1, 2)
277
+ value_states = value_states.transpose(1, 2)
278
+
279
+ dropout_rate = self.attention_dropout if self.training else 0.0
280
+
281
+ # In PEFT, usually we cast the layer norms in float32 for training stability reasons
282
+ # therefore the input hidden states gets silently casted in float32. Hence, we need
283
+ # cast them back in the correct dtype just to be sure everything works as expected.
284
+ # This might slowdown training & inference so it is recommended to not cast the LayerNorms
285
+ # in fp32. (LlamaRMSNorm handles it correctly)
286
+
287
+ input_dtype = query_states.dtype
288
+ if input_dtype == torch.float32:
289
+ if torch.is_autocast_enabled():
290
+ target_dtype = torch.get_autocast_gpu_dtype()
291
+ # Handle the case where the model is quantized
292
+ elif hasattr(self.config, "_pre_quantization_dtype"):
293
+ target_dtype = self.config._pre_quantization_dtype
294
+ else:
295
+ target_dtype = self.q_proj.weight.dtype
296
+
297
+ logger.warning_once(
298
+ f"The input hidden states seems to be silently casted in float32, this might be related to"
299
+ f" the fact you have upcasted embedding or layer norm layers in float32. We will cast back the input in"
300
+ f" {target_dtype}."
301
+ )
302
+
303
+ query_states = query_states.to(target_dtype)
304
+ key_states = key_states.to(target_dtype)
305
+ value_states = value_states.to(target_dtype)
306
+
307
+ attn_output = self._flash_attention_forward(
308
+ query_states, key_states, value_states, attention_mask, q_len, dropout=dropout_rate
309
+ )
310
+
311
+ attn_output = attn_output.reshape(bsz, q_len, self.hidden_size).contiguous()
312
+ attn_output = self.o_proj(attn_output)
313
+
314
+ if not output_attentions:
315
+ attn_weights = None
316
+
317
+ return attn_output, attn_weights, past_key_value
318
+
319
+ def _flash_attention_forward(
320
+ self, query_states, key_states, value_states, attention_mask, query_length, dropout=0.0, softmax_scale=None
321
+ ):
322
+ """
323
+ Calls the forward method of Flash Attention - if the input hidden states contain at least one padding token
324
+ first unpad the input, then computes the attention scores and pad the final attention scores.
325
+
326
+ Args:
327
+ query_states (`torch.Tensor`):
328
+ Input query states to be passed to Flash Attention API
329
+ key_states (`torch.Tensor`):
330
+ Input key states to be passed to Flash Attention API
331
+ value_states (`torch.Tensor`):
332
+ Input value states to be passed to Flash Attention API
333
+ attention_mask (`torch.Tensor`):
334
+ The padding mask - corresponds to a tensor of size `(batch_size, seq_len)` where 0 stands for the
335
+ position of padding tokens and 1 for the position of non-padding tokens.
336
+ dropout (`int`, *optional*):
337
+ Attention dropout
338
+ softmax_scale (`float`, *optional*):
339
+ The scaling of QK^T before applying softmax. Default to 1 / sqrt(head_dim)
340
+ """
341
+ if not self._flash_attn_uses_top_left_mask:
342
+ causal = self.is_causal
343
+ else:
344
+ # 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__.
345
+ causal = self.is_causal and query_length != 1
346
+
347
+ # Contains at least one padding token in the sequence
348
+ if attention_mask is not None:
349
+ batch_size = query_states.shape[0]
350
+ query_states, key_states, value_states, indices_q, cu_seq_lens, max_seq_lens = self._upad_input(
351
+ query_states, key_states, value_states, attention_mask, query_length
352
+ )
353
+
354
+ cu_seqlens_q, cu_seqlens_k = cu_seq_lens
355
+ max_seqlen_in_batch_q, max_seqlen_in_batch_k = max_seq_lens
356
+
357
+ attn_output_unpad = flash_attn_varlen_func(
358
+ query_states,
359
+ key_states,
360
+ value_states,
361
+ cu_seqlens_q=cu_seqlens_q,
362
+ cu_seqlens_k=cu_seqlens_k,
363
+ max_seqlen_q=max_seqlen_in_batch_q,
364
+ max_seqlen_k=max_seqlen_in_batch_k,
365
+ dropout_p=dropout,
366
+ softmax_scale=softmax_scale,
367
+ causal=causal,
368
+ )
369
+
370
+ attn_output = pad_input(attn_output_unpad, indices_q, batch_size, query_length)
371
+ else:
372
+ attn_output = flash_attn_func(
373
+ query_states, key_states, value_states, dropout, softmax_scale=softmax_scale, causal=causal
374
+ )
375
+
376
+ return attn_output
377
+
378
+ def _upad_input(self, query_layer, key_layer, value_layer, attention_mask, query_length):
379
+ indices_k, cu_seqlens_k, max_seqlen_in_batch_k = _get_unpad_data(attention_mask)
380
+ batch_size, kv_seq_len, num_key_value_heads, head_dim = key_layer.shape
381
+
382
+ key_layer = index_first_axis(
383
+ key_layer.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim), indices_k
384
+ )
385
+ value_layer = index_first_axis(
386
+ value_layer.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim), indices_k
387
+ )
388
+ if query_length == kv_seq_len:
389
+ query_layer = index_first_axis(
390
+ query_layer.reshape(batch_size * kv_seq_len, self.num_heads, head_dim), indices_k
391
+ )
392
+ cu_seqlens_q = cu_seqlens_k
393
+ max_seqlen_in_batch_q = max_seqlen_in_batch_k
394
+ indices_q = indices_k
395
+ elif query_length == 1:
396
+ max_seqlen_in_batch_q = 1
397
+ cu_seqlens_q = torch.arange(
398
+ batch_size + 1, dtype=torch.int32, device=query_layer.device
399
+ ) # There is a memcpy here, that is very bad.
400
+ indices_q = cu_seqlens_q[:-1]
401
+ query_layer = query_layer.squeeze(1)
402
+ else:
403
+ # The -q_len: slice assumes left padding.
404
+ attention_mask = attention_mask[:, -query_length:]
405
+ query_layer, indices_q, cu_seqlens_q, max_seqlen_in_batch_q = unpad_input(query_layer, attention_mask)
406
+
407
+ return (
408
+ query_layer,
409
+ key_layer,
410
+ value_layer,
411
+ indices_q,
412
+ (cu_seqlens_q, cu_seqlens_k),
413
+ (max_seqlen_in_batch_q, max_seqlen_in_batch_k),
414
+ )
415
+
416
+
417
+ class LlamaSdpaAttention(LlamaAttention):
418
+ """
419
+ Llama attention module using torch.nn.functional.scaled_dot_product_attention. This module inherits from
420
+ `LlamaAttention` as the weights of the module stays untouched. The only changes are on the forward pass to adapt to
421
+ SDPA API.
422
+ """
423
+
424
+ # Adapted from LlamaAttention.forward
425
+ def forward(
426
+ self,
427
+ hidden_states: torch.Tensor,
428
+ modality_indicators: torch.Tensor,
429
+ attention_mask: Optional[torch.Tensor] = None,
430
+ position_ids: Optional[torch.LongTensor] = None,
431
+ past_key_value: Optional[Cache] = None,
432
+ output_attentions: bool = False,
433
+ use_cache: bool = False,
434
+ ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
435
+ if output_attentions:
436
+ # TODO: Improve this warning with e.g. `model.config.attn_implementation = "manual"` once this is implemented.
437
+ logger.warning_once(
438
+ "LlamaModel is using LlamaSdpaAttention, but `torch.nn.functional.scaled_dot_product_attention` does not support `output_attentions=True`. Falling back to the manual attention implementation, "
439
+ '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.'
440
+ )
441
+ return super().forward(
442
+ hidden_states=hidden_states,
443
+ modality_indicators=modality_indicators,
444
+ attention_mask=attention_mask,
445
+ position_ids=position_ids,
446
+ past_key_value=past_key_value,
447
+ output_attentions=output_attentions,
448
+ use_cache=use_cache,
449
+ )
450
+
451
+ bsz, q_len, _ = hidden_states.size()
452
+
453
+ query_states = self.q_proj(hidden_states)
454
+ key_states = self.k_proj(hidden_states, modality_indicators)
455
+ value_states = self.v_proj(hidden_states, modality_indicators)
456
+
457
+ query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
458
+ key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
459
+ value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
460
+
461
+ kv_seq_len = key_states.shape[-2]
462
+ if past_key_value is not None:
463
+ kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx)
464
+ cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
465
+
466
+ query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)
467
+
468
+ if past_key_value is not None:
469
+ cache_kwargs = {"sin": sin, "cos": cos} # Specific to RoPE models
470
+ key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
471
+
472
+ key_states = repeat_kv(key_states, self.num_key_value_groups)
473
+ value_states = repeat_kv(value_states, self.num_key_value_groups)
474
+
475
+ if attention_mask is not None:
476
+ if attention_mask.size() != (bsz, 1, q_len, kv_seq_len):
477
+ raise ValueError(
478
+ f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}"
479
+ )
480
+
481
+ # SDPA with memory-efficient backend is currently (torch==2.1.2) bugged with non-contiguous inputs with custom attn_mask,
482
+ # Reference: https://github.com/pytorch/pytorch/issues/112577.
483
+ if query_states.device.type == "cuda" and attention_mask is not None:
484
+ query_states = query_states.contiguous()
485
+ key_states = key_states.contiguous()
486
+ value_states = value_states.contiguous()
487
+
488
+ attn_output = torch.nn.functional.scaled_dot_product_attention(
489
+ query_states,
490
+ key_states,
491
+ value_states,
492
+ attn_mask=attention_mask,
493
+ dropout_p=self.attention_dropout if self.training else 0.0,
494
+ # 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.
495
+ is_causal=self.is_causal and attention_mask is None and q_len > 1,
496
+ )
497
+
498
+ attn_output = attn_output.transpose(1, 2).contiguous()
499
+ attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)
500
+
501
+ attn_output = self.o_proj(attn_output)
502
+
503
+ return attn_output, None, past_key_value
504
+
505
+
506
+
507
+ LLAMA_ATTENTION_CLASSES = {
508
+ "eager": LlamaAttention,
509
+ "flash_attention_2": LlamaFlashAttention2,
510
+ "sdpa": LlamaSdpaAttention,
511
+ }
512
+
513
+ class LlamaDecoderLayer(nn.Module):
514
+ def __init__(self, config: LlamaConfig, layer_idx):
515
+ super().__init__()
516
+ self.hidden_size = config.hidden_size
517
+ self.self_attn = LlamaAttention(config=config)
518
+ self.self_attn = LLAMA_ATTENTION_CLASSES[config._attn_implementation](config=config, layer_idx=layer_idx)
519
+ self.mlp = LlamaMLP(config)
520
+ self.input_layernorm = MultiwayNetwork(module_provider=partial(
521
+ LlamaRMSNorm, hidden_size=config.hidden_size, eps=config.rms_norm_eps
522
+ ))
523
+ self.post_attention_layernorm = MultiwayNetwork(module_provider=partial(
524
+ LlamaRMSNorm, hidden_size=config.hidden_size, eps=config.rms_norm_eps
525
+ ))
526
+
527
+ def forward(
528
+ self,
529
+ hidden_states: torch.Tensor,
530
+ modality_indicators: torch.Tensor = None,
531
+ attention_mask: Optional[torch.Tensor] = None,
532
+ position_ids: Optional[torch.LongTensor] = None,
533
+ past_key_value: Optional[Tuple[torch.Tensor]] = None,
534
+ output_attentions: Optional[bool] = False,
535
+ use_cache: Optional[bool] = False,
536
+ ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]:
537
+ """
538
+ Args:
539
+ hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`
540
+ attention_mask (`torch.FloatTensor`, *optional*): attention mask of size
541
+ `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values.
542
+ output_attentions (`bool`, *optional*):
543
+ Whether or not to return the attentions tensors of all attention layers. See `attentions` under
544
+ returned tensors for more detail.
545
+ use_cache (`bool`, *optional*):
546
+ If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding
547
+ (see `past_key_values`).
548
+ past_key_value (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states
549
+ """
550
+
551
+ residual = hidden_states
552
+
553
+ hidden_states = self.input_layernorm(hidden_states, modality_indicators)
554
+
555
+ # Self Attention
556
+ hidden_states, self_attn_weights, present_key_value = self.self_attn(
557
+ hidden_states=hidden_states,
558
+ modality_indicators=modality_indicators,
559
+ attention_mask=attention_mask,
560
+ position_ids=position_ids,
561
+ past_key_value=past_key_value,
562
+ output_attentions=output_attentions,
563
+ use_cache=use_cache,
564
+ )
565
+ hidden_states = residual + hidden_states
566
+
567
+ # Fully Connected
568
+ residual = hidden_states
569
+ hidden_states = self.post_attention_layernorm(hidden_states, modality_indicators)
570
+ hidden_states = self.mlp(hidden_states)
571
+ hidden_states = residual + hidden_states
572
+
573
+ outputs = (hidden_states,)
574
+
575
+ if output_attentions:
576
+ outputs += (self_attn_weights,)
577
+
578
+ if use_cache:
579
+ outputs += (present_key_value,)
580
+
581
+ return outputs
582
+
583
+
584
+ def model_forward(
585
+ self,
586
+ input_ids: torch.LongTensor = None,
587
+ modality_indicators: torch.Tensor = None,
588
+ attention_mask: Optional[torch.Tensor] = None,
589
+ position_ids: Optional[torch.LongTensor] = None,
590
+ past_key_values: Optional[List[torch.FloatTensor]] = None,
591
+ inputs_embeds: Optional[torch.FloatTensor] = None,
592
+ use_cache: Optional[bool] = None,
593
+ output_attentions: Optional[bool] = None,
594
+ output_hidden_states: Optional[bool] = None,
595
+ return_dict: Optional[bool] = None,
596
+ ) -> Union[Tuple, BaseModelOutputWithPast]:
597
+ output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
598
+ output_hidden_states = (
599
+ output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
600
+ )
601
+ use_cache = use_cache if use_cache is not None else self.config.use_cache
602
+
603
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
604
+
605
+ # retrieve input_ids and inputs_embeds
606
+ if input_ids is not None and inputs_embeds is not None:
607
+ raise ValueError("You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time")
608
+ elif input_ids is not None:
609
+ batch_size, seq_length = input_ids.shape
610
+ elif inputs_embeds is not None:
611
+ batch_size, seq_length, _ = inputs_embeds.shape
612
+ else:
613
+ raise ValueError("You have to specify either decoder_input_ids or decoder_inputs_embeds")
614
+
615
+ seq_length_with_past = seq_length
616
+ past_key_values_length = 0
617
+
618
+ if past_key_values is not None:
619
+ past_key_values_length = past_key_values[0][0].shape[2]
620
+ seq_length_with_past = seq_length_with_past + past_key_values_length
621
+
622
+ if position_ids is None:
623
+ device = input_ids.device if input_ids is not None else inputs_embeds.device
624
+ position_ids = torch.arange(
625
+ past_key_values_length, seq_length + past_key_values_length, dtype=torch.long, device=device
626
+ )
627
+ position_ids = position_ids.unsqueeze(0).view(-1, seq_length)
628
+ else:
629
+ position_ids = position_ids.view(-1, seq_length).long()
630
+
631
+ if inputs_embeds is None:
632
+ inputs_embeds = self.embed_tokens(input_ids)
633
+ # embed positions
634
+ if attention_mask is None:
635
+ attention_mask = torch.ones(
636
+ (batch_size, seq_length_with_past), dtype=torch.bool, device=inputs_embeds.device
637
+ )
638
+
639
+ if self._use_flash_attention_2:
640
+ # 2d mask is passed through the layers
641
+ attention_mask = attention_mask if (attention_mask is not None and 0 in attention_mask) else None
642
+ elif self._use_sdpa and not output_attentions:
643
+ # output_attentions=True can not be supported when using SDPA, and we fall back on
644
+ # the manual implementation that requires a 4D causal mask in all cases.
645
+ attention_mask = _prepare_4d_causal_attention_mask_for_sdpa(
646
+ attention_mask,
647
+ (batch_size, seq_length),
648
+ inputs_embeds,
649
+ past_key_values_length,
650
+ )
651
+ else:
652
+ # 4d mask is passed through the layers
653
+ attention_mask = _prepare_4d_causal_attention_mask(
654
+ attention_mask, (batch_size, seq_length), inputs_embeds, past_key_values_length
655
+ )
656
+
657
+ hidden_states = inputs_embeds
658
+
659
+ if self.gradient_checkpointing and self.training:
660
+ if use_cache:
661
+ logger.warning_once(
662
+ "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
663
+ )
664
+ use_cache = False
665
+
666
+ # decoder layers
667
+ all_hidden_states = () if output_hidden_states else None
668
+ all_self_attns = () if output_attentions else None
669
+ next_decoder_cache = () if use_cache else None
670
+
671
+ for idx, decoder_layer in enumerate(self.layers):
672
+ if output_hidden_states:
673
+ all_hidden_states += (hidden_states,)
674
+
675
+ past_key_value = past_key_values[idx] if past_key_values is not None else None
676
+
677
+ if self.gradient_checkpointing and self.training:
678
+
679
+ def create_custom_forward(module):
680
+ def custom_forward(*inputs):
681
+ # None for past_key_value
682
+ return module(*inputs, past_key_value, output_attentions)
683
+
684
+ return custom_forward
685
+
686
+ layer_outputs = torch.utils.checkpoint.checkpoint(
687
+ create_custom_forward(decoder_layer),
688
+ hidden_states,
689
+ modality_indicators,
690
+ attention_mask,
691
+ position_ids,
692
+ )
693
+ else:
694
+ layer_outputs = decoder_layer(
695
+ hidden_states,
696
+ modality_indicators=modality_indicators,
697
+ attention_mask=attention_mask,
698
+ position_ids=position_ids,
699
+ past_key_value=past_key_value,
700
+ output_attentions=output_attentions,
701
+ use_cache=use_cache,
702
+ )
703
+
704
+ hidden_states = layer_outputs[0]
705
+
706
+ if use_cache:
707
+ next_decoder_cache += (layer_outputs[2 if output_attentions else 1],)
708
+
709
+ if output_attentions:
710
+ all_self_attns += (layer_outputs[1],)
711
+
712
+ hidden_states = self.norm(hidden_states)
713
+
714
+ # add hidden states from the last decoder layer
715
+ if output_hidden_states:
716
+ all_hidden_states += (hidden_states,)
717
+
718
+ next_cache = next_decoder_cache if use_cache else None
719
+ if not return_dict:
720
+ return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None)
721
+ return BaseModelOutputWithPast(
722
+ last_hidden_state=hidden_states,
723
+ past_key_values=next_cache,
724
+ hidden_states=all_hidden_states,
725
+ attentions=all_self_attns,
726
+ )
727
+
728
+
729
+ def causal_model_forward(
730
+ self,
731
+ input_ids: torch.LongTensor = None,
732
+ modality_indicators: torch.Tensor = None,
733
+ attention_mask: Optional[torch.Tensor] = None,
734
+ position_ids: Optional[torch.LongTensor] = None,
735
+ past_key_values: Optional[List[torch.FloatTensor]] = None,
736
+ inputs_embeds: Optional[torch.FloatTensor] = None,
737
+ labels: Optional[torch.LongTensor] = None,
738
+ use_cache: Optional[bool] = None,
739
+ output_attentions: Optional[bool] = None,
740
+ output_hidden_states: Optional[bool] = None,
741
+ return_dict: Optional[bool] = None,
742
+ ) -> Union[Tuple, CausalLMOutputWithPast]:
743
+ r"""
744
+ Args:
745
+ labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
746
+ Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
747
+ config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
748
+ (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
749
+
750
+ Returns:
751
+
752
+ Example:
753
+
754
+ ```python
755
+ >>> from transformers import AutoTokenizer, LlamaForCausalLM
756
+
757
+ >>> model = LlamaForCausalLM.from_pretrained(PATH_TO_CONVERTED_WEIGHTS)
758
+ >>> tokenizer = AutoTokenizer.from_pretrained(PATH_TO_CONVERTED_TOKENIZER)
759
+
760
+ >>> prompt = "Hey, are you conscious? Can you talk to me?"
761
+ >>> inputs = tokenizer(prompt, return_tensors="pt")
762
+
763
+ >>> # Generate
764
+ >>> generate_ids = model.generate(inputs.input_ids, max_length=30)
765
+ >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
766
+ "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you."
767
+ ```"""
768
+
769
+ output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
770
+ output_hidden_states = (
771
+ output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
772
+ )
773
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
774
+
775
+ # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
776
+ outputs = self.model(
777
+ input_ids=input_ids,
778
+ modality_indicators=modality_indicators,
779
+ attention_mask=attention_mask,
780
+ position_ids=position_ids,
781
+ past_key_values=past_key_values,
782
+ inputs_embeds=inputs_embeds,
783
+ use_cache=use_cache,
784
+ output_attentions=output_attentions,
785
+ output_hidden_states=output_hidden_states,
786
+ return_dict=return_dict,
787
+ )
788
+
789
+ hidden_states = outputs[0]
790
+ if self.config.pretraining_tp > 1:
791
+ lm_head_slices = self.lm_head.weight.split(self.vocab_size // self.config.pretraining_tp, dim=0)
792
+ logits = [F.linear(hidden_states, lm_head_slices[i]) for i in range(self.config.pretraining_tp)]
793
+ logits = torch.cat(logits, dim=-1)
794
+ else:
795
+ logits = self.lm_head(hidden_states)
796
+ logits = logits.float()
797
+
798
+ loss = None
799
+ if labels is not None:
800
+ # Shift so that tokens < n predict n
801
+ shift_logits = logits[..., :-1, :].contiguous()
802
+ shift_labels = labels[..., 1:].contiguous()
803
+ # Flatten the tokens
804
+ loss_fct = CrossEntropyLoss()
805
+ shift_logits = shift_logits.view(-1, self.config.vocab_size)
806
+ shift_labels = shift_labels.view(-1)
807
+ # Enable model parallelism
808
+ shift_labels = shift_labels.to(shift_logits.device)
809
+ loss = loss_fct(shift_logits, shift_labels)
810
+
811
+ if not return_dict:
812
+ output = (logits,) + outputs[1:]
813
+ return (loss,) + output if loss is not None else output
814
+
815
+ return CausalLMOutputWithPast(
816
+ loss=loss,
817
+ logits=logits,
818
+ past_key_values=outputs.past_key_values,
819
+ hidden_states=outputs.hidden_states,
820
+ attentions=outputs.attentions,
821
+ )
822
+
823
+ def replace_llama_modality_adaptive():
824
+ transformers.models.llama.configuration_llama.LlamaConfig = LlamaConfig
825
+ transformers.models.llama.modeling_llama.LlamaAttention = LlamaAttention
826
+ transformers.models.llama.modeling_llama.LlamaFlashAttention2 = LlamaFlashAttention2
827
+ transformers.models.llama.modeling_llama.LlamaSdpaAttention = LlamaSdpaAttention
828
+ transformers.models.llama.modeling_llama.LlamaDecoderLayer = LlamaDecoderLayer
829
+ transformers.models.llama.modeling_llama.LlamaModel.forward = model_forward
830
+ transformers.models.llama.modeling_llama.LlamaForCausalLM.forward = causal_model_forward
831
+
832
+
833
+ if __name__ == "__main__":
834
+ replace_llama_modality_adaptive()
835
+ config = transformers.LlamaConfig.from_pretrained('/cpfs01/shared/public/test/vicuna-7b-v1.5/')
836
+ model = transformers.LlamaForCausalLM(config)
837
+ print(model)
modeling_mplug_owl2.py ADDED
@@ -0,0 +1,469 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright 2023 Haotian Liu & Qinghao Ye (Modified from LLaVA)
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.
5
+ # You may obtain a copy of the License at
6
+ #
7
+ # http://www.apache.org/licenses/LICENSE-2.0
8
+ #
9
+ # Unless required by applicable law or agreed to in writing, software
10
+ # distributed under the License is distributed on an "AS IS" BASIS,
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
+ from abc import ABC, abstractmethod
16
+ from typing import List, Optional, Tuple, Union
17
+ from datasets import load_dataset
18
+ import torch
19
+ import torch.nn as nn
20
+ from torch.nn import CrossEntropyLoss
21
+ import numpy as np
22
+ import copy
23
+ import os
24
+ import sys
25
+ from PIL import Image
26
+ import requests
27
+ from io import BytesIO
28
+
29
+ dir_path = os.path.dirname(os.path.realpath(__file__))
30
+ sys.path.insert(0, dir_path)
31
+
32
+ from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer, CLIPImageProcessor, LlamaConfig, LlamaModel, LlamaForCausalLM
33
+ from transformers.modeling_outputs import CausalLMOutputWithPast
34
+
35
+ from configuration_mplug_owl2 import MPLUGOwl2Config, MplugOwlVisionConfig, MplugOwlVisualAbstractorConfig
36
+ from visual_encoder import MplugOwlVisionModel, MplugOwlVisualAbstractorModel
37
+ from modeling_llama2 import replace_llama_modality_adaptive
38
+ IGNORE_INDEX = -100
39
+ IMAGE_TOKEN_INDEX = -200
40
+ DEFAULT_IMAGE_TOKEN = "<|image|>"
41
+ from icecream import ic
42
+
43
+ def tokenizer_image_token(prompt, tokenizer, image_token_index=IMAGE_TOKEN_INDEX, return_tensors=None):
44
+ prompt_chunks = [tokenizer(chunk).input_ids if len(chunk) > 0 else [] for chunk in prompt.split(DEFAULT_IMAGE_TOKEN)]
45
+
46
+ def insert_separator(X, sep):
47
+ return [ele for sublist in zip(X, [sep]*len(X)) for ele in sublist][:-1]
48
+
49
+ input_ids = []
50
+ offset = 0
51
+ if len(prompt_chunks) > 0 and len(prompt_chunks[0]) > 0 and prompt_chunks[0][0] == tokenizer.bos_token_id:
52
+ offset = 1
53
+ input_ids.append(prompt_chunks[0][0])
54
+
55
+ for x in insert_separator(prompt_chunks, [image_token_index] * (offset + 1)):
56
+ input_ids.extend(x[offset:])
57
+
58
+ if return_tensors is not None:
59
+ if return_tensors == 'pt':
60
+ return torch.tensor(input_ids, dtype=torch.long)
61
+ raise ValueError(f'Unsupported tensor type: {return_tensors}')
62
+ return input_ids
63
+
64
+ def expand2square(pil_img, background_color):
65
+ from PIL import Image
66
+ width, height = pil_img.size
67
+ if width == height:
68
+ return pil_img
69
+ elif width > height:
70
+ result = Image.new(pil_img.mode, (width, width), background_color)
71
+ result.paste(pil_img, (0, (width - height) // 2))
72
+ return result
73
+ else:
74
+ result = Image.new(pil_img.mode, (height, height), background_color)
75
+ result.paste(pil_img, ((height - width) // 2, 0))
76
+ return result
77
+
78
+ def norm_cdf(x):
79
+ return 0.5 * (1 + torch.erf(x / torch.sqrt(torch.tensor(2.0))))
80
+
81
+ def optimize_score_map_pytorch_cuda(c, seed=0, original_seed=20020, num_iterations=100):
82
+ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
83
+
84
+ torch.manual_seed(seed)
85
+ np.random.seed(seed)
86
+
87
+ c = torch.tensor(c, dtype=torch.float32, device=device, requires_grad=False)
88
+ initial_scores = torch.rand(c.shape[0], device=device, requires_grad=True)
89
+
90
+ optimizer = torch.optim.Adam([initial_scores], lr=0.1)
91
+
92
+ for _ in range(num_iterations):
93
+ optimizer.zero_grad()
94
+ sum_log_diff = torch.sum(c * torch.log(torch.maximum(torch.sigmoid(initial_scores[:, None] - initial_scores), torch.tensor(1e-6, device=device))))
95
+ sum_squares = torch.sum(initial_scores ** 2) / 2
96
+
97
+ loss = -(sum_log_diff - sum_squares)
98
+ loss.backward()
99
+ optimizer.step()
100
+
101
+ optimized_scores = initial_scores.detach().cpu().numpy()
102
+ min_score, max_score = np.min(optimized_scores), np.max(optimized_scores)
103
+
104
+ # Scale scores to 0-100
105
+ scaled_scores = 100 * (optimized_scores - min_score) / (max_score - min_score)
106
+
107
+ # Reset the seed
108
+ np.random.seed(original_seed)
109
+ return scaled_scores[-1]
110
+
111
+ def softmax(logits):
112
+ # exp_logits = np.exp(logits - np.max(logits))
113
+ probs = np.exp(logits) / np.sum(np.exp(logits))
114
+ return probs
115
+ # return exp_logits / exp_logits.sum()
116
+
117
+ def update_matrix(anchor_matrix, scores, indices):
118
+ n = anchor_matrix.shape[0]
119
+ new_row = np.zeros((1, n))
120
+ new_col = np.zeros((n + 1, 1))
121
+ new_row[0, indices] = scores
122
+ new_col[indices, 0] = 1-scores # Assuming symmetric preference for simplicity
123
+ anchor_matrix = np.vstack([anchor_matrix, new_row])
124
+ anchor_matrix = np.hstack([anchor_matrix, new_col])
125
+
126
+ return anchor_matrix
127
+
128
+
129
+ class MPLUGOwl2MetaModel:
130
+ def __init__(self, config):
131
+ super(MPLUGOwl2MetaModel, self).__init__(config)
132
+ self.vision_model = MplugOwlVisionModel(
133
+ MplugOwlVisionConfig(**config.visual_config["visual_model"])
134
+ )
135
+ self.visual_abstractor = MplugOwlVisualAbstractorModel(
136
+ MplugOwlVisualAbstractorConfig(**config.visual_config["visual_abstractor"]), config.hidden_size
137
+ )
138
+
139
+ def get_vision_tower(self):
140
+ vision_model = getattr(self, 'vision_model', None)
141
+ if type(vision_model) is list:
142
+ vision_model = vision_model[0]
143
+ return vision_model
144
+
145
+ def get_visual_abstractor(self):
146
+ visual_abstractor = getattr(self, 'visual_abstractor', None)
147
+ if type(visual_abstractor) is list:
148
+ visual_abstractor = visual_abstractor[0]
149
+ return visual_abstractor
150
+
151
+
152
+ class MPLUGOwl2MetaForCausalLM(ABC):
153
+ @abstractmethod
154
+ def get_model(self):
155
+ pass
156
+
157
+ def encode_images(self, images):
158
+ image_features = self.get_model().vision_model(images).last_hidden_state
159
+ image_features = self.get_model().visual_abstractor(encoder_hidden_states=image_features).last_hidden_state
160
+ return image_features
161
+
162
+ def prepare_inputs_labels_for_multimodal(
163
+ self, input_ids, attention_mask, past_key_values, labels, images
164
+ ):
165
+ if images is None or input_ids.shape[1] == 1:
166
+ if past_key_values is not None and images is not None and input_ids.shape[1] == 1:
167
+ attention_mask = torch.ones((attention_mask.shape[0], past_key_values[-1][-1].shape[-2] + 1), dtype=attention_mask.dtype, device=attention_mask.device)
168
+ multiway_indices = torch.zeros_like(input_ids).long().to(self.device)
169
+ return input_ids, multiway_indices, attention_mask, past_key_values, None, labels
170
+
171
+ if type(images) is list or images.ndim == 5:
172
+ concat_images = torch.cat([image for image in images], dim=0)
173
+ image_features = self.encode_images(concat_images)
174
+ split_sizes = [image.shape[0] for image in images]
175
+ image_features = torch.split(image_features, split_sizes, dim=0)
176
+ image_features = [x.flatten(0, 1) for x in image_features]
177
+ else:
178
+ image_features = self.encode_images(images)
179
+
180
+ new_input_embeds = []
181
+ new_modality_indicators = []
182
+ new_labels = [] if labels is not None else None
183
+ cur_image_idx = 0
184
+ for batch_idx, cur_input_ids in enumerate(input_ids):
185
+ if (cur_input_ids == IMAGE_TOKEN_INDEX).sum() == 0:
186
+ # multimodal LLM, but the current sample is not multimodal
187
+ # FIXME: this is a hacky fix, for deepspeed zero3 to work
188
+ half_len = cur_input_ids.shape[0] // 2
189
+ cur_image_features = image_features[cur_image_idx]
190
+ cur_input_embeds_1 = self.get_model().embed_tokens(cur_input_ids[:half_len])
191
+ cur_input_embeds_2 = self.get_model().embed_tokens(cur_input_ids[half_len:])
192
+ cur_input_embeds = torch.cat([cur_input_embeds_1, cur_image_features[0:0], cur_input_embeds_2], dim=0)
193
+ new_input_embeds.append(cur_input_embeds)
194
+
195
+ cur_modality_indicators = torch.zeros(len(cur_input_embeds)).long().to(self.device)
196
+ new_modality_indicators.append(cur_modality_indicators)
197
+ if labels is not None:
198
+ new_labels.append(labels[batch_idx])
199
+ cur_image_idx += 1
200
+ continue
201
+ image_token_indices = torch.where(cur_input_ids == IMAGE_TOKEN_INDEX)[0]
202
+ cur_new_input_embeds = []
203
+ cur_modality_indicators = []
204
+ if labels is not None:
205
+ cur_labels = labels[batch_idx]
206
+ cur_new_labels = []
207
+ assert cur_labels.shape == cur_input_ids.shape
208
+ while image_token_indices.numel() > 0:
209
+ # print("cur_image_idx", cur_image_idx)
210
+ cur_image_features = image_features[cur_image_idx]
211
+ image_token_start = image_token_indices[0]
212
+ cur_new_input_embeds.append(self.get_model().embed_tokens(cur_input_ids[:image_token_start]))
213
+ cur_new_input_embeds.append(cur_image_features)
214
+
215
+ # Add modality indicator
216
+ assert image_token_start == len(cur_input_ids[:image_token_start])
217
+ cur_modality_indicators.append(torch.zeros(len(cur_input_ids[:image_token_start])).long())
218
+ cur_modality_indicators.append(torch.ones(len(cur_image_features)).long())
219
+
220
+ if labels is not None:
221
+ cur_new_labels.append(cur_labels[:image_token_start])
222
+ cur_new_labels.append(torch.full((cur_image_features.shape[0],), IGNORE_INDEX, device=labels.device, dtype=labels.dtype))
223
+ cur_labels = cur_labels[image_token_start+1:]
224
+ cur_image_idx += 1
225
+ cur_input_ids = cur_input_ids[image_token_start+1:]
226
+ image_token_indices = torch.where(cur_input_ids == IMAGE_TOKEN_INDEX)[0]
227
+ if cur_input_ids.numel() > 0:
228
+ cur_new_input_embeds.append(self.get_model().embed_tokens(cur_input_ids))
229
+ cur_modality_indicators.append(torch.zeros(len(cur_input_ids)).long())
230
+ if labels is not None:
231
+ cur_new_labels.append(cur_labels)
232
+ cur_new_input_embeds = [x.to(device=self.device) for x in cur_new_input_embeds]
233
+ cur_new_input_embeds = torch.cat(cur_new_input_embeds, dim=0)
234
+ new_input_embeds.append(cur_new_input_embeds)
235
+
236
+ # Modality
237
+ cur_modality_indicators = [x.to(device=self.device) for x in cur_modality_indicators]
238
+ cur_modality_indicators = torch.cat(cur_modality_indicators, dim=0)
239
+ new_modality_indicators.append(cur_modality_indicators)
240
+
241
+
242
+ if labels is not None:
243
+ cur_new_labels = torch.cat(cur_new_labels, dim=0)
244
+ new_labels.append(cur_new_labels)
245
+
246
+ if any(x.shape != new_input_embeds[0].shape for x in new_input_embeds):
247
+ max_len = max(x.shape[0] for x in new_input_embeds)
248
+
249
+ # Embedding
250
+ new_input_embeds_align = []
251
+ for cur_new_embed in new_input_embeds:
252
+ cur_new_embed = torch.cat((cur_new_embed, torch.zeros((max_len - cur_new_embed.shape[0], cur_new_embed.shape[1]), dtype=cur_new_embed.dtype, device=cur_new_embed.device)), dim=0)
253
+ new_input_embeds_align.append(cur_new_embed)
254
+ new_input_embeds = torch.stack(new_input_embeds_align, dim=0)
255
+
256
+ # Modality
257
+ new_modality_indicators_align = []
258
+ for cur_modality_indicator in new_modality_indicators:
259
+ cur_new_embed = torch.cat((cur_modality_indicator, torch.zeros(max_len - cur_modality_indicator.shape[0], dtype=cur_modality_indicator.dtype, device=cur_modality_indicator.device)), dim=0)
260
+ new_modality_indicators_align.append(cur_new_embed)
261
+ new_modality_indicators = torch.stack(new_modality_indicators_align, dim=0)
262
+
263
+ # Label
264
+ if labels is not None:
265
+ new_labels_align = []
266
+ _new_labels = new_labels
267
+ for cur_new_label in new_labels:
268
+ cur_new_label = torch.cat((cur_new_label, torch.full((max_len - cur_new_label.shape[0],), IGNORE_INDEX, dtype=cur_new_label.dtype, device=cur_new_label.device)), dim=0)
269
+ new_labels_align.append(cur_new_label)
270
+ new_labels = torch.stack(new_labels_align, dim=0)
271
+
272
+ # Attention Mask
273
+ if attention_mask is not None:
274
+ new_attention_mask = []
275
+ for cur_attention_mask, cur_new_labels, cur_new_labels_align in zip(attention_mask, _new_labels, new_labels):
276
+ new_attn_mask_pad_left = torch.full((cur_new_labels.shape[0] - labels.shape[1],), True, dtype=attention_mask.dtype, device=attention_mask.device)
277
+ new_attn_mask_pad_right = torch.full((cur_new_labels_align.shape[0] - cur_new_labels.shape[0],), False, dtype=attention_mask.dtype, device=attention_mask.device)
278
+ cur_new_attention_mask = torch.cat((new_attn_mask_pad_left, cur_attention_mask, new_attn_mask_pad_right), dim=0)
279
+ new_attention_mask.append(cur_new_attention_mask)
280
+ attention_mask = torch.stack(new_attention_mask, dim=0)
281
+ assert attention_mask.shape == new_labels.shape
282
+ else:
283
+ new_input_embeds = torch.stack(new_input_embeds, dim=0)
284
+ new_modality_indicators = torch.stack(new_modality_indicators, dim=0)
285
+ if labels is not None:
286
+ new_labels = torch.stack(new_labels, dim=0)
287
+
288
+ if attention_mask is not None:
289
+ new_attn_mask_pad_left = torch.full((attention_mask.shape[0], new_input_embeds.shape[1] - input_ids.shape[1]), True, dtype=attention_mask.dtype, device=attention_mask.device)
290
+ attention_mask = torch.cat((new_attn_mask_pad_left, attention_mask), dim=1)
291
+ assert attention_mask.shape == new_input_embeds.shape[:2]
292
+ return None, new_modality_indicators, attention_mask, past_key_values, new_input_embeds, new_labels
293
+
294
+
295
+
296
+ class MPLUGOwl2LlamaModel(MPLUGOwl2MetaModel, LlamaModel):
297
+ config_class = MPLUGOwl2Config
298
+
299
+ def __init__(self, config: MPLUGOwl2Config):
300
+ super(MPLUGOwl2LlamaModel, self).__init__(config)
301
+
302
+
303
+ class MPLUGOwl2LlamaForCausalLM(LlamaForCausalLM, MPLUGOwl2MetaForCausalLM):
304
+ config_class = MPLUGOwl2Config
305
+
306
+ def __init__(self, config):
307
+ super(LlamaForCausalLM, self).__init__(config)
308
+ self.model = MPLUGOwl2LlamaModel(config)
309
+ self.tokenizer = AutoTokenizer.from_pretrained("VQA-CityU/Compare2Score_1")
310
+ self.image_processor = CLIPImageProcessor.from_pretrained("VQA-CityU/Compare2Score_1")
311
+
312
+ self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
313
+ self.preferential_ids_ = [id_[1] for id_ in self.tokenizer(["inferior", "worse", "similar", "better", "superior"])["input_ids"]]
314
+ self.anchor_images = load_dataset("VQA-CityU/Anchor_images")
315
+
316
+ self.weight_tensor = np.array([0., 0.25, 0.5, 0.75, 1.], dtype=np.float16)
317
+ self.anchor_matrix = np.array(
318
+ [[5.0000000e-01, 2.5912809e-01, 3.3130276e-04, 1.6087297e-06, 1.1803027e-09],
319
+ [7.4087191e-01, 5.0000000e-01, 2.4985345e-01, 9.9954158e-02, 1.8675303e-08],
320
+ [9.9966872e-01, 7.5014657e-01, 5.0000000e-01, 4.9968880e-01, 2.4852838e-01],
321
+ [9.9999839e-01, 9.0004587e-01, 5.0031120e-01, 5.0000000e-01, 2.5400183e-01],
322
+ [1.0000000e+00, 1.0000000e+00, 7.5147164e-01, 7.4599814e-01, 5.0000000e-01]],
323
+ dtype=np.float32)
324
+ anchor_intervals = 5#16
325
+ num_anchor_image_per_interval = 1
326
+ num_anchor_image = anchor_intervals * num_anchor_image_per_interval
327
+ self.anchor_indices = np.arange(0,num_anchor_image)
328
+ # Initialize weights and apply final processing
329
+ self.post_init()
330
+
331
+
332
+ def get_model(self):
333
+ return self.model
334
+
335
+ def download_image(self, url):
336
+ response = requests.get(url)
337
+ return Image.open(BytesIO(response.content)).convert('RGB')
338
+
339
+ def load_image(self, path):
340
+ if path.startswith('http://') or path.startswith('https://'):
341
+ return self.download_image(path)
342
+ return Image.open(path).convert('RGB')
343
+
344
+ def score(self, image_path):
345
+ prompt = "USER: <|image|> <|image|> Compared with the first image, what is your quality rating for second image? \nASSISTANT: The quality of the second image is"
346
+ input_ids = tokenizer_image_token(prompt, self.tokenizer, IMAGE_TOKEN_INDEX, return_tensors='pt').unsqueeze(0).to(self.device)
347
+
348
+ anchor_images = [item['image'] for item in self.anchor_images['train']]
349
+
350
+ probabilities = []
351
+ for index in self.anchor_indices:
352
+ anchor_image = anchor_images[index]
353
+ image = self.load_image(image_path)
354
+ images = [anchor_image, image]
355
+ images = [expand2square(img, tuple(int(x*255) for x in self.image_processor.image_mean)) for img in images]
356
+ image_tensor = self.image_processor.preprocess(images, return_tensors='pt')['pixel_values'].half().to(self.device)
357
+
358
+ with torch.inference_mode():
359
+ output_logits = model(input_ids, images=image_tensor)["logits"][:, -1, self.preferential_ids_]
360
+ output_logits = output_logits.cpu().detach().numpy() / 100
361
+ print(output_logits)
362
+ probabilities.append(np.dot(softmax(output_logits), self.weight_tensor))
363
+ updated_matrix = update_matrix(self.anchor_matrix, np.squeeze(np.array(probabilities)), self.anchor_indices)
364
+ score = optimize_score_map_pytorch_cuda(updated_matrix, seed=0, original_seed=20020, num_iterations=100)
365
+ print(score)
366
+ return score
367
+
368
+ def forward(
369
+ self,
370
+ input_ids: torch.LongTensor = None,
371
+ # modality_indicators: torch.LongTensor = None,
372
+ attention_mask: Optional[torch.Tensor] = None,
373
+ past_key_values: Optional[List[torch.FloatTensor]] = None,
374
+ inputs_embeds: Optional[torch.FloatTensor] = None,
375
+ labels: Optional[torch.LongTensor] = None,
376
+ use_cache: Optional[bool] = None,
377
+ output_attentions: Optional[bool] = None,
378
+ output_hidden_states: Optional[bool] = None,
379
+ images: Optional[torch.FloatTensor] = None,
380
+ return_dict: Optional[bool] = None,
381
+ ) -> Union[Tuple, CausalLMOutputWithPast]:
382
+ output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
383
+ output_hidden_states = (
384
+ output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
385
+ )
386
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
387
+ input_ids, modality_indicators, attention_mask, past_key_values, inputs_embeds, labels = \
388
+ self.prepare_inputs_labels_for_multimodal(input_ids, attention_mask, past_key_values, labels, images)
389
+
390
+ # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
391
+ outputs = self.model(
392
+ input_ids=input_ids,
393
+ modality_indicators=modality_indicators,
394
+ attention_mask=attention_mask,
395
+ past_key_values=past_key_values,
396
+ inputs_embeds=inputs_embeds,
397
+ use_cache=use_cache,
398
+ output_attentions=output_attentions,
399
+ output_hidden_states=output_hidden_states,
400
+ return_dict=return_dict
401
+ )
402
+
403
+ hidden_states = outputs[0]
404
+ logits = self.lm_head(hidden_states)
405
+
406
+ loss = None
407
+ if labels is not None:
408
+ # Shift so that tokens < n predict n
409
+ shift_logits = logits[..., :-1, :].contiguous()
410
+ shift_labels = labels[..., 1:].contiguous()
411
+ # Flatten the tokens
412
+ loss_fct = CrossEntropyLoss()
413
+ shift_logits = shift_logits.view(-1, self.config.vocab_size)
414
+ shift_labels = shift_labels.view(-1)
415
+ # Enable model/pipeline parallelism
416
+ shift_labels = shift_labels.to(shift_logits.device)
417
+ loss = loss_fct(shift_logits, shift_labels)
418
+
419
+ if not return_dict:
420
+ output = (logits,) + outputs[1:]
421
+ return (loss,) + output if loss is not None else output
422
+
423
+ return CausalLMOutputWithPast(
424
+ loss=loss,
425
+ logits=logits,
426
+ past_key_values=outputs.past_key_values,
427
+ hidden_states=outputs.hidden_states,
428
+ attentions=outputs.attentions,
429
+ )
430
+
431
+ def prepare_inputs_for_generation(
432
+ self, input_ids, past_key_values=None, attention_mask=None, inputs_embeds=None, **kwargs
433
+ ):
434
+ if past_key_values:
435
+ input_ids = input_ids[:, -1:]
436
+
437
+ # if `inputs_embeds` are passed, we only want to use them in the 1st generation step
438
+ if inputs_embeds is not None and past_key_values is None:
439
+ model_inputs = {"inputs_embeds": inputs_embeds}
440
+ else:
441
+ model_inputs = {"input_ids": input_ids}
442
+
443
+ model_inputs.update(
444
+ {
445
+ "past_key_values": past_key_values,
446
+ "use_cache": kwargs.get("use_cache"),
447
+ "attention_mask": attention_mask,
448
+ "images": kwargs.get("images", None),
449
+ }
450
+ )
451
+ return model_inputs
452
+
453
+ AutoConfig.register("mplug_owl2", MPLUGOwl2Config)
454
+ AutoModelForCausalLM.register(MPLUGOwl2Config, MPLUGOwl2LlamaForCausalLM)
455
+
456
+ replace_llama_modality_adaptive()
457
+
458
+ if __name__ == "__main__":
459
+ # config = MPLUGOwl2Config.from_pretrained('VQA-CityU/Compare2Score_1')
460
+ from icecream import ic
461
+ # config = MPLUGOwl2Config()
462
+ # model = AutoModelForCausalLM(config)
463
+ model = AutoModelForCausalLM.from_pretrained('VQA-CityU/Compare2Score_1', trust_remote_code=True,
464
+ torch_dtype=torch.float16, device_map="auto")
465
+
466
+ model.score("/home/zhw/IQA/code/NeurIPS24/Q-Align/playground/data/TID2013/distorted_images/i01_01_5.bmp")
467
+ url = "https://raw.githubusercontent.com/Q-Future/Q-Align/main/fig/singapore_flyer.jpg"
468
+ model.score(url)
469
+
visual_encoder.py ADDED
@@ -0,0 +1,922 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import math
2
+ from typing import Any, Optional, Tuple, Union
3
+
4
+ from transformers.modeling_outputs import BaseModelOutput, BaseModelOutputWithPooling, BaseModelOutputWithPastAndCrossAttentions
5
+ from transformers.modeling_utils import PreTrainedModel
6
+ from transformers.pytorch_utils import find_pruneable_heads_and_indices, prune_linear_layer
7
+
8
+ import numpy as np
9
+ import torch
10
+ import torch.nn as nn
11
+ import torch.utils.checkpoint
12
+ from icecream import ic
13
+
14
+ def get_abs_pos(abs_pos, tgt_size):
15
+ # abs_pos: L, C
16
+ # tgt_size: M
17
+ # return: M, C
18
+ src_size = int(math.sqrt(abs_pos.size(0)))
19
+ tgt_size = int(math.sqrt(tgt_size))
20
+ dtype = abs_pos.dtype
21
+
22
+ if src_size != tgt_size:
23
+ return F.interpolate(
24
+ abs_pos.float().reshape(1, src_size, src_size, -1).permute(0, 3, 1, 2),
25
+ size=(tgt_size, tgt_size),
26
+ mode="bicubic",
27
+ align_corners=False,
28
+ ).permute(0, 2, 3, 1).flatten(0, 2).to(dtype=dtype)
29
+ else:
30
+ return abs_pos
31
+
32
+ # https://github.com/facebookresearch/mae/blob/efb2a8062c206524e35e47d04501ed4f544c0ae8/util/pos_embed.py#L20
33
+ def get_2d_sincos_pos_embed(embed_dim, grid_size, cls_token=False):
34
+ """
35
+ grid_size: int of the grid height and width
36
+ return:
37
+ pos_embed: [grid_size*grid_size, embed_dim] or [1+grid_size*grid_size, embed_dim] (w/ or w/o cls_token)
38
+ """
39
+ grid_h = np.arange(grid_size, dtype=np.float32)
40
+ grid_w = np.arange(grid_size, dtype=np.float32)
41
+ grid = np.meshgrid(grid_w, grid_h) # here w goes first
42
+ grid = np.stack(grid, axis=0)
43
+
44
+ grid = grid.reshape([2, 1, grid_size, grid_size])
45
+ pos_embed = get_2d_sincos_pos_embed_from_grid(embed_dim, grid)
46
+ if cls_token:
47
+ pos_embed = np.concatenate([np.zeros([1, embed_dim]), pos_embed], axis=0)
48
+ return pos_embed
49
+
50
+
51
+ def get_2d_sincos_pos_embed_from_grid(embed_dim, grid):
52
+ assert embed_dim % 2 == 0
53
+
54
+ # use half of dimensions to encode grid_h
55
+ emb_h = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[0]) # (H*W, D/2)
56
+ emb_w = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[1]) # (H*W, D/2)
57
+
58
+ emb = np.concatenate([emb_h, emb_w], axis=1) # (H*W, D)
59
+ return emb
60
+
61
+
62
+ def get_1d_sincos_pos_embed_from_grid(embed_dim, pos):
63
+ """
64
+ embed_dim: output dimension for each position
65
+ pos: a list of positions to be encoded: size (M,)
66
+ out: (M, D)
67
+ """
68
+ assert embed_dim % 2 == 0
69
+ omega = np.arange(embed_dim // 2, dtype=np.float32)
70
+ omega /= embed_dim / 2.
71
+ omega = 1. / 10000**omega # (D/2,)
72
+
73
+ pos = pos.reshape(-1) # (M,)
74
+ out = np.einsum('m,d->md', pos, omega) # (M, D/2), outer product
75
+
76
+ emb_sin = np.sin(out) # (M, D/2)
77
+ emb_cos = np.cos(out) # (M, D/2)
78
+
79
+ emb = np.concatenate([emb_sin, emb_cos], axis=1) # (M, D)
80
+ return emb
81
+
82
+
83
+
84
+ class MplugOwlVisionEmbeddings(nn.Module):
85
+ def __init__(self, config):
86
+ super().__init__()
87
+ self.config = config
88
+ self.hidden_size = config.hidden_size
89
+ self.image_size = config.image_size
90
+ self.patch_size = config.patch_size
91
+
92
+ self.cls_token = nn.Parameter(torch.randn(1, 1, self.hidden_size))
93
+
94
+ self.patch_embed = nn.Conv2d(
95
+ in_channels=3,
96
+ out_channels=self.hidden_size,
97
+ kernel_size=self.patch_size,
98
+ stride=self.patch_size,
99
+ bias=False,
100
+ )
101
+
102
+ self.num_patches = (self.image_size // self.patch_size) ** 2
103
+
104
+ self.position_embedding = nn.Parameter(torch.randn(1, self.num_patches + 1, self.hidden_size))
105
+
106
+ self.pre_layernorm = nn.LayerNorm(self.hidden_size, eps=config.layer_norm_eps)
107
+
108
+ def forward(self, pixel_values: torch.FloatTensor) -> torch.Tensor:
109
+ batch_size = pixel_values.size(0)
110
+ image_embeds = self.patch_embed(pixel_values)
111
+ image_embeds = image_embeds.flatten(2).transpose(1, 2)
112
+
113
+ class_embeds = self.cls_token.expand(batch_size, 1, -1).to(image_embeds.dtype)
114
+ embeddings = torch.cat([class_embeds, image_embeds], dim=1)
115
+ embeddings = embeddings + self.position_embedding[:, : embeddings.size(1)].to(image_embeds.dtype)
116
+ embeddings = self.pre_layernorm(embeddings)
117
+ return embeddings
118
+
119
+
120
+
121
+ class MplugOwlVisionAttention(nn.Module):
122
+ """Multi-headed attention from 'Attention Is All You Need' paper"""
123
+
124
+ def __init__(self, config):
125
+ super().__init__()
126
+ self.config = config
127
+ self.hidden_size = config.hidden_size
128
+ self.num_heads = config.num_attention_heads
129
+ self.head_dim = self.hidden_size // self.num_heads
130
+ if self.head_dim * self.num_heads != self.hidden_size:
131
+ raise ValueError(
132
+ f"hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size} and `num_heads`:"
133
+ f" {self.num_heads})."
134
+ )
135
+ self.scale = self.head_dim**-0.5
136
+ self.dropout = nn.Dropout(config.attention_dropout)
137
+
138
+ self.query_key_value = nn.Linear(self.hidden_size, 3 * self.hidden_size)
139
+ self.dense = nn.Linear(self.hidden_size, self.hidden_size)
140
+
141
+ def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int):
142
+ return tensor.view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2).contiguous()
143
+
144
+ def forward(
145
+ self,
146
+ hidden_states: torch.Tensor,
147
+ head_mask: Optional[torch.Tensor] = None,
148
+ output_attentions: Optional[bool] = False,
149
+ ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
150
+ """Input shape: Batch x Time x Channel"""
151
+
152
+ bsz, seq_len, embed_dim = hidden_states.size()
153
+
154
+ mixed_qkv = self.query_key_value(hidden_states)
155
+
156
+ mixed_qkv = mixed_qkv.reshape(bsz, seq_len, self.num_heads, 3, embed_dim // self.num_heads).permute(
157
+ 3, 0, 2, 1, 4
158
+ ) # [3, b, np, sq, hn]
159
+ query_states, key_states, value_states = (
160
+ mixed_qkv[0],
161
+ mixed_qkv[1],
162
+ mixed_qkv[2],
163
+ )
164
+ # if self.config.use_flash_attn and flash_attn_func is not None:
165
+ if False:
166
+ # [b*sq, np, hn]
167
+ query_states = query_states.permute(0, 2, 1, 3).contiguous()
168
+ query_states = query_states.view(query_states.size(0) * query_states.size(1), query_states.size(2), -1)
169
+
170
+ key_states = key_states.permute(0, 2, 1, 3).contiguous()
171
+ key_states = key_states.view(key_states.size(0) * key_states.size(1), key_states.size(2), -1)
172
+
173
+ value_states = value_states.permute(0, 2, 1, 3).contiguous()
174
+ value_states = value_states.view(value_states.size(0) * value_states.size(1), value_states.size(2), -1)
175
+
176
+ cu_seqlens = torch.arange(
177
+ 0, (bsz + 1) * seq_len, step=seq_len, dtype=torch.int32, device=query_states.device
178
+ )
179
+
180
+ context_layer = flash_attn_func(
181
+ query_states,
182
+ key_states,
183
+ value_states,
184
+ cu_seqlens,
185
+ cu_seqlens,
186
+ seq_len,
187
+ seq_len,
188
+ self.dropout if self.training else 0.0,
189
+ softmax_scale=self.scale,
190
+ causal=False,
191
+ return_attn_probs=False,
192
+ )
193
+ # [b*sq, np, hn] => [b, sq, np, hn]
194
+ context_layer = context_layer.view(bsz, seq_len, context_layer.size(1), context_layer.size(2))
195
+ else:
196
+ # Take the dot product between "query" and "key" to get the raw attention scores.
197
+ attention_scores = torch.matmul(query_states, key_states.transpose(-1, -2))
198
+
199
+ attention_scores = attention_scores * self.scale
200
+
201
+ # Normalize the attention scores to probabilities.
202
+ attention_probs = torch.softmax(attention_scores, dim=-1)
203
+
204
+ # This is actually dropping out entire tokens to attend to, which might
205
+ # seem a bit unusual, but is taken from the original Transformer paper.
206
+ attention_probs = self.dropout(attention_probs)
207
+
208
+ # Mask heads if we want to
209
+ if head_mask is not None:
210
+ attention_probs = attention_probs * head_mask
211
+
212
+ context_layer = torch.matmul(attention_probs, value_states).permute(0, 2, 1, 3)
213
+
214
+ new_context_layer_shape = context_layer.size()[:-2] + (self.hidden_size,)
215
+ context_layer = context_layer.reshape(new_context_layer_shape)
216
+
217
+ output = self.dense(context_layer)
218
+
219
+ outputs = (output, attention_probs) if output_attentions else (output, None)
220
+
221
+ return outputs
222
+
223
+
224
+ class QuickGELU(nn.Module):
225
+ def forward(self, x: torch.Tensor):
226
+ return x * torch.sigmoid(1.702 * x)
227
+
228
+
229
+ class MplugOwlMLP(nn.Module):
230
+ def __init__(self, config):
231
+ super().__init__()
232
+ self.config = config
233
+ self.activation_fn = QuickGELU()
234
+ self.fc1 = nn.Linear(config.hidden_size, config.intermediate_size)
235
+ self.fc2 = nn.Linear(config.intermediate_size, config.hidden_size)
236
+
237
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
238
+ hidden_states = self.fc1(hidden_states)
239
+ hidden_states = self.activation_fn(hidden_states)
240
+ hidden_states = self.fc2(hidden_states)
241
+ return hidden_states
242
+
243
+
244
+ class MplugOwlVisionEncoderLayer(nn.Module):
245
+ def __init__(self, config):
246
+ super().__init__()
247
+ self.hidden_size = config.hidden_size
248
+ self.self_attn = MplugOwlVisionAttention(config)
249
+ self.input_layernorm = nn.LayerNorm(self.hidden_size, eps=config.layer_norm_eps)
250
+ self.mlp = MplugOwlMLP(config)
251
+ self.post_attention_layernorm = nn.LayerNorm(self.hidden_size, eps=config.layer_norm_eps)
252
+
253
+ def forward(
254
+ self,
255
+ hidden_states: torch.Tensor,
256
+ attention_mask: torch.Tensor,
257
+ output_attentions: Optional[bool] = False,
258
+ ) -> Tuple[torch.FloatTensor]:
259
+ """
260
+ Args:
261
+ hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`
262
+ attention_mask (`torch.FloatTensor`): attention mask of size
263
+ `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values.
264
+ `(config.encoder_attention_heads,)`.
265
+ output_attentions (`bool`, *optional*):
266
+ Whether or not to return the attentions tensors of all attention layers. See `attentions` under
267
+ returned tensors for more detail.
268
+ """
269
+ residual = hidden_states
270
+
271
+ hidden_states = self.input_layernorm(hidden_states)
272
+ hidden_states, attn_weights = self.self_attn(
273
+ hidden_states=hidden_states,
274
+ head_mask=attention_mask,
275
+ output_attentions=output_attentions,
276
+ )
277
+ hidden_states = hidden_states + residual
278
+ residual = hidden_states
279
+ hidden_states = self.post_attention_layernorm(hidden_states)
280
+ hidden_states = self.mlp(hidden_states)
281
+
282
+ hidden_states = hidden_states + residual
283
+
284
+ outputs = (hidden_states,)
285
+
286
+ if output_attentions:
287
+ outputs += (attn_weights,)
288
+
289
+ return outputs
290
+
291
+
292
+ class MplugOwlVisionEncoder(nn.Module):
293
+ """
294
+ Transformer encoder consisting of `config.num_hidden_layers` self attention layers. Each layer is a
295
+ [`MplugOwlVisionEncoderLayer`].
296
+
297
+ Args:
298
+ config (`MplugOwlVisionConfig`):
299
+ The corresponding vision configuration for the `MplugOwlEncoder`.
300
+ """
301
+
302
+ def __init__(self, config):
303
+ super().__init__()
304
+ self.config = config
305
+ self.layers = nn.ModuleList([MplugOwlVisionEncoderLayer(config) for _ in range(config.num_hidden_layers)])
306
+ self.gradient_checkpointing = True
307
+
308
+ def forward(
309
+ self,
310
+ inputs_embeds,
311
+ attention_mask: Optional[torch.Tensor] = None,
312
+ output_attentions: Optional[bool] = None,
313
+ output_hidden_states: Optional[bool] = None,
314
+ return_dict: Optional[bool] = None,
315
+ ) -> Union[Tuple, BaseModelOutput]:
316
+ r"""
317
+ Args:
318
+ inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
319
+ Embedded representation of the inputs. Should be float, not int tokens.
320
+ attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
321
+ Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
322
+
323
+ - 1 for tokens that are **not masked**,
324
+ - 0 for tokens that are **masked**.
325
+
326
+ [What are attention masks?](../glossary#attention-mask)
327
+ output_attentions (`bool`, *optional*):
328
+ Whether or not to return the attentions tensors of all attention layers. See `attentions` under
329
+ returned tensors for more detail.
330
+ output_hidden_states (`bool`, *optional*):
331
+ Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors
332
+ for more detail.
333
+ return_dict (`bool`, *optional*):
334
+ Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
335
+ """
336
+ output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
337
+ output_hidden_states = (
338
+ output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
339
+ )
340
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
341
+
342
+ encoder_states = () if output_hidden_states else None
343
+ all_attentions = () if output_attentions else None
344
+
345
+ hidden_states = inputs_embeds
346
+ for idx, encoder_layer in enumerate(self.layers):
347
+ if output_hidden_states:
348
+ encoder_states = encoder_states + (hidden_states,)
349
+ if self.gradient_checkpointing and self.training:
350
+
351
+ def create_custom_forward(module):
352
+ def custom_forward(*inputs):
353
+ return module(*inputs, output_attentions)
354
+
355
+ return custom_forward
356
+
357
+ layer_outputs = torch.utils.checkpoint.checkpoint(
358
+ create_custom_forward(encoder_layer),
359
+ hidden_states,
360
+ attention_mask,
361
+ )
362
+ else:
363
+ layer_outputs = encoder_layer(
364
+ hidden_states,
365
+ attention_mask,
366
+ output_attentions=output_attentions,
367
+ )
368
+
369
+ hidden_states = layer_outputs[0]
370
+
371
+ if output_attentions:
372
+ all_attentions = all_attentions + (layer_outputs[1],)
373
+
374
+ if output_hidden_states:
375
+ encoder_states = encoder_states + (hidden_states,)
376
+
377
+ if not return_dict:
378
+ return tuple(v for v in [hidden_states, encoder_states, all_attentions] if v is not None)
379
+ return BaseModelOutput(
380
+ last_hidden_state=hidden_states, hidden_states=encoder_states, attentions=all_attentions
381
+ )
382
+
383
+
384
+ class MplugOwlVisionModel(PreTrainedModel):
385
+ main_input_name = "pixel_values"
386
+ _no_split_modules = ["MplugOwlVisionEncoderLayer"]
387
+
388
+ def __init__(self, config):
389
+ super().__init__(config)
390
+ self.config = config
391
+ self.hidden_size = config.hidden_size
392
+
393
+ self.embeddings = MplugOwlVisionEmbeddings(config)
394
+ self.encoder = MplugOwlVisionEncoder(config)
395
+ self.post_layernorm = nn.LayerNorm(self.hidden_size, eps=config.layer_norm_eps)
396
+
397
+ self.post_init()
398
+
399
+
400
+ def forward(
401
+ self,
402
+ pixel_values: Optional[torch.FloatTensor] = None,
403
+ output_attentions: Optional[bool] = None,
404
+ output_hidden_states: Optional[bool] = None,
405
+ return_dict: Optional[bool] = None,
406
+ ) -> Union[Tuple, BaseModelOutputWithPooling]:
407
+ r"""
408
+ Returns:
409
+
410
+ """
411
+ output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
412
+ output_hidden_states = (
413
+ output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
414
+ )
415
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
416
+
417
+ if pixel_values is None:
418
+ raise ValueError("You have to specify pixel_values")
419
+
420
+ hidden_states = self.embeddings(pixel_values)
421
+
422
+ encoder_outputs = self.encoder(
423
+ inputs_embeds=hidden_states,
424
+ output_attentions=output_attentions,
425
+ output_hidden_states=output_hidden_states,
426
+ return_dict=return_dict,
427
+ )
428
+
429
+ last_hidden_state = encoder_outputs[0]
430
+ last_hidden_state = self.post_layernorm(last_hidden_state)
431
+
432
+ pooled_output = last_hidden_state[:, 0, :]
433
+ pooled_output = self.post_layernorm(pooled_output)
434
+
435
+ if not return_dict:
436
+ return (last_hidden_state, pooled_output) + encoder_outputs[1:]
437
+
438
+ return BaseModelOutputWithPooling(
439
+ last_hidden_state=last_hidden_state,
440
+ pooler_output=pooled_output,
441
+ hidden_states=encoder_outputs.hidden_states,
442
+ attentions=encoder_outputs.attentions,
443
+ )
444
+
445
+ def get_input_embeddings(self):
446
+ return self.embeddings
447
+
448
+
449
+ class MplugOwlVisualAbstractorMLP(nn.Module):
450
+ def __init__(self, config):
451
+ super().__init__()
452
+ self.config = config
453
+ in_features = config.hidden_size
454
+ self.act = nn.SiLU()
455
+
456
+ self.w1 = nn.Linear(in_features, config.intermediate_size)
457
+ self.w2 = nn.Linear(config.intermediate_size, in_features)
458
+ self.w3 = nn.Linear(in_features, config.intermediate_size)
459
+ self.ffn_ln = nn.LayerNorm(config.intermediate_size, eps=config.layer_norm_eps)
460
+
461
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
462
+ hidden_states = self.act(self.w1(hidden_states)) * self.w3(hidden_states)
463
+ hidden_states = self.ffn_ln(hidden_states)
464
+ hidden_states = self.w2(hidden_states)
465
+ return hidden_states
466
+
467
+
468
+ class MplugOwlVisualAbstractorMultiHeadAttention(nn.Module):
469
+ def __init__(self, config):
470
+ super().__init__()
471
+ self.config = config
472
+ if config.hidden_size % config.num_attention_heads != 0:
473
+ raise ValueError(
474
+ "The hidden size (%d) is not a multiple of the number of attention heads (%d)"
475
+ % (config.hidden_size, config.num_attention_heads)
476
+ )
477
+
478
+ self.num_attention_heads = config.num_attention_heads
479
+ self.attention_head_size = int(config.hidden_size / config.num_attention_heads)
480
+ self.all_head_size = self.num_attention_heads * self.attention_head_size
481
+
482
+ self.query = nn.Linear(config.hidden_size, self.all_head_size)
483
+ self.key = nn.Linear(config.encoder_hidden_size, self.all_head_size)
484
+ self.value = nn.Linear(config.encoder_hidden_size, self.all_head_size)
485
+
486
+ self.dropout = nn.Dropout(config.attention_probs_dropout_prob)
487
+ self.save_attention = False
488
+
489
+ # self.q_pos_embed = nn.Parameter(
490
+ # torch.from_numpy(get_1d_sincos_pos_embed_from_grid(config.hidden_size, np.arange(config.num_learnable_queries, dtype=np.float32))).float()
491
+ # ).requires_grad_(False)
492
+ # grids = config.grid_size
493
+ # self.k_pos_embed = nn.Parameter(
494
+ # torch.from_numpy(get_2d_sincos_pos_embed(config.hidden_size, grids, cls_token=True)).float()
495
+ # ).requires_grad_(False)
496
+ grids = config.grid_size
497
+ self.register_buffer(
498
+ 'q_pos_embed',
499
+ torch.from_numpy(get_1d_sincos_pos_embed_from_grid(config.hidden_size, np.arange(config.num_learnable_queries, dtype=np.float32))).float()
500
+ )
501
+ self.register_buffer(
502
+ 'k_pos_embed',
503
+ torch.from_numpy(get_2d_sincos_pos_embed(config.hidden_size, grids, cls_token=True)).float()
504
+ )
505
+
506
+
507
+ def save_attn_gradients(self, attn_gradients):
508
+ self.attn_gradients = attn_gradients
509
+
510
+ def get_attn_gradients(self):
511
+ return self.attn_gradients
512
+
513
+ def save_attention_map(self, attention_map):
514
+ self.attention_map = attention_map
515
+
516
+ def get_attention_map(self):
517
+ return self.attention_map
518
+
519
+ def transpose_for_scores(self, x):
520
+ new_x_shape = x.size()[:-1] + (self.num_attention_heads, self.attention_head_size)
521
+ x = x.view(*new_x_shape)
522
+ return x.permute(0, 2, 1, 3)
523
+
524
+ def forward(
525
+ self,
526
+ hidden_states,
527
+ attention_mask=None,
528
+ head_mask=None,
529
+ encoder_hidden_states=None,
530
+ encoder_attention_mask=None,
531
+ past_key_value=None,
532
+ output_attentions=False,
533
+ ):
534
+ # If this is instantiated as a cross-attention module, the keys
535
+ # and values come from an encoder; the attention mask needs to be
536
+ # such that the encoder's padding tokens are not attended to.
537
+
538
+ qk_pos_embed = torch.cat([self.q_pos_embed, self.k_pos_embed], dim = 0).unsqueeze(0).to(dtype=hidden_states.dtype)
539
+
540
+ key_layer = self.transpose_for_scores(self.key(encoder_hidden_states + qk_pos_embed))
541
+ value_layer = self.transpose_for_scores(self.value(encoder_hidden_states))
542
+ attention_mask = encoder_attention_mask
543
+
544
+ mixed_query_layer = self.query(hidden_states + self.q_pos_embed.unsqueeze(0).to(dtype=hidden_states.dtype))
545
+
546
+ query_layer = self.transpose_for_scores(mixed_query_layer)
547
+
548
+ past_key_value = (key_layer, value_layer)
549
+
550
+ # Take the dot product between "query" and "key" to get the raw attention scores.
551
+ attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2))
552
+
553
+ attention_scores = attention_scores / math.sqrt(self.attention_head_size)
554
+
555
+ if attention_mask is not None:
556
+ # Apply the attention mask is (precomputed for all layers in BertModel forward() function)
557
+ attention_scores = attention_scores + attention_mask
558
+
559
+ # Normalize the attention scores to probabilities.
560
+ attention_probs = nn.Softmax(dim=-1)(attention_scores)
561
+
562
+ if self.save_attention:
563
+ self.save_attention_map(attention_probs)
564
+ attention_probs.register_hook(self.save_attn_gradients)
565
+
566
+ # This is actually dropping out entire tokens to attend to, which might
567
+ # seem a bit unusual, but is taken from the original Transformer paper.
568
+ attention_probs_dropped = self.dropout(attention_probs)
569
+
570
+ # Mask heads if we want to
571
+ if head_mask is not None:
572
+ attention_probs_dropped = attention_probs_dropped * head_mask
573
+
574
+ context_layer = torch.matmul(attention_probs_dropped, value_layer)
575
+
576
+ context_layer = context_layer.permute(0, 2, 1, 3).contiguous()
577
+ new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,)
578
+ context_layer = context_layer.view(*new_context_layer_shape)
579
+
580
+ outputs = (context_layer, attention_probs) if output_attentions else (context_layer,)
581
+
582
+ outputs = outputs + (past_key_value,)
583
+ return outputs
584
+
585
+
586
+ class MplugOwlVisualAbstractorCrossOutput(nn.Module):
587
+ def __init__(self, config):
588
+ super().__init__()
589
+ dim = config.hidden_size
590
+ self.out_proj = nn.Linear(dim, dim, bias=True)
591
+ self.norm2 = nn.LayerNorm(dim)
592
+ self.mlp = MplugOwlVisualAbstractorMLP(config)
593
+
594
+ def forward(self, hidden_states: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor:
595
+ input_tensor = input_tensor + self.out_proj(hidden_states)
596
+ input_tensor = input_tensor + self.mlp(self.norm2(input_tensor))
597
+ return input_tensor
598
+
599
+
600
+ class MplugOwlVisualAbstractorAttention(nn.Module):
601
+ def __init__(self, config):
602
+ super().__init__()
603
+ self.attention = MplugOwlVisualAbstractorMultiHeadAttention(config)
604
+ self.output = MplugOwlVisualAbstractorCrossOutput(config)
605
+ self.pruned_heads = set()
606
+ self.norm1 = nn.LayerNorm(config.hidden_size)
607
+ self.normk = nn.LayerNorm(config.hidden_size)
608
+
609
+ def prune_heads(self, heads):
610
+ if len(heads) == 0:
611
+ return
612
+ heads, index = find_pruneable_heads_and_indices(
613
+ heads, self.attention.num_attention_heads, self.attention.attention_head_size, self.pruned_heads
614
+ )
615
+
616
+ # Prune linear layers
617
+ self.attention.query = prune_linear_layer(self.attention.query, index)
618
+ self.attention.key = prune_linear_layer(self.attention.key, index)
619
+ self.attention.value = prune_linear_layer(self.attention.value, index)
620
+ self.output.dense = prune_linear_layer(self.output.out_proj, index, dim=1)
621
+
622
+ # Update hyper params and store pruned heads
623
+ self.attention.num_attention_heads = self.attention.num_attention_heads - len(heads)
624
+ self.attention.all_head_size = self.attention.attention_head_size * self.attention.num_attention_heads
625
+ self.pruned_heads = self.pruned_heads.union(heads)
626
+
627
+ def forward(
628
+ self,
629
+ hidden_states: torch.Tensor,
630
+ attention_mask: Optional[torch.FloatTensor] = None,
631
+ head_mask: Optional[torch.FloatTensor] = None,
632
+ encoder_hidden_states: Optional[torch.FloatTensor] = None,
633
+ encoder_attention_mask: Optional[torch.FloatTensor] = None,
634
+ past_key_value: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
635
+ output_attentions: Optional[bool] = False,
636
+ ) -> Tuple[torch.Tensor]:
637
+ # HACK we apply norm on q and k
638
+ hidden_states = self.norm1(hidden_states)
639
+ encoder_hidden_states = self.normk(encoder_hidden_states)
640
+ encoder_hidden_states = torch.cat([hidden_states, encoder_hidden_states], dim=1)
641
+ encoder_attention_mask = torch.cat([attention_mask, encoder_attention_mask], dim=-1)
642
+ self_outputs = self.attention(
643
+ hidden_states,
644
+ attention_mask,
645
+ head_mask,
646
+ encoder_hidden_states,
647
+ encoder_attention_mask,
648
+ past_key_value,
649
+ output_attentions,
650
+ )
651
+ attention_output = self.output(self_outputs[0], hidden_states)
652
+ # add attentions if we output them
653
+ outputs = (attention_output,) + self_outputs[1:]
654
+ return outputs
655
+
656
+
657
+ class MplugOwlVisualAbstractorLayer(nn.Module):
658
+ def __init__(self, config, layer_idx):
659
+ super().__init__()
660
+ self.chunk_size_feed_forward = config.chunk_size_feed_forward
661
+ self.seq_len_dim = 1
662
+
663
+ self.layer_idx = layer_idx
664
+
665
+ self.crossattention = MplugOwlVisualAbstractorAttention(config)
666
+ self.has_cross_attention = True
667
+
668
+ def forward(
669
+ self,
670
+ hidden_states,
671
+ attention_mask=None,
672
+ head_mask=None,
673
+ encoder_hidden_states=None,
674
+ encoder_attention_mask=None,
675
+ output_attentions=False,
676
+ ):
677
+ if encoder_hidden_states is None:
678
+ raise ValueError("encoder_hidden_states must be given for cross-attention layers")
679
+ cross_attention_outputs = self.crossattention(
680
+ hidden_states,
681
+ attention_mask,
682
+ head_mask,
683
+ encoder_hidden_states,
684
+ encoder_attention_mask,
685
+ output_attentions=output_attentions,
686
+ )
687
+ query_attention_output = cross_attention_outputs[0]
688
+
689
+ outputs = (query_attention_output,)
690
+ return outputs
691
+
692
+
693
+ class MplugOwlVisualAbstractorEncoder(nn.Module):
694
+ def __init__(self, config):
695
+ super().__init__()
696
+ self.config = config
697
+ self.layers = nn.ModuleList(
698
+ [MplugOwlVisualAbstractorLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
699
+ )
700
+ self.gradient_checkpointing = True
701
+
702
+ def forward(
703
+ self,
704
+ hidden_states,
705
+ attention_mask=None,
706
+ head_mask=None,
707
+ encoder_hidden_states=None,
708
+ encoder_attention_mask=None,
709
+ past_key_values=None,
710
+ output_attentions=False,
711
+ output_hidden_states=False,
712
+ return_dict=True,
713
+ ):
714
+ all_hidden_states = () if output_hidden_states else None
715
+
716
+ for i in range(self.config.num_hidden_layers):
717
+ layer_module = self.layers[i]
718
+ if output_hidden_states:
719
+ all_hidden_states = all_hidden_states + (hidden_states,)
720
+
721
+ layer_head_mask = head_mask[i] if head_mask is not None else None
722
+ past_key_value = past_key_values[i] if past_key_values is not None else None
723
+
724
+ if getattr(self.config, "gradient_checkpointing", False) and self.training:
725
+
726
+ def create_custom_forward(module):
727
+ def custom_forward(*inputs):
728
+ return module(*inputs, past_key_value, output_attentions)
729
+
730
+ return custom_forward
731
+
732
+ layer_outputs = torch.utils.checkpoint.checkpoint(
733
+ create_custom_forward(layer_module),
734
+ hidden_states,
735
+ attention_mask,
736
+ layer_head_mask,
737
+ encoder_hidden_states,
738
+ encoder_attention_mask,
739
+ )
740
+ else:
741
+ layer_outputs = layer_module(
742
+ hidden_states,
743
+ attention_mask,
744
+ layer_head_mask,
745
+ encoder_hidden_states,
746
+ encoder_attention_mask,
747
+ output_attentions,
748
+ )
749
+
750
+ hidden_states = layer_outputs[0]
751
+
752
+ return BaseModelOutput(
753
+ last_hidden_state=hidden_states,
754
+ )
755
+
756
+
757
+ class MplugOwlVisualAbstractorModel(PreTrainedModel):
758
+ _no_split_modules = ["MplugOwlVisualAbstractorLayer"]
759
+ def __init__(self, config, language_hidden_size):
760
+ super().__init__(config)
761
+ self.config = config
762
+
763
+ self.encoder = MplugOwlVisualAbstractorEncoder(config)
764
+ self.visual_fc = torch.nn.Linear(config.hidden_size, language_hidden_size)
765
+ self.query_embeds = torch.nn.Parameter(torch.randn(1, config.num_learnable_queries, config.hidden_size))
766
+ self.vit_eos = torch.nn.Parameter(torch.randn(1, 1, language_hidden_size))
767
+
768
+ self.post_init()
769
+
770
+ def _prune_heads(self, heads_to_prune):
771
+ """
772
+ Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} See base
773
+ class PreTrainedModel
774
+ """
775
+ for layer, heads in heads_to_prune.items():
776
+ self.encoder.layer[layer].attention.prune_heads(heads)
777
+
778
+ def get_extended_attention_mask(
779
+ self,
780
+ attention_mask: torch.Tensor,
781
+ input_shape: Tuple[int],
782
+ device: torch.device,
783
+ ) -> torch.Tensor:
784
+ """
785
+ Makes broadcastable attention and causal masks so that future and masked tokens are ignored.
786
+
787
+ Arguments:
788
+ attention_mask (`torch.Tensor`):
789
+ Mask with ones indicating tokens to attend to, zeros for tokens to ignore.
790
+ input_shape (`Tuple[int]`):
791
+ The shape of the input to the model.
792
+ device: (`torch.device`):
793
+ The device of the input to the model.
794
+
795
+ Returns:
796
+ `torch.Tensor` The extended attention mask, with a the same dtype as `attention_mask.dtype`.
797
+ """
798
+ # We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length]
799
+ # ourselves in which case we just need to make it broadcastable to all heads.
800
+ if attention_mask.dim() == 3:
801
+ extended_attention_mask = attention_mask[:, None, :, :]
802
+ elif attention_mask.dim() == 2:
803
+ # Provided a padding mask of dimensions [batch_size, seq_length]
804
+ # - the model is an encoder, so make the mask broadcastable to [batch_size, num_heads, seq_length, seq_length]
805
+ extended_attention_mask = attention_mask[:, None, None, :]
806
+ else:
807
+ raise ValueError(
808
+ "Wrong shape for input_ids (shape {}) or attention_mask (shape {})".format(
809
+ input_shape, attention_mask.shape
810
+ )
811
+ )
812
+
813
+ # Since attention_mask is 1.0 for positions we want to attend and 0.0 for
814
+ # masked positions, this operation will create a tensor which is 0.0 for
815
+ # positions we want to attend and -10000.0 for masked positions.
816
+ # Since we are adding it to the raw scores before the softmax, this is
817
+ # effectively the same as removing these entirely.
818
+ extended_attention_mask = extended_attention_mask.to(dtype=self.dtype) # fp16 compatibility
819
+ extended_attention_mask = (1.0 - extended_attention_mask) * -10000.0
820
+ return extended_attention_mask
821
+
822
+ def forward(
823
+ self,
824
+ attention_mask=None,
825
+ head_mask=None,
826
+ encoder_hidden_states=None,
827
+ encoder_attention_mask=None,
828
+ past_key_values=None,
829
+ output_attentions=None,
830
+ output_hidden_states=None,
831
+ return_dict=None,
832
+ ):
833
+ r"""
834
+ encoder_hidden_states (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, `optional`):
835
+ Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention if
836
+ the model is configured as a decoder.
837
+ encoder_attention_mask (`torch.FloatTensor` of shape `(batch_size, sequence_length)`, `optional`):
838
+ Mask to avoid performing attention on the padding token indices of the encoder input. This mask is used in
839
+ the cross-attention if the model is configured as a decoder. Mask values selected in `[0, 1]`:
840
+ - 1 for tokens that are **not masked**,
841
+ - 0 for tokens that are **masked**.
842
+ past_key_values (`tuple(tuple(torch.FloatTensor))` of length `config.n_layers` with each tuple having 4 tensors of:
843
+ shape `(batch_size, num_heads, sequence_length - 1, embed_size_per_head)`): Contains precomputed key and
844
+ value hidden states of the attention blocks. Can be used to speed up decoding. If `past_key_values` are
845
+ used, the user can optionally input only the last `decoder_input_ids` (those that don't have their past key
846
+ value states given to this model) of shape `(batch_size, 1)` instead of all `decoder_input_ids` of shape
847
+ `(batch_size, sequence_length)`.
848
+ """
849
+ output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
850
+ output_hidden_states = (
851
+ output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
852
+ )
853
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
854
+
855
+ query_embeds = self.query_embeds.repeat(encoder_hidden_states.shape[0], 1, 1)
856
+ embedding_output = query_embeds
857
+ input_shape = embedding_output.size()[:-1]
858
+ batch_size, seq_length = input_shape
859
+ device = embedding_output.device
860
+
861
+ # We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length]
862
+ # ourselves in which case we just need to make it broadcastable to all heads.
863
+ if attention_mask is None:
864
+ attention_mask = torch.ones(
865
+ (query_embeds.shape[0], query_embeds.shape[1]), dtype=torch.long, device=query_embeds.device
866
+ )
867
+ extended_attention_mask = self.get_extended_attention_mask(attention_mask, input_shape, device)
868
+
869
+ # If a 2D or 3D attention mask is provided for the cross-attention
870
+ # we need to make broadcastable to [batch_size, num_heads, seq_length, seq_length]
871
+ if encoder_hidden_states is not None:
872
+ if type(encoder_hidden_states) == list:
873
+ encoder_batch_size, encoder_sequence_length, _ = encoder_hidden_states[0].size()
874
+ else:
875
+ (
876
+ encoder_batch_size,
877
+ encoder_sequence_length,
878
+ _,
879
+ ) = encoder_hidden_states.size()
880
+ encoder_hidden_shape = (encoder_batch_size, encoder_sequence_length)
881
+
882
+ if type(encoder_attention_mask) == list:
883
+ encoder_extended_attention_mask = [self.invert_attention_mask(mask) for mask in encoder_attention_mask]
884
+ elif encoder_attention_mask is None:
885
+ encoder_attention_mask = torch.ones(encoder_hidden_shape, device=device)
886
+ encoder_extended_attention_mask = self.invert_attention_mask(encoder_attention_mask)
887
+ else:
888
+ encoder_extended_attention_mask = self.invert_attention_mask(encoder_attention_mask)
889
+ else:
890
+ encoder_extended_attention_mask = None
891
+
892
+ # Prepare head mask if needed
893
+ # 1.0 in head_mask indicate we keep the head
894
+ # attention_probs has shape bsz x n_heads x N x N
895
+ # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads]
896
+ # and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length]
897
+ head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers)
898
+
899
+ encoder_outputs = self.encoder(
900
+ embedding_output,
901
+ attention_mask=extended_attention_mask,
902
+ head_mask=head_mask,
903
+ encoder_hidden_states=encoder_hidden_states,
904
+ encoder_attention_mask=encoder_extended_attention_mask,
905
+ past_key_values=past_key_values,
906
+ output_attentions=output_attentions,
907
+ output_hidden_states=output_hidden_states,
908
+ return_dict=return_dict,
909
+ )
910
+ sequence_output = encoder_outputs[0]
911
+ pooled_output = sequence_output[:, 0, :]
912
+
913
+ sequence_output = self.visual_fc(sequence_output)
914
+ sequence_output = torch.cat([sequence_output, self.vit_eos.repeat(sequence_output.shape[0], 1, 1)], dim=1)
915
+
916
+ return BaseModelOutputWithPooling(
917
+ last_hidden_state=sequence_output,
918
+ pooler_output=pooled_output,
919
+ hidden_states=encoder_outputs.hidden_states,
920
+ )
921
+
922
+