Aluode/PerceptionLabPortable
0
1# coding=utf-82# Copyright 2024 Microsoft Research and HuggingFace Inc. team.3#4# Licensed under the Apache License, Version 2.0 (the "License");5# you may not use this file except in compliance with the License.6# You may obtain a copy of the License at7#8# http://www.apache.org/licenses/LICENSE-2.09#10# Unless required by applicable law or agreed to in writing, software11# distributed under the License is distributed on an "AS IS" BASIS,12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.13# See the License for the specific language governing permissions and14# limitations under the License.15"""PyTorch UDOP model."""16 17import collections18import logging19import math20import random21from abc import ABC, abstractmethod22from collections.abc import Sequence23from copy import deepcopy24from dataclasses import dataclass25from typing import Any, Optional, Union26 27import torch28from torch import Tensor, nn29from torch.nn import CrossEntropyLoss30 31from transformers import UdopConfig32from transformers.modeling_outputs import (33 Seq2SeqLMOutput,34 Seq2SeqModelOutput,35)36 37from ...activations import ACT2FN38from ...cache_utils import Cache, DynamicCache, EncoderDecoderCache39from ...generation import GenerationMixin40from ...modeling_attn_mask_utils import AttentionMaskConverter41from ...modeling_layers import GradientCheckpointingLayer42from ...modeling_utils import PreTrainedModel43from ...pytorch_utils import find_pruneable_heads_and_indices, prune_linear_layer44from ...utils import (45 ModelOutput,46 auto_docstring,47 is_torch_flex_attn_available,48 is_torchdynamo_compiling,49)50from ...utils.deprecation import deprecate_kwarg51 52 53if is_torch_flex_attn_available():54 from torch.nn.attention.flex_attention import BlockMask55 56 from ...integrations.flex_attention import make_flex_block_causal_mask57 58 59logger = logging.getLogger(__name__)60 61 62@dataclass63@auto_docstring(64 custom_intro="""65 Class for the model's outputs that may also contain a past key/values (to speed up sequential decoding). Includes66 an additional attention mask.67 """68)69class BaseModelOutputWithAttentionMask(ModelOutput):70 r"""71 last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):72 Sequence of hidden-states at the output of the last layer of the model. If `past_key_values` is used only73 the last hidden-state of the sequences of shape `(batch_size, 1, hidden_size)` is output.74 attention_mask (`torch.FloatTensor` of shape `(batch_size, sequence_length)`, *optional*):75 Attention mask used in the model's forward pass to avoid performing attention on padding token indices.76 Mask values selected in `[0, 1]`:77 - 1 for tokens that are **not masked**,78 - 0 for tokens that are **masked**.79 past_key_values (`Cache`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):80 It is a [`~cache_utils.Cache`] instance. For more details, see our [kv cache guide](https://huggingface.co/docs/transformers/en/kv_cache).81 82 Contains pre-computed hidden-states (key and values in the83 self-attention blocks and optionally if `config.is_encoder_decoder=True` in the cross-attention blocks)84 that can be used (see `past_key_values` input) to speed up sequential decoding.85 hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):86 Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, +87 one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`. Hidden-states of88 the model at the output of each layer plus the optional initial embedding outputs.89 attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):90 Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,91 sequence_length)`. Attentions weights after the attention softmax, used to compute the weighted average in92 the self-attention heads.93 cross_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` and `config.add_cross_attention=True` is passed or when `config.output_attentions=True`):94 Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,95 sequence_length)`. Attentions weights of the decoder's cross-attention layer, after the attention softmax,96 used to compute the weighted average in the cross-attention heads.97 """98 99 last_hidden_state: Optional[torch.FloatTensor] = None100 attention_mask: Optional[torch.FloatTensor] = None101 past_key_values: Optional[Cache] = None102 hidden_states: Optional[tuple[torch.FloatTensor]] = None103 attentions: Optional[tuple[torch.FloatTensor]] = None104 cross_attentions: Optional[tuple[torch.FloatTensor]] = None105 106 107def get_visual_bbox(image_size=224, patch_size=16):108 image_feature_pool_shape = [image_size // patch_size, image_size // patch_size]109 visual_bbox_x = torch.arange(0, 1.0 * (image_feature_pool_shape[1] + 1), 1.0)110 visual_bbox_x /= image_feature_pool_shape[1]111 112 visual_bbox_y = torch.arange(0, 1.0 * (image_feature_pool_shape[0] + 1), 1.0)113 visual_bbox_y /= image_feature_pool_shape[0]114 115 visual_bbox_input = torch.stack(116 [117 visual_bbox_x[:-1].repeat(image_feature_pool_shape[0], 1),118 visual_bbox_y[:-1].repeat(image_feature_pool_shape[1], 1).transpose(0, 1),119 visual_bbox_x[1:].repeat(image_feature_pool_shape[0], 1),120 visual_bbox_y[1:].repeat(image_feature_pool_shape[1], 1).transpose(0, 1),121 ],122 dim=-1,123 )124 125 visual_bbox_input = visual_bbox_input.view(-1, 4)126 127 return visual_bbox_input128 129 130def pad_sequence(seq, target_len, pad_value=0):131 if isinstance(seq, torch.Tensor):132 n = seq.shape[0]133 else:134 n = len(seq)135 seq = torch.tensor(seq)136 m = target_len - n137 if m > 0:138 ret = torch.stack([pad_value] * m).to(seq)139 seq = torch.cat([seq, ret], dim=0)140 return seq[:target_len]141 142 143def combine_image_text_embeddings(144 image_embeddings,145 inputs_embeds,146 bbox,147 visual_bbox,148 attention_mask=None,149 num_patches=14,150 max_len=0,151 image_size=224,152 patch_size=16,153):154 """155 Combine the image and text embeddings for the input to the encoder/decoder of UDOP.156 157 First, the image embeddings are created by checking for each visual patch if it is inside the bounding box of a158 token. If it is, the visual patch is combined with the token embedding. Then, the visual bounding boxes are combined159 with the text bounding boxes. Finally, the visual bounding boxes are combined with the text attention mask.160 """161 162 sequence_length = num_patches163 ocr_points_x = torch.clip(164 torch.floor((bbox[:, :, 0] + bbox[:, :, 2]) / 2.0 * sequence_length).long(), 0, sequence_length - 1165 )166 ocr_points_y = (167 torch.clip(torch.floor((bbox[:, :, 1] + bbox[:, :, 3]) / 2.0 * sequence_length).long(), 0, sequence_length - 1)168 * sequence_length169 )170 ocr_points = ocr_points_x + ocr_points_y171 # make sure bounding boxes are of type float to calculate means172 bbox = bbox.to(torch.float64)173 target_seg = (bbox.mean(-1) == 0.0) | (bbox.mean(-1) == 1.0)174 repeated_vision_embeds = torch.gather(175 image_embeddings, 1, ocr_points.unsqueeze(-1).repeat(1, 1, image_embeddings.size(-1))176 )177 repeated_vision_embeds[target_seg] = 0.0178 inputs_embeds += repeated_vision_embeds179 180 patch_inds = torch.full_like(image_embeddings[:, :, 0], True).bool()181 ind = torch.cat(182 [183 torch.arange(len(ocr_points))[:, None].repeat(1, ocr_points.size(-1))[:, :, None].to(ocr_points),184 ocr_points[:, :, None],185 ],186 dim=-1,187 )188 ind = ind.flatten(0, 1)189 rows, cols = zip(*ind)190 patch_inds[rows, cols] = False191 192 input_vision_patches = [image_embeddings[i][patch_inds[i]] for i in range(len(patch_inds))]193 194 if visual_bbox is None:195 visual_bbox = get_visual_bbox(image_size=image_size, patch_size=patch_size)196 visual_bbox = visual_bbox.unsqueeze(0).repeat(image_embeddings.size(0), 1, 1)197 visual_bbox = visual_bbox.to(image_embeddings.device)198 199 visual_bbox = [visual_bbox[i][patch_inds[i]] for i in range(len(patch_inds))]200 if attention_mask is not None:201 visual_attention_mask = [torch.tensor([1] * len(item)).to(attention_mask) for item in visual_bbox]202 203 if max_len == 0:204 max_len = image_embeddings.size(1)205 else:206 max_len = max_len - inputs_embeds.size(1)207 inputs_vision_patches = torch.stack(208 [pad_sequence(item, max_len, torch.zeros_like(image_embeddings[0, 0])) for item in input_vision_patches]209 )210 visual_bbox = torch.stack([pad_sequence(item, max_len, torch.zeros_like(bbox[0, 0])) for item in visual_bbox])211 if attention_mask is not None:212 visual_attention_mask = torch.stack(213 [pad_sequence(item, max_len, torch.zeros_like(attention_mask[0, 0])) for item in visual_attention_mask]214 )215 216 inputs_embeds = torch.cat([inputs_embeds, inputs_vision_patches], 1)217 bbox = torch.cat([bbox, visual_bbox], 1)218 if attention_mask is not None:219 attention_mask = torch.cat([attention_mask, visual_attention_mask], 1)220 return inputs_embeds, bbox, attention_mask221 222 223class UdopPatchEmbeddings(nn.Module):224 """2D Image to Patch Embeddings"""225 226 def __init__(self, config):227 super().__init__()228 image_size, patch_size = config.image_size, config.patch_size229 num_channels, hidden_size = config.num_channels, config.hidden_size230 231 image_size = image_size if isinstance(image_size, collections.abc.Iterable) else (image_size, image_size)232 patch_size = patch_size if isinstance(patch_size, collections.abc.Iterable) else (patch_size, patch_size)233 num_patches = (image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0])234 self.image_size = image_size235 self.patch_size = patch_size236 self.num_channels = num_channels237 self.num_patches = num_patches238 239 self.proj = nn.Conv2d(num_channels, hidden_size, kernel_size=patch_size, stride=patch_size)240 241 def forward(self, pixel_values):242 batch_size, num_channels, height, width = pixel_values.shape243 if height != self.image_size[0] or width != self.image_size[1]:244 raise ValueError(245 f"Input image size ({height}*{width}) doesn't match model ({self.image_size[0]}*{self.image_size[1]})."246 )247 embeddings = self.proj(pixel_values)248 embeddings = embeddings.flatten(2).transpose(1, 2)249 return embeddings250 251 252@auto_docstring253class UdopPreTrainedModel(PreTrainedModel):254 config: UdopConfig255 base_model_prefix = "transformer"256 supports_gradient_checkpointing = True257 258 _can_compile_fullgraph = False259 _keep_in_fp32_modules = ["wo"]260 261 def _init_weights(self, module):262 """Initialize the weights"""263 factor = self.config.initializer_factor # Used for testing weights initialization264 if isinstance(module, UdopLayerNorm):265 module.weight.data.fill_(factor * 1.0)266 elif isinstance(module, nn.Embedding):267 module.weight.data.normal_(mean=0.0, std=factor)268 if module.padding_idx is not None:269 module.weight.data[module.padding_idx].zero_()270 elif isinstance(module, nn.Conv2d):271 # Upcast the input in `fp32` and cast it back to desired `dtype` to avoid272 # `trunc_normal_cpu` not implemented in `half` issues273 module.weight.data = nn.init.trunc_normal_(module.weight.data.to(torch.float32), mean=0.0, std=factor).to(274 module.weight.dtype275 )276 if module.bias is not None:277 module.bias.data.zero_()278 elif isinstance(module, RelativePositionBiasBase):279 factor = self.config.initializer_factor280 d_model = self.config.d_model281 module.relative_attention_bias.weight.data.normal_(mean=0.0, std=factor * ((d_model) ** -0.5))282 elif isinstance(module, UdopModel):283 # Mesh TensorFlow embeddings initialization284 # See https://github.com/tensorflow/mesh/blob/fa19d69eafc9a482aff0b59ddd96b025c0cb207d/mesh_tensorflow/layers.py#L1624285 module.shared.weight.data.normal_(mean=0.0, std=factor * 1.0)286 elif isinstance(module, UdopForConditionalGeneration):287 if hasattr(module, "lm_head") and not self.config.tie_word_embeddings:288 module.lm_head.weight.data.normal_(mean=0.0, std=factor * 1.0)289 elif isinstance(module, UdopDenseActDense):290 # Mesh TensorFlow FF initialization291 # See https://github.com/tensorflow/mesh/blob/master/mesh_tensorflow/transformer/transformer_layers.py#L56292 # and https://github.com/tensorflow/mesh/blob/fa19d69eafc9a482aff0b59ddd96b025c0cb207d/mesh_tensorflow/layers.py#L89293 module.wi.weight.data.normal_(mean=0.0, std=factor * ((self.config.d_model) ** -0.5))294 if hasattr(module.wi, "bias") and module.wi.bias is not None:295 module.wi.bias.data.zero_()296 module.wo.weight.data.normal_(mean=0.0, std=factor * ((self.config.d_ff) ** -0.5))297 if hasattr(module.wo, "bias") and module.wo.bias is not None:298 module.wo.bias.data.zero_()299 elif isinstance(module, UdopDenseGatedActDense):300 module.wi_0.weight.data.normal_(mean=0.0, std=factor * ((self.config.d_model) ** -0.5))301 if hasattr(module.wi_0, "bias") and module.wi_0.bias is not None:302 module.wi_0.bias.data.zero_()303 module.wi_1.weight.data.normal_(mean=0.0, std=factor * ((self.config.d_model) ** -0.5))304 if hasattr(module.wi_1, "bias") and module.wi_1.bias is not None:305 module.wi_1.bias.data.zero_()306 module.wo.weight.data.normal_(mean=0.0, std=factor * ((self.config.d_ff) ** -0.5))307 if hasattr(module.wo, "bias") and module.wo.bias is not None:308 module.wo.bias.data.zero_()309 elif isinstance(module, UdopAttention):310 # Mesh TensorFlow attention initialization to avoid scaling before softmax311 # See https://github.com/tensorflow/mesh/blob/fa19d69eafc9a482aff0b59ddd96b025c0cb207d/mesh_tensorflow/transformer/attention.py#L136312 d_model = self.config.d_model313 key_value_proj_dim = self.config.d_kv314 n_heads = self.config.num_heads315 module.q.weight.data.normal_(mean=0.0, std=factor * ((d_model * key_value_proj_dim) ** -0.5))316 module.k.weight.data.normal_(mean=0.0, std=factor * (d_model**-0.5))317 module.v.weight.data.normal_(mean=0.0, std=factor * (d_model**-0.5))318 module.o.weight.data.normal_(mean=0.0, std=factor * ((n_heads * key_value_proj_dim) ** -0.5))319 if module.has_relative_attention_bias:320 module.relative_attention_bias.weight.data.normal_(mean=0.0, std=factor * ((d_model) ** -0.5))321 322 # Copied from transformers.models.prophetnet.modeling_prophetnet.ProphetNetPreTrainedModel._shift_right with ProphetNet->Udop323 def _shift_right(self, input_ids):324 decoder_start_token_id = self.config.decoder_start_token_id325 pad_token_id = self.config.pad_token_id326 327 assert decoder_start_token_id is not None, (328 "self.model.config.decoder_start_token_id has to be defined. In Udop it is usually set to the"329 " pad_token_id. See Udop docs for more information"330 )331 332 # shift inputs to the right333 shifted_input_ids = input_ids.new_zeros(input_ids.shape)334 shifted_input_ids[..., 1:] = input_ids[..., :-1].clone()335 shifted_input_ids[..., 0] = decoder_start_token_id336 337 assert pad_token_id is not None, "self.model.config.pad_token_id has to be defined."338 # replace possible -100 values in labels by `pad_token_id`339 shifted_input_ids.masked_fill_(shifted_input_ids == -100, pad_token_id)340 341 assert torch.all(shifted_input_ids >= 0).item(), "Verify that `shifted_input_ids` has only positive values"342 343 return shifted_input_ids344 345 346# Copied from transformers.models.t5.modeling_t5.T5LayerNorm with T5->Udop347class UdopLayerNorm(nn.Module):348 def __init__(self, hidden_size, eps=1e-6):349 """350 Construct a layernorm module in the Udop style. No bias and no subtraction of mean.351 """352 super().__init__()353 self.weight = nn.Parameter(torch.ones(hidden_size))354 self.variance_epsilon = eps355 356 def forward(self, hidden_states):357 # Udop uses a layer_norm which only scales and doesn't shift, which is also known as Root Mean358 # Square Layer Normalization https://huggingface.co/papers/1910.07467 thus variance is calculated359 # w/o mean and there is no bias. Additionally we want to make sure that the accumulation for360 # half-precision inputs is done in fp32361 362 variance = hidden_states.to(torch.float32).pow(2).mean(-1, keepdim=True)363 hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)364 365 # convert into half-precision if necessary366 if self.weight.dtype in [torch.float16, torch.bfloat16]:367 hidden_states = hidden_states.to(self.weight.dtype)368 369 return self.weight * hidden_states370 371 372# Copied from transformers.models.t5.modeling_t5.T5DenseActDense with T5->Udop373class UdopDenseActDense(nn.Module):374 def __init__(self, config: UdopConfig):375 super().__init__()376 self.wi = nn.Linear(config.d_model, config.d_ff, bias=False)377 self.wo = nn.Linear(config.d_ff, config.d_model, bias=False)378 self.dropout = nn.Dropout(config.dropout_rate)379 self.act = ACT2FN[config.dense_act_fn]380 381 def forward(self, hidden_states):382 hidden_states = self.wi(hidden_states)383 hidden_states = self.act(hidden_states)384 hidden_states = self.dropout(hidden_states)385 if (386 isinstance(self.wo.weight, torch.Tensor)387 and hidden_states.dtype != self.wo.weight.dtype388 and self.wo.weight.dtype != torch.int8389 ):390 hidden_states = hidden_states.to(self.wo.weight.dtype)391 hidden_states = self.wo(hidden_states)392 return hidden_states393 394 395# Copied from transformers.models.t5.modeling_t5.T5DenseGatedActDense with T5->Udop396class UdopDenseGatedActDense(nn.Module):397 def __init__(self, config: UdopConfig):398 super().__init__()399 self.wi_0 = nn.Linear(config.d_model, config.d_ff, bias=False)400 self.wi_1 = nn.Linear(config.d_model, config.d_ff, bias=False)401 self.wo = nn.Linear(config.d_ff, config.d_model, bias=False)402 self.dropout = nn.Dropout(config.dropout_rate)403 self.act = ACT2FN[config.dense_act_fn]404 405 def forward(self, hidden_states):406 hidden_gelu = self.act(self.wi_0(hidden_states))407 hidden_linear = self.wi_1(hidden_states)408 hidden_states = hidden_gelu * hidden_linear409 hidden_states = self.dropout(hidden_states)410 411 # To make 8bit quantization work for google/flan-t5-xxl, self.wo is kept in float32.412 # See https://github.com/huggingface/transformers/issues/20287413 # we also make sure the weights are not in `int8` in case users will force `_keep_in_fp32_modules` to be `None``414 if (415 isinstance(self.wo.weight, torch.Tensor)416 and hidden_states.dtype != self.wo.weight.dtype417 and self.wo.weight.dtype != torch.int8418 ):419 hidden_states = hidden_states.to(self.wo.weight.dtype)420 421 hidden_states = self.wo(hidden_states)422 return hidden_states423 424 425# Copied from transformers.models.t5.modeling_t5.T5LayerFF with T5->Udop426class UdopLayerFF(nn.Module):427 def __init__(self, config: UdopConfig):428 super().__init__()429 if config.is_gated_act:430 self.DenseReluDense = UdopDenseGatedActDense(config)431 else:432 self.DenseReluDense = UdopDenseActDense(config)433 434 self.layer_norm = UdopLayerNorm(config.d_model, eps=config.layer_norm_epsilon)435 self.dropout = nn.Dropout(config.dropout_rate)436 437 def forward(self, hidden_states):438 forwarded_states = self.layer_norm(hidden_states)439 forwarded_states = self.DenseReluDense(forwarded_states)440 hidden_states = hidden_states + self.dropout(forwarded_states)441 return hidden_states442 443 444# Copied from transformers.models.t5.modeling_t5.T5Attention with T5->Udop445class UdopAttention(nn.Module):446 def __init__(447 self,448 config: UdopConfig,449 has_relative_attention_bias=False,450 layer_idx: Optional[int] = None,451 ):452 super().__init__()453 self.is_decoder = config.is_decoder454 self.has_relative_attention_bias = has_relative_attention_bias455 self.relative_attention_num_buckets = config.relative_attention_num_buckets456 self.relative_attention_max_distance = config.relative_attention_max_distance457 self.d_model = config.d_model458 self.key_value_proj_dim = config.d_kv459 self.n_heads = config.num_heads460 self.dropout = config.dropout_rate461 self.inner_dim = self.n_heads * self.key_value_proj_dim462 self.layer_idx = layer_idx463 if layer_idx is None and self.is_decoder:464 logger.warning_once(465 f"Instantiating a decoder {self.__class__.__name__} without passing `layer_idx` is not recommended and "466 "will to errors during the forward call, if caching is used. Please make sure to provide a `layer_idx` "467 "when creating this class."468 )469 470 # Mesh TensorFlow initialization to avoid scaling before softmax471 self.q = nn.Linear(self.d_model, self.inner_dim, bias=False)472 self.k = nn.Linear(self.d_model, self.inner_dim, bias=False)473 self.v = nn.Linear(self.d_model, self.inner_dim, bias=False)474 self.o = nn.Linear(self.inner_dim, self.d_model, bias=False)475 476 if self.has_relative_attention_bias:477 self.relative_attention_bias = nn.Embedding(self.relative_attention_num_buckets, self.n_heads)478 self.pruned_heads = set()479 self.gradient_checkpointing = False480 481 def prune_heads(self, heads):482 if len(heads) == 0:483 return484 heads, index = find_pruneable_heads_and_indices(485 heads, self.n_heads, self.key_value_proj_dim, self.pruned_heads486 )487 # Prune linear layers488 self.q = prune_linear_layer(self.q, index)489 self.k = prune_linear_layer(self.k, index)490 self.v = prune_linear_layer(self.v, index)491 self.o = prune_linear_layer(self.o, index, dim=1)492 # Update hyper params493 self.n_heads = self.n_heads - len(heads)494 self.inner_dim = self.key_value_proj_dim * self.n_heads495 self.pruned_heads = self.pruned_heads.union(heads)496 497 @staticmethod498 def _relative_position_bucket(relative_position, bidirectional=True, num_buckets=32, max_distance=128):499 """500 Adapted from Mesh Tensorflow:501 https://github.com/tensorflow/mesh/blob/0cb87fe07da627bf0b7e60475d59f95ed6b5be3d/mesh_tensorflow/transformer/transformer_layers.py#L593502 503 Translate relative position to a bucket number for relative attention. The relative position is defined as504 memory_position - query_position, i.e. the distance in tokens from the attending position to the attended-to505 position. If bidirectional=False, then positive relative positions are invalid. We use smaller buckets for506 small absolute relative_position and larger buckets for larger absolute relative_positions. All relative507 positions >=max_distance map to the same bucket. All relative positions <=-max_distance map to the same bucket.508 This should allow for more graceful generalization to longer sequences than the model has been trained on509 510 Args:511 relative_position: an int32 Tensor512 bidirectional: a boolean - whether the attention is bidirectional513 num_buckets: an integer514 max_distance: an integer515 516 Returns:517 a Tensor with the same shape as relative_position, containing int32 values in the range [0, num_buckets)518 """519 relative_buckets = 0520 if bidirectional:521 num_buckets //= 2522 relative_buckets += (relative_position > 0).to(torch.long) * num_buckets523 relative_position = torch.abs(relative_position)524 else:525 relative_position = -torch.min(relative_position, torch.zeros_like(relative_position))526 # now relative_position is in the range [0, inf)527 528 # half of the buckets are for exact increments in positions529 max_exact = num_buckets // 2530 is_small = relative_position < max_exact531 532 # The other half of the buckets are for logarithmically bigger bins in positions up to max_distance533 relative_position_if_large = max_exact + (534 torch.log(relative_position.float() / max_exact)535 / math.log(max_distance / max_exact)536 * (num_buckets - max_exact)537 ).to(torch.long)538 relative_position_if_large = torch.min(539 relative_position_if_large, torch.full_like(relative_position_if_large, num_buckets - 1)540 )541 542 relative_buckets += torch.where(is_small, relative_position, relative_position_if_large)543 return relative_buckets544 545 def compute_bias(self, query_length, key_length, device=None, cache_position=None):546 """Compute binned relative position bias"""547 if device is None:548 device = self.relative_attention_bias.weight.device549 if cache_position is None:550 context_position = torch.arange(query_length, dtype=torch.long, device=device)[:, None]551 else:552 context_position = cache_position[:, None].to(device)553 memory_position = torch.arange(key_length, dtype=torch.long, device=device)[None, :]554 relative_position = memory_position - context_position # shape (query_length, key_length)555 relative_position_bucket = self._relative_position_bucket(556 relative_position, # shape (query_length, key_length)557 bidirectional=(not self.is_decoder),558 num_buckets=self.relative_attention_num_buckets,559 max_distance=self.relative_attention_max_distance,560 )561 values = self.relative_attention_bias(relative_position_bucket) # shape (query_length, key_length, num_heads)562 values = values.permute([2, 0, 1]).unsqueeze(0) # shape (1, num_heads, query_length, key_length)563 return values564 565 @deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58")566 def forward(567 self,568 hidden_states,569 mask=None,570 key_value_states=None,571 position_bias=None,572 past_key_values=None,573 layer_head_mask=None,574 query_length=None,575 use_cache=False,576 output_attentions=False,577 cache_position=None,578 ):579 """580 Self-attention (if key_value_states is None) or attention over source sentence (provided by key_value_states).581 """582 # Input is (batch_size, seq_length, dim)583 # Mask is (batch_size, 1, 1, key_length) (non-causal encoder) or (batch_size, 1, seq_length, key_length) (causal decoder)584 batch_size, seq_length = hidden_states.shape[:2]585 586 # if key_value_states are provided this layer is used as a cross-attention layer for the decoder587 is_cross_attention = key_value_states is not None588 589 query_states = self.q(hidden_states)590 query_states = query_states.view(batch_size, -1, self.n_heads, self.key_value_proj_dim).transpose(1, 2)591 592 # Check is encoder-decoder model is being used. Otherwise we'll get `DynamicCache`593 is_updated = False594 if isinstance(past_key_values, EncoderDecoderCache):595 is_updated = past_key_values.is_updated.get(self.layer_idx)596 if is_cross_attention:597 # after the first generated id, we can subsequently re-use all key/value_states from cache598 curr_past_key_value = past_key_values.cross_attention_cache599 else:600 curr_past_key_value = past_key_values.self_attention_cache601 else:602 curr_past_key_value = past_key_values603 604 current_states = key_value_states if is_cross_attention else hidden_states605 if is_cross_attention and past_key_values is not None and is_updated:606 # reuse k,v, cross_attentions607 key_states = curr_past_key_value.layers[self.layer_idx].keys608 value_states = curr_past_key_value.layers[self.layer_idx].values609 else:610 key_states = self.k(current_states)611 value_states = self.v(current_states)612 key_states = key_states.view(batch_size, -1, self.n_heads, self.key_value_proj_dim).transpose(1, 2)613 value_states = value_states.view(batch_size, -1, self.n_heads, self.key_value_proj_dim).transpose(1, 2)614 615 if past_key_values is not None:616 # save all key/value_states to cache to be re-used for fast auto-regressive generation617 cache_position = cache_position if not is_cross_attention else None618 key_states, value_states = curr_past_key_value.update(619 key_states, value_states, self.layer_idx, {"cache_position": cache_position}620 )621 # set flag that curr layer for cross-attn is already updated so we can re-use in subsequent calls622 if is_cross_attention and isinstance(past_key_values, EncoderDecoderCache):623 past_key_values.is_updated[self.layer_idx] = True624 625 # compute scores, equivalent of torch.einsum("bnqd,bnkd->bnqk", query_states, key_states), compatible with onnx op>9626 scores = torch.matmul(query_states, key_states.transpose(3, 2))627 628 if position_bias is None:629 key_length = key_states.shape[-2]630 # cache position is 0-indexed so we add 1 to get the real length of queries (aka with past)631 real_seq_length = query_length if query_length is not None else cache_position[-1] + 1632 if not self.has_relative_attention_bias:633 position_bias = torch.zeros(634 (1, self.n_heads, seq_length, key_length), device=scores.device, dtype=scores.dtype635 )636 if self.gradient_checkpointing and self.training:637 position_bias.requires_grad = True638 else:639 position_bias = self.compute_bias(640 real_seq_length, key_length, device=scores.device, cache_position=cache_position641 )642 position_bias = position_bias[:, :, -seq_length:, :]643 644 if mask is not None:645 causal_mask = mask[:, :, :, : key_states.shape[-2]]646 position_bias = position_bias + causal_mask647 648 if self.pruned_heads:649 mask = torch.ones(position_bias.shape[1])650 mask[list(self.pruned_heads)] = 0651 position_bias_masked = position_bias[:, mask.bool()]652 else:653 position_bias_masked = position_bias654 655 scores += position_bias_masked656 657 # (batch_size, n_heads, seq_length, key_length)658 attn_weights = nn.functional.softmax(scores.float(), dim=-1).type_as(scores)659 attn_weights = nn.functional.dropout(attn_weights, p=self.dropout, training=self.training)660 661 # Mask heads if we want to662 if layer_head_mask is not None:663 attn_weights = attn_weights * layer_head_mask664 665 attn_output = torch.matmul(attn_weights, value_states)666 667 attn_output = attn_output.transpose(1, 2).contiguous()668 attn_output = attn_output.view(batch_size, -1, self.inner_dim)669 attn_output = self.o(attn_output)670 671 outputs = (attn_output, position_bias)672 673 if output_attentions:674 outputs = outputs + (attn_weights,)675 return outputs676 677 678# Copied from transformers.models.t5.modeling_t5.T5LayerSelfAttention with T5->Udop679class UdopLayerSelfAttention(nn.Module):680 def __init__(self, config, has_relative_attention_bias=False, layer_idx: Optional[int] = None):681 super().__init__()682 self.SelfAttention = UdopAttention(683 config, has_relative_attention_bias=has_relative_attention_bias, layer_idx=layer_idx684 )685 self.layer_norm = UdopLayerNorm(config.d_model, eps=config.layer_norm_epsilon)686 self.dropout = nn.Dropout(config.dropout_rate)687 688 @deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58")689 def forward(690 self,691 hidden_states,692 attention_mask=None,693 position_bias=None,694 layer_head_mask=None,695 past_key_values=None,696 use_cache=False,697 output_attentions=False,698 cache_position=None,699 ):700 normed_hidden_states = self.layer_norm(hidden_states)701 attention_output = self.SelfAttention(702 normed_hidden_states,703 mask=attention_mask,704 position_bias=position_bias,705 layer_head_mask=layer_head_mask,706 past_key_values=past_key_values,707 use_cache=use_cache,708 output_attentions=output_attentions,709 cache_position=cache_position,710 )711 hidden_states = hidden_states + self.dropout(attention_output[0])712 outputs = (hidden_states,) + attention_output[1:] # add attentions if we output them713 return outputs714 715 716# Copied from transformers.models.t5.modeling_t5.T5LayerCrossAttention with T5->Udop717class UdopLayerCrossAttention(nn.Module):718 def __init__(self, config, layer_idx: Optional[int] = None):719 super().__init__()720 self.EncDecAttention = UdopAttention(config, has_relative_attention_bias=False, layer_idx=layer_idx)721 self.layer_norm = UdopLayerNorm(config.d_model, eps=config.layer_norm_epsilon)722 self.dropout = nn.Dropout(config.dropout_rate)723 724 @deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58")725 def forward(726 self,727 hidden_states,728 key_value_states,729 attention_mask=None,730 position_bias=None,731 layer_head_mask=None,732 past_key_values=None,733 use_cache=False,734 query_length=None,735 output_attentions=False,736 cache_position=None,737 ):738 normed_hidden_states = self.layer_norm(hidden_states)739 attention_output = self.EncDecAttention(740 normed_hidden_states,741 mask=attention_mask,742 key_value_states=key_value_states,743 position_bias=position_bias,744 layer_head_mask=layer_head_mask,745 past_key_values=past_key_values,746 use_cache=use_cache,747 query_length=query_length,748 output_attentions=output_attentions,749 cache_position=cache_position,750 )751 layer_output = hidden_states + self.dropout(attention_output[0])752 outputs = (layer_output,) + attention_output[1:] # add attentions if we output them753 return outputs754 755 756# Copied from transformers.models.t5.modeling_t5.T5Block with T5->Udop757class UdopBlock(GradientCheckpointingLayer):758 def __init__(self, config, has_relative_attention_bias=False, layer_idx: Optional[int] = None):759 super().__init__()760 self.is_decoder = config.is_decoder761 self.layer = nn.ModuleList()762 self.layer.append(763 UdopLayerSelfAttention(764 config, has_relative_attention_bias=has_relative_attention_bias, layer_idx=layer_idx765 )766 )767 if self.is_decoder:768 self.layer.append(UdopLayerCrossAttention(config, layer_idx=layer_idx))769 770 self.layer.append(UdopLayerFF(config))771 772 @deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58")773 def forward(774 self,775 hidden_states,776 attention_mask=None,777 position_bias=None,778 encoder_hidden_states=None,779 encoder_attention_mask=None,780 encoder_decoder_position_bias=None,781 layer_head_mask=None,782 cross_attn_layer_head_mask=None,783 past_key_values=None,784 use_cache=False,785 output_attentions=False,786 return_dict=True,787 cache_position=None,788 ):789 self_attention_outputs = self.layer[0](790 hidden_states,791 attention_mask=attention_mask,792 position_bias=position_bias,793 layer_head_mask=layer_head_mask,794 past_key_values=past_key_values,795 use_cache=use_cache,796 output_attentions=output_attentions,797 cache_position=cache_position,798 )799 hidden_states = self_attention_outputs[0]800 attention_outputs = self_attention_outputs[1:] # Keep self-attention outputs and relative position weights801 802 # clamp inf values to enable fp16 training803 if hidden_states.dtype == torch.float16:804 clamp_value = torch.where(805 torch.isinf(hidden_states).any(),806 torch.finfo(hidden_states.dtype).max - 1000,807 torch.finfo(hidden_states.dtype).max,808 )809 hidden_states = torch.clamp(hidden_states, min=-clamp_value, max=clamp_value)810 811 do_cross_attention = self.is_decoder and encoder_hidden_states is not None812 if do_cross_attention:813 cross_attention_outputs = self.layer[1](814 hidden_states,815 key_value_states=encoder_hidden_states,816 attention_mask=encoder_attention_mask,817 position_bias=encoder_decoder_position_bias,818 layer_head_mask=cross_attn_layer_head_mask,819 past_key_values=past_key_values,820 query_length=cache_position[-1] + 1,821 use_cache=use_cache,822 output_attentions=output_attentions,823 )824 hidden_states = cross_attention_outputs[0]825 826 # clamp inf values to enable fp16 training827 if hidden_states.dtype == torch.float16:828 clamp_value = torch.where(829 torch.isinf(hidden_states).any(),830 torch.finfo(hidden_states.dtype).max - 1000,831 torch.finfo(hidden_states.dtype).max,832 )833 hidden_states = torch.clamp(hidden_states, min=-clamp_value, max=clamp_value)834 835 # Keep cross-attention outputs and relative position weights836 attention_outputs = attention_outputs + cross_attention_outputs[1:]837 838 # Apply Feed Forward layer839 hidden_states = self.layer[-1](hidden_states)840 841 # clamp inf values to enable fp16 training842 if hidden_states.dtype == torch.float16:843 clamp_value = torch.where(844 torch.isinf(hidden_states).any(),845 torch.finfo(hidden_states.dtype).max - 1000,846 torch.finfo(hidden_states.dtype).max,847 )848 hidden_states = torch.clamp(hidden_states, min=-clamp_value, max=clamp_value)849 850 outputs = (hidden_states,)851 852 return (853 outputs + attention_outputs854 ) # hidden-states, (self-attention position bias), (self-attention weights), (cross-attention position bias), (cross-attention weights)855 856 857class UdopCellEmbeddings(nn.Module):858 def __init__(self, max_2d_position_embeddings=501, hidden_size=1024):859 super().__init__()860 self.max_2d_position_embeddings = max_2d_position_embeddings861 862 self.x_position_embeddings = nn.Embedding(max_2d_position_embeddings, hidden_size)863 self.y_position_embeddings = nn.Embedding(max_2d_position_embeddings, hidden_size)864 865 def forward(self, bbox):866 bbox = torch.clip(bbox, 0.0, 1.0)867 bbox = (bbox * (self.max_2d_position_embeddings - 1)).long()868 left_position_embeddings = self.x_position_embeddings(bbox[:, :, 0])869 upper_position_embeddings = self.y_position_embeddings(bbox[:, :, 1])870 right_position_embeddings = self.x_position_embeddings(bbox[:, :, 2])871 lower_position_embeddings = self.y_position_embeddings(bbox[:, :, 3])872 873 embeddings = (874 left_position_embeddings875 + upper_position_embeddings876 + right_position_embeddings877 + lower_position_embeddings878 )879 880 return embeddings881 882 883# get function for bucket computation884# protected member access seems to be lesser evil than copy paste whole function885get_relative_position_bucket = UdopAttention._relative_position_bucket886AUGMENTATION_RANGE = (0.80, 1.25)887 888 889class RelativePositionBiasBase(nn.Module, ABC):890 """891 Base class of relative biases.892 893 Args:894 num_heads (`int`):895 Number of attention heads in the model, it will create embeddings of size `num_heads`, which will be added to the scores of each token pair.896 relative_attention_num_buckets (`int`, *optional*, defaults to 32):897 Pair token metric (distance in the sequence, distance in pixels etc.) will be bucketed, parameter is defining number of such898 buckets.899 bidirectional (`bool`, *optional*, defaults to `True`):900 Whether the distance should be bidirectional for a pair of tokens. If `False`, then distance(tok1, tok2) == distance(tok2, tok1).901 scaling_factor (`int`, *optional*, defaults to 1):902 Defining factor which will be used to scale relative distance.903 max_distance (`int`, *optional*, defaults to 128):904 All distances above this value will end up in the one/same bucket.905 augmentation (`bool`, *optional*, defaults to `False`):906 Whether to multiply relative distances by a random scalar.907 expand (`bool`, *optional*, defaults to `False`):908 Whether to expand an existing pretrained model with subsequent additions of prefix_bucket.909 """910 911 def __init__(912 self,913 num_heads=None,914 relative_attention_num_buckets=32,915 bidirectional=True,916 scaling_factor=1,917 max_distance=128,918 level="tokens",919 augmentation=False,920 prefix_bucket=False,921 expand=False,922 ):923 super().__init__()924 self.prefix_bucket = prefix_bucket925 self.augmentation = augmentation926 self.level = level927 self.max_distance = max_distance928 self.scaling_factor = scaling_factor929 self.bidirectional = bidirectional930 self.num_heads = num_heads931 self.expand = expand932 self.relative_attention_num_buckets = relative_attention_num_buckets933 extra_head = 2 if prefix_bucket and not self.expand else 0934 self.relative_attention_bias = nn.Embedding(self.relative_attention_num_buckets + extra_head, self.num_heads)935 936 @abstractmethod937 def prepare_input(938 self,939 attention_mask: Optional[Tensor] = None,940 bbox: Optional[dict[str, Any]] = None,941 ) -> Tensor:942 pass943 944 def get_bucket(self, attention_mask: Optional[Tensor] = None, bbox: Optional[dict[str, Any]] = None) -> Tensor:945 relative_position = self.prepare_input(attention_mask, bbox)946 rp_bucket: Tensor = get_relative_position_bucket(947 relative_position,948 bidirectional=self.bidirectional,949 num_buckets=self.relative_attention_num_buckets,950 max_distance=self.max_distance,951 )952 return rp_bucket953 954 def get_relative_position(self, positions):955 context_position = positions[:, :, None]956 memory_position = positions[:, None, :]957 relative_position = memory_position - context_position958 if self.augmentation and self.training:959 relative_position *= random.uniform(*AUGMENTATION_RANGE)960 relative_position *= self.scaling_factor961 962 return relative_position.to(torch.long)963 964 def forward(self, attention_mask: Optional[Tensor] = None, bbox: Optional[dict[str, Any]] = None) -> Tensor:965 # re-using pretrained model with subsequent addition of prefix_bucket966 if self.expand and self.prefix_bucket:967 new_bias = nn.Embedding(self.relative_attention_num_buckets + 2, self.num_heads)968 new_bias.weight.data[: self.relative_attention_num_buckets] = self.relative_attention_bias.weight.data969 new_bias.weight.data[self.relative_attention_num_buckets :] = 0.1970 self.relative_attention_bias = new_bias971 self.expand = False972 973 rp_bucket = self.get_bucket(attention_mask, bbox)974 975 if self.prefix_bucket:976 if rp_bucket.size(0) == 1 and attention_mask.size(0) > 1:977 rp_bucket = rp_bucket.repeat(attention_mask.size(0), 1, 1)978 # based on assumption that prefix bboxes are negative979 is_prefix = bbox[:, :, 1] < 0980 num_prefix = is_prefix.sum(-1)981 for idx, num_prefix_row in enumerate(num_prefix.cpu().numpy()):982 rp_bucket[idx, :num_prefix_row, num_prefix_row:] = self.relative_attention_num_buckets983 rp_bucket[idx, num_prefix_row:, :num_prefix_row] = self.relative_attention_num_buckets + 1984 985 values: Tensor = self.relative_attention_bias(rp_bucket)986 if values.dim() != 4:987 raise ValueError("Wrong dimension of values tensor")988 values = values.permute([0, 3, 1, 2])989 990 return values991 992 993class RelativePositionBias1D(RelativePositionBiasBase):994 def __init__(self, scaling_factor=1, max_distance=128, **kwargs):995 """996 Reimplementation of T5 relative position bias. Distance between given tokens is their distance in the sequence.997 Parameters are the same as in base class998 """999 super().__init__(scaling_factor=scaling_factor, max_distance=max_distance, **kwargs)1000 1001 def prepare_input(self, attention_mask: Optional[Tensor] = None, bbox: Optional[dict[str, Any]] = None) -> Tensor:1002 if self.scaling_factor != 1:1003 raise ValueError("No need to scale 1d features")1004 relative_position = self.get_relative_position(1005 torch.arange(attention_mask.size(1), dtype=torch.long, device=attention_mask.device)[None, :]1006 )1007 1008 return relative_position1009 1010 1011class RelativePositionBiasHorizontal(RelativePositionBiasBase):1012 def __init__(self, scaling_factor=100, max_distance=100, **kwargs):1013 """1014 Represents in the bucket embeddings horizontal distance between two tokens. Parameters are the same as in base1015 class1016 """1017 super().__init__(scaling_factor=scaling_factor, max_distance=max_distance, **kwargs)1018 1019 def prepare_input(self, attention_mask: Optional[Tensor] = None, bbox: Optional[dict[str, Any]] = None) -> Tensor:1020 if not self.scaling_factor > 1.0:1021 raise ValueError("Need to scale the values of bboxes, as there are in small (0,1) range")1022 if bbox is None:1023 raise ValueError("Bbox is required for horizontal relative position bias")1024 # get x positions of left point of bbox1025 horizontal_position: Tensor = bbox[:, :, [0, 2]].mean(dim=-1)1026 1027 return self.get_relative_position(horizontal_position)1028 1029 1030class RelativePositionBiasVertical(RelativePositionBiasBase):1031 def __init__(self, scaling_factor=100, max_distance=100, **kwargs):1032 """1033 Represents in the bucket embeddings vertical distance between two tokens. Parameters are the same as in base1034 class1035 """1036 super().__init__(scaling_factor=scaling_factor, max_distance=max_distance, **kwargs)1037 1038 def prepare_input(self, attention_mask: Optional[Tensor] = None, bbox: Optional[dict[str, Any]] = None) -> Tensor:1039 if not self.scaling_factor > 1.0:1040 raise ValueError("Need to scale the values of bboxes, as there are in small (0,1) range")1041 if bbox is None:1042 raise ValueError("Bbox is required for vertical relative position bias")1043 # get y positions of middle of bbox1044 vertical_position: Tensor = bbox[:, :, [1, 3]].mean(dim=-1)1045 1046 return self.get_relative_position(vertical_position)1047 1048 1049class RelativePositionBiasAggregated(nn.Module):1050 def __init__(self, modules: Sequence[RelativePositionBiasBase]):1051 """1052 Class which sums up various computed biases.1053 1054 Args:1055 modules (Sequence[RelativePositionBiasBase]):1056 List of relative bias modules.1057 """1058 super().__init__()1059 self.biases = nn.ModuleList(modules)1060 1061 def forward(1062 self, attention_mask: Optional[Tensor] = None, bbox: Optional[dict[str, Any]] = None1063 ) -> Union[float, Tensor]:1064 output = 0.01065 for bias in self.biases: # type: ignore1066 output = bias(attention_mask, bbox) + output1067 1068 return output1069 1070 1071BIAS_CLASSES = {1072 "1d": RelativePositionBias1D,1073 "horizontal": RelativePositionBiasHorizontal,1074 "vertical": RelativePositionBiasVertical,1075}1076 1077 1078def create_relative_bias(config: UdopConfig) -> Sequence[RelativePositionBiasBase]:1079 """1080 Creates empty list or one/multiple relative biases.1081 1082 :param config: Model's configuration :return: Sequence with created bias modules.1083 """1084 bias_list = []1085 if hasattr(config, "relative_bias_args"):1086 for bias_kwargs_org in config.relative_bias_args:1087 bias_kwargs = deepcopy(bias_kwargs_org)1088 bias_type = bias_kwargs.pop("type")1089 model_num_heads = config.num_heads if hasattr(config, "num_heads") else config.num_attention_heads1090 if "num_heads" in bias_kwargs:1091 if bias_kwargs["num_heads"] != model_num_heads:1092 raise ValueError("Number of heads must match num of heads in the model")1093 else:1094 bias_kwargs["num_heads"] = model_num_heads1095 bias_list.append(BIAS_CLASSES[bias_type](**bias_kwargs)) # type: ignore1096 1097 return bias_list1098 1099 1100class UdopStack(UdopPreTrainedModel):1101 """1102 This class is based on `T5Stack`, but modified to take into account the image modality as well as 2D position1103 embeddings.1104 """1105 1106 def __init__(self, config, embed_tokens=None, embed_patches=None):1107 super().__init__(config)1108 1109 self.embed_tokens = embed_tokens1110 self.embed_patches = embed_patches1111 self.is_decoder = config.is_decoder1112 self._max_length = config.max_length1113 self.num_layers = config.num_layers1114 1115 self.block = nn.ModuleList(1116 [UdopBlock(config, has_relative_attention_bias=bool(i == 0), layer_idx=i) for i in range(self.num_layers)]1117 )1118 self.final_layer_norm = UdopLayerNorm(config.d_model, eps=config.layer_norm_epsilon)1119 1120 self.dropout = nn.Dropout(config.dropout_rate)1121 1122 if not self.is_decoder:1123 self.cell_2d_embedding = UdopCellEmbeddings(config.max_2d_position_embeddings, config.hidden_size)1124 1125 # get weights from encoder position bias1126 self.relative_bias = self._get_relative_bias(config)1127 1128 def _tie_weights(self):1129 for bias in self.relative_bias.biases:1130 if isinstance(bias, RelativePositionBias1D):1131 self._tie_or_clone_weights(1132 bias.relative_attention_bias, self.block[0].layer[0].SelfAttention.relative_attention_bias1133 )1134 1135 @staticmethod1136 def _get_relative_bias(config: UdopConfig) -> RelativePositionBiasAggregated:1137 relative_bias_list = create_relative_bias(config)1138 return RelativePositionBiasAggregated(relative_bias_list)1139 1140 def get_output_embeddings(self):1141 return self.embed_tokens1142 1143 def set_input_embeddings(self, new_embeddings):1144 self.embed_tokens = new_embeddings1145 1146 def forward(1147 self,1148 input_ids=None,1149 attention_mask=None,1150 bbox=None,1151 encoder_hidden_states=None,1152 encoder_attention_mask=None,1153 inputs_embeds=None,1154 pixel_values=None,1155 visual_bbox=None,1156 image_embeddings=None,1157 position_bias=None,1158 head_mask=None,1159 cross_attn_head_mask=None,1160 past_key_values=None,1161 use_cache=None,1162 output_attentions=None,1163 output_hidden_states=None,1164 return_dict=None,1165 cache_position=None,1166 ):1167 use_cache = use_cache if use_cache is not None else self.config.use_cache1168 output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions1169 output_hidden_states = (1170 output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states1171 )1172 return_dict = return_dict if return_dict is not None else self.config.use_return_dict1173 1174 # input embeddings processing1175 1176 if input_ids is not None and inputs_embeds is not None:1177 err_msg_prefix = "decoder_" if self.is_decoder else ""1178 raise ValueError(1179 f"You cannot specify both {err_msg_prefix}inputs and {err_msg_prefix}inputs_embeds at the same time"1180 )1181 elif input_ids is not None and torch.numel(input_ids) > 0:1182 input_shape = input_ids.size()1183 input_ids = input_ids.view(-1, input_shape[-1])1184 elif inputs_embeds is None and input_ids is not None and torch.numel(input_ids) == 0:1185 input_ids = torch.full((4, 1024), self.config.pad_token_id, device=input_ids.device, dtype=input_ids.dtype)1186 attention_mask = torch.zeros((4, 1024), device=input_ids.device, dtype=input_ids.dtype)1187 bbox = torch.zeros((4, 1024, 4), device=input_ids.device, dtype=input_ids.dtype)1188 input_shape = input_ids.size()1189 position_bias = torch.zeros_like(self.get_extended_attention_mask(attention_mask, input_shape))1190 # encoder_attention_mask = attention_mask1191 logger.warning("Empty batch")1192 elif inputs_embeds is not None:1193 input_shape = inputs_embeds.size()[:-1]1194 else:1195 err_msg_prefix = "decoder_" if self.is_decoder else ""1196 raise ValueError(f"You have to specify either {err_msg_prefix}inputs or {err_msg_prefix}inputs_embeds")1197 1198 if inputs_embeds is None:1199 if self.embed_tokens is None:1200 raise ValueError("You have to initialize the model with valid token embeddings")