Xenova/tiny-random-Florence2ForConditionalGeneration
61.3k
1# coding=utf-82# Copyright 2024 Microsoft and the HuggingFace Inc. team. All rights reserved.3#4# Licensed under the Apache License, Version 2.0 (the "License");5# you may not use this file except in compliance with the License.6# You may obtain a copy of the License 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 16""" PyTorch Florence-2 model."""17from dataclasses import dataclass18from typing import List, Optional, Tuple, Union19 20import math21import torch22import torch.utils.checkpoint23from torch import nn24import torch.nn.functional as F25import torch.utils.checkpoint as checkpoint26from torch.nn import CrossEntropyLoss 27from collections import OrderedDict28from einops import rearrange29from timm.models.layers import DropPath, trunc_normal_30 31from transformers.modeling_utils import PreTrainedModel32from transformers.utils import (33 ModelOutput,34 add_start_docstrings,35 add_start_docstrings_to_model_forward,36 is_flash_attn_2_available,37 logging,38 replace_return_docstrings,39 is_flash_attn_2_available,40 is_flash_attn_greater_or_equal_2_10,41)42from .configuration_florence2 import Florence2Config 43from .configuration_florence2 import Florence2LanguageConfig44from .configuration_florence2 import Florence2VisionConfig45 46 47from transformers.activations import ACT2FN48from transformers.modeling_attn_mask_utils import (49 _prepare_4d_attention_mask,50 _prepare_4d_attention_mask_for_sdpa,51 _prepare_4d_causal_attention_mask,52 _prepare_4d_causal_attention_mask_for_sdpa,53)54from transformers.modeling_outputs import (55 BaseModelOutput,56 BaseModelOutputWithPastAndCrossAttentions,57 Seq2SeqLMOutput,58 Seq2SeqModelOutput,59)60 61 62if is_flash_attn_2_available():63 from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input # noqa64 65logger = logging.get_logger(__name__)66 67_CONFIG_FOR_DOC = "Florence2Config"68 69class LearnedAbsolutePositionEmbedding2D(nn.Module):70 """71 This module learns positional embeddings up to a fixed maximum size.72 """73 74 def __init__(self, embedding_dim=256, num_pos=50):75 super().__init__()76 self.row_embeddings = nn.Embedding(num_pos, embedding_dim // 2)77 self.column_embeddings = nn.Embedding(num_pos, embedding_dim - (embedding_dim // 2))78 79 def forward(self, pixel_values):80 """81 pixel_values: (batch_size, height, width, num_channels) 82 returns: (batch_size, height, width, embedding_dim * 2)83 """84 if len(pixel_values.shape) != 4:85 raise ValueError('pixel_values must be a 4D tensor')86 height, width = pixel_values.shape[1:3]87 width_values = torch.arange(width, device=pixel_values.device)88 height_values = torch.arange(height, device=pixel_values.device)89 x_emb = self.column_embeddings(width_values)90 y_emb = self.row_embeddings(height_values)91 # (height, width, embedding_dim * 2)92 pos = torch.cat([x_emb.unsqueeze(0).repeat(height, 1, 1), y_emb.unsqueeze(1).repeat(1, width, 1)], dim=-1)93 # (embedding_dim * 2, height, width)94 pos = pos.permute(2, 0, 1)95 pos = pos.unsqueeze(0)96 # (batch_size, embedding_dim * 2, height, width)97 pos = pos.repeat(pixel_values.shape[0], 1, 1, 1)98 # (batch_size, height, width, embedding_dim * 2)99 pos = pos.permute(0, 2, 3, 1)100 return pos101 102class PositionalEmbeddingCosine1D(nn.Module):103 """104 This class implements a very simple positional encoding. It follows closely105 the encoder from the link below:106 https://pytorch.org/tutorials/beginner/translation_transformer.html107 108 Args:109 embed_dim: The dimension of the embeddings.110 dropout_prob: The dropout probability.111 max_seq_len: The maximum length to precompute the positional encodings.112 """113 def __init__(114 self,115 embed_dim: int = 512,116 max_seq_len: int = 1024) -> None:117 super(PositionalEmbeddingCosine1D, self).__init__()118 self.embed_dim = embed_dim119 self.max_seq_len = max_seq_len120 # Generate the sinusoidal arrays.121 factor = math.log(10000)122 denominator = torch.exp(123 -factor * torch.arange(0, self.embed_dim, 2) / self.embed_dim)124 # Matrix where rows correspond to a positional embedding as a function125 # of the position index (i.e., the row index).126 frequencies = \127 torch.arange(0, self.max_seq_len) \128 .reshape(self.max_seq_len, 1) * denominator129 pos_idx_to_embed = torch.zeros((self.max_seq_len, self.embed_dim))130 # Populate uneven entries.131 pos_idx_to_embed[:, 0::2] = torch.sin(frequencies)132 pos_idx_to_embed[:, 1::2] = torch.cos(frequencies)133 # Save the positional embeddings in a constant buffer.134 self.register_buffer("pos_idx_to_embed", pos_idx_to_embed)135 136 def forward(self, seq_embeds: torch.Tensor) -> torch.Tensor:137 """138 Args:139 seq_embeds: The sequence embeddings in order. Allowed size:140 1. [T, D], where T is the length of the sequence, and D is the141 frame embedding dimension.142 2. [B, T, D], where B is the batch size and T and D are the143 same as above.144 145 Returns a tensor of with the same dimensions as the input: i.e.,146 [1, T, D] or [T, D].147 """148 shape_len = len(seq_embeds.shape)149 assert 2 <= shape_len <= 3150 len_seq = seq_embeds.size(-2)151 assert len_seq <= self.max_seq_len152 pos_embeds = self.pos_idx_to_embed[0:seq_embeds.size(-2), :]153 # Adapt pre-computed positional embeddings to the input.154 if shape_len == 3:155 pos_embeds = pos_embeds.view(156 (1, pos_embeds.size(0), pos_embeds.size(1)))157 return pos_embeds158 159 160class LearnedAbsolutePositionEmbedding1D(nn.Module):161 """162 Learnable absolute positional embeddings for 1D sequences.163 164 Args:165 embed_dim: The dimension of the embeddings.166 max_seq_len: The maximum length to precompute the positional encodings.167 """168 def __init__(169 self,170 embedding_dim: int = 512,171 num_pos: int = 1024) -> None:172 super(LearnedAbsolutePositionEmbedding1D, self).__init__()173 self.embeddings = nn.Embedding(num_pos, embedding_dim)174 self.num_pos = num_pos175 176 def forward(self, seq_embeds: torch.Tensor) -> torch.Tensor:177 """178 Args:179 seq_embeds: The sequence embeddings in order. Allowed size:180 1. [T, D], where T is the length of the sequence, and D is the181 frame embedding dimension.182 2. [B, T, D], where B is the batch size and T and D are the183 same as above.184 185 Returns a tensor of with the same dimensions as the input: i.e.,186 [1, T, D] or [T, D].187 """188 shape_len = len(seq_embeds.shape)189 assert 2 <= shape_len <= 3190 len_seq = seq_embeds.size(-2)191 assert len_seq <= self.num_pos192 # [T, D]193 pos_embeds = self.embeddings(torch.arange(len_seq).to(seq_embeds.device))194 # Adapt pre-computed positional embeddings to the input.195 if shape_len == 3:196 pos_embeds = pos_embeds.view(197 (1, pos_embeds.size(0), pos_embeds.size(1)))198 return pos_embeds199 200 201 202class MySequential(nn.Sequential):203 def forward(self, *inputs):204 for module in self._modules.values():205 if type(inputs) == tuple:206 inputs = module(*inputs)207 else:208 inputs = module(inputs)209 return inputs210 211 212class PreNorm(nn.Module):213 def __init__(self, norm, fn, drop_path=None):214 super().__init__()215 self.norm = norm216 self.fn = fn217 self.drop_path = drop_path218 219 def forward(self, x, *args, **kwargs):220 shortcut = x221 if self.norm != None:222 x, size = self.fn(self.norm(x), *args, **kwargs)223 else:224 x, size = self.fn(x, *args, **kwargs)225 226 if self.drop_path:227 x = self.drop_path(x)228 229 x = shortcut + x230 231 return x, size232 233 234class Mlp(nn.Module):235 def __init__(236 self,237 in_features,238 hidden_features=None,239 out_features=None,240 act_layer=nn.GELU,241 ):242 super().__init__()243 out_features = out_features or in_features244 hidden_features = hidden_features or in_features245 self.net = nn.Sequential(OrderedDict([246 ("fc1", nn.Linear(in_features, hidden_features)),247 ("act", act_layer()),248 ("fc2", nn.Linear(hidden_features, out_features))249 ]))250 251 def forward(self, x, size):252 return self.net(x), size253 254 255class DepthWiseConv2d(nn.Module):256 def __init__(257 self,258 dim_in,259 kernel_size,260 padding,261 stride,262 bias=True,263 ):264 super().__init__()265 self.dw = nn.Conv2d(266 dim_in, dim_in,267 kernel_size=kernel_size,268 padding=padding,269 groups=dim_in,270 stride=stride,271 bias=bias272 )273 274 def forward(self, x, size):275 B, N, C = x.shape276 H, W = size277 assert N == H * W278 279 x = self.dw(x.transpose(1, 2).view(B, C, H, W))280 size = (x.size(-2), x.size(-1))281 x = x.flatten(2).transpose(1, 2)282 return x, size283 284 285class ConvEmbed(nn.Module):286 """ Image to Patch Embedding287 """288 289 def __init__(290 self,291 patch_size=7,292 in_chans=3,293 embed_dim=64,294 stride=4,295 padding=2,296 norm_layer=None,297 pre_norm=True298 ):299 super().__init__()300 self.patch_size = patch_size301 302 self.proj = nn.Conv2d(303 in_chans, embed_dim,304 kernel_size=patch_size,305 stride=stride,306 padding=padding307 )308 309 dim_norm = in_chans if pre_norm else embed_dim310 self.norm = norm_layer(dim_norm) if norm_layer else None311 312 self.pre_norm = pre_norm313 314 def forward(self, x, size):315 H, W = size316 if len(x.size()) == 3:317 if self.norm and self.pre_norm:318 x = self.norm(x)319 x = rearrange(320 x, 'b (h w) c -> b c h w',321 h=H, w=W322 )323 324 x = self.proj(x)325 326 _, _, H, W = x.shape327 x = rearrange(x, 'b c h w -> b (h w) c')328 if self.norm and not self.pre_norm:329 x = self.norm(x)330 331 return x, (H, W)332 333 334class ChannelAttention(nn.Module):335 336 def __init__(self, dim, groups=8, qkv_bias=True):337 super().__init__()338 339 self.groups = groups340 self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)341 self.proj = nn.Linear(dim, dim)342 343 def forward(self, x, size):344 B, N, C = x.shape345 346 qkv = self.qkv(x).reshape(B, N, 3, self.groups, C // self.groups).permute(2, 0, 3, 1, 4)347 q, k, v = qkv[0], qkv[1], qkv[2]348 349 q = q * (float(N) ** -0.5)350 attention = q.transpose(-1, -2) @ k351 attention = attention.softmax(dim=-1)352 x = (attention @ v.transpose(-1, -2)).transpose(-1, -2)353 x = x.transpose(1, 2).reshape(B, N, C)354 x = self.proj(x)355 return x, size356 357 358class ChannelBlock(nn.Module):359 360 def __init__(self, dim, groups, mlp_ratio=4., qkv_bias=True,361 drop_path_rate=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm,362 conv_at_attn=True, conv_at_ffn=True):363 super().__init__()364 365 drop_path = DropPath(drop_path_rate) if drop_path_rate > 0. else nn.Identity()366 367 self.conv1 = PreNorm(None, DepthWiseConv2d(dim, 3, 1, 1)) if conv_at_attn else None368 self.channel_attn = PreNorm(369 norm_layer(dim),370 ChannelAttention(dim, groups=groups, qkv_bias=qkv_bias),371 drop_path372 )373 self.conv2 = PreNorm(None, DepthWiseConv2d(dim, 3, 1, 1)) if conv_at_ffn else None374 self.ffn = PreNorm(375 norm_layer(dim),376 Mlp(in_features=dim, hidden_features=int(dim*mlp_ratio), act_layer=act_layer),377 drop_path378 )379 380 def forward(self, x, size):381 if self.conv1:382 x, size = self.conv1(x, size)383 x, size = self.channel_attn(x, size)384 385 if self.conv2:386 x, size = self.conv2(x, size)387 x, size = self.ffn(x, size)388 389 return x, size390 391 392def window_partition(x, window_size: int):393 B, H, W, C = x.shape394 x = x.view(B, H // window_size, window_size, W // window_size, window_size, C)395 windows = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size, window_size, C)396 return windows397 398 399def window_reverse(windows, batch_size: int, window_size: int, H: int, W: int):400 B = batch_size 401 # this will cause onnx conversion failed for dynamic axis, because treated as constant402 # int(windows.shape[0] / (H * W / window_size / window_size)) 403 x = windows.view(B, H // window_size, W // window_size, window_size, window_size, -1)404 x = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, H, W, -1)405 return x406 407 408class WindowAttention(nn.Module):409 def __init__(self, dim, num_heads, window_size, qkv_bias=True):410 411 super().__init__()412 self.dim = dim413 self.window_size = window_size414 self.num_heads = num_heads415 head_dim = dim // num_heads416 self.scale = float(head_dim) ** -0.5417 418 self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)419 self.proj = nn.Linear(dim, dim)420 421 self.softmax = nn.Softmax(dim=-1)422 423 def forward(self, x, size):424 425 H, W = size426 B, L, C = x.shape427 assert L == H * W, "input feature has wrong size"428 429 x = x.view(B, H, W, C)430 431 pad_l = pad_t = 0432 pad_r = (self.window_size - W % self.window_size) % self.window_size433 pad_b = (self.window_size - H % self.window_size) % self.window_size434 x = F.pad(x, (0, 0, pad_l, pad_r, pad_t, pad_b))435 _, Hp, Wp, _ = x.shape436 437 x = window_partition(x, self.window_size)438 x = x.view(-1, self.window_size * self.window_size, C)439 440 # W-MSA/SW-MSA441 # attn_windows = self.attn(x_windows)442 443 B_, N, C = x.shape444 qkv = self.qkv(x).reshape(B_, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)445 q, k, v = qkv[0], qkv[1], qkv[2]446 447 q = q * self.scale448 attn = (q @ k.transpose(-2, -1))449 attn = self.softmax(attn)450 451 x = (attn @ v).transpose(1, 2).reshape(B_, N, C)452 x = self.proj(x)453 454 # merge windows455 x = x.view(456 -1, self.window_size, self.window_size, C457 )458 x = window_reverse(x, B, self.window_size, Hp, Wp)459 460 if pad_r > 0 or pad_b > 0:461 x = x[:, :H, :W, :].contiguous()462 463 x = x.view(B, H * W, C)464 465 return x, size466 467 468class SpatialBlock(nn.Module):469 470 def __init__(self, dim, num_heads, window_size,471 mlp_ratio=4., qkv_bias=True, drop_path_rate=0., act_layer=nn.GELU,472 norm_layer=nn.LayerNorm, conv_at_attn=True, conv_at_ffn=True):473 super().__init__()474 475 drop_path = DropPath(drop_path_rate) if drop_path_rate > 0. else nn.Identity()476 477 self.conv1 = PreNorm(None, DepthWiseConv2d(dim, 3, 1, 1)) if conv_at_attn else None478 self.window_attn = PreNorm(479 norm_layer(dim),480 WindowAttention(dim, num_heads, window_size, qkv_bias=qkv_bias),481 drop_path482 )483 self.conv2 = PreNorm(None, DepthWiseConv2d(dim, 3, 1, 1)) if conv_at_ffn else None484 self.ffn = PreNorm(485 norm_layer(dim),486 Mlp(in_features=dim, hidden_features=int(dim*mlp_ratio), act_layer=act_layer),487 drop_path488 )489 490 def forward(self, x, size):491 if self.conv1:492 x, size = self.conv1(x, size)493 x, size = self.window_attn(x, size)494 495 if self.conv2:496 x, size = self.conv2(x, size)497 x, size = self.ffn(x, size)498 return x, size499 500 501class DaViT(nn.Module):502 """ DaViT: Dual-Attention Transformer503 504 Args:505 in_chans (int): Number of input image channels. Default: 3.506 num_classes (int): Number of classes for classification head. Default: 1000.507 patch_size (tuple(int)): Patch size of convolution in different stages. Default: (7, 2, 2, 2).508 patch_stride (tuple(int)): Patch stride of convolution in different stages. Default: (4, 2, 2, 2).509 patch_padding (tuple(int)): Patch padding of convolution in different stages. Default: (3, 0, 0, 0).510 patch_prenorm (tuple(bool)): If True, perform norm before convlution layer. Default: (True, False, False, False).511 embed_dims (tuple(int)): Patch embedding dimension in different stages. Default: (64, 128, 192, 256).512 num_heads (tuple(int)): Number of spatial attention heads in different stages. Default: (4, 8, 12, 16).513 num_groups (tuple(int)): Number of channel groups in different stages. Default: (4, 8, 12, 16).514 window_size (int): Window size. Default: 7.515 mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. Default: 4.516 qkv_bias (bool): If True, add a learnable bias to query, key, value. Default: True.517 drop_path_rate (float): Stochastic depth rate. Default: 0.1.518 norm_layer (nn.Module): Normalization layer. Default: nn.LayerNorm.519 enable_checkpoint (bool): If True, enable checkpointing. Default: False.520 conv_at_attn (bool): If True, performe depthwise convolution before attention layer. Default: True.521 conv_at_ffn (bool): If True, performe depthwise convolution before ffn layer. Default: True.522 """523 524 def __init__(525 self,526 in_chans=3,527 num_classes=1000,528 depths=(1, 1, 3, 1),529 patch_size=(7, 2, 2, 2),530 patch_stride=(4, 2, 2, 2),531 patch_padding=(3, 0, 0, 0),532 patch_prenorm=(False, False, False, False),533 embed_dims=(64, 128, 192, 256),534 num_heads=(3, 6, 12, 24),535 num_groups=(3, 6, 12, 24),536 window_size=7,537 mlp_ratio=4.,538 qkv_bias=True,539 drop_path_rate=0.1,540 norm_layer=nn.LayerNorm,541 enable_checkpoint=False,542 conv_at_attn=True,543 conv_at_ffn=True,544 ):545 super().__init__()546 547 self.num_classes = num_classes548 self.embed_dims = embed_dims549 self.num_heads = num_heads550 self.num_groups = num_groups551 self.num_stages = len(self.embed_dims)552 self.enable_checkpoint = enable_checkpoint553 assert self.num_stages == len(self.num_heads) == len(self.num_groups)554 555 num_stages = len(embed_dims)556 dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths)*2)]557 558 depth_offset = 0559 convs = []560 blocks = []561 for i in range(num_stages):562 conv_embed = ConvEmbed(563 patch_size=patch_size[i],564 stride=patch_stride[i],565 padding=patch_padding[i],566 in_chans=in_chans if i == 0 else self.embed_dims[i - 1],567 embed_dim=self.embed_dims[i],568 norm_layer=norm_layer,569 pre_norm=patch_prenorm[i]570 )571 convs.append(conv_embed)572 573 block = MySequential(574 *[575 MySequential(OrderedDict([576 (577 'spatial_block', SpatialBlock(578 embed_dims[i],579 num_heads[i],580 window_size,581 drop_path_rate=dpr[depth_offset+j*2],582 qkv_bias=qkv_bias,583 mlp_ratio=mlp_ratio,584 conv_at_attn=conv_at_attn,585 conv_at_ffn=conv_at_ffn,586 )587 ),588 (589 'channel_block', ChannelBlock(590 embed_dims[i],591 num_groups[i],592 drop_path_rate=dpr[depth_offset+j*2+1],593 qkv_bias=qkv_bias,594 mlp_ratio=mlp_ratio,595 conv_at_attn=conv_at_attn,596 conv_at_ffn=conv_at_ffn,597 )598 )599 ])) for j in range(depths[i])600 ]601 )602 blocks.append(block)603 depth_offset += depths[i]*2604 605 self.convs = nn.ModuleList(convs)606 self.blocks = nn.ModuleList(blocks)607 608 self.norms = norm_layer(self.embed_dims[-1])609 self.avgpool = nn.AdaptiveAvgPool1d(1)610 self.head = nn.Linear(self.embed_dims[-1], num_classes) if num_classes > 0 else nn.Identity()611 612 self.apply(self._init_weights)613 614 @property615 def dim_out(self):616 return self.embed_dims[-1]617 618 def _init_weights(self, m):619 if isinstance(m, nn.Linear):620 trunc_normal_(m.weight, std=0.02)621 if m.bias is not None:622 nn.init.constant_(m.bias, 0)623 elif isinstance(m, nn.Conv2d):624 nn.init.normal_(m.weight, std=0.02)625 for name, _ in m.named_parameters():626 if name in ['bias']:627 nn.init.constant_(m.bias, 0)628 elif isinstance(m, nn.LayerNorm):629 nn.init.constant_(m.weight, 1.0)630 nn.init.constant_(m.bias, 0)631 elif isinstance(m, nn.BatchNorm2d):632 nn.init.constant_(m.weight, 1.0)633 nn.init.constant_(m.bias, 0)634 635 def forward_features_unpool(self, x):636 """637 forward until avg pooling 638 Args:639 x (_type_): input image tensor640 """641 input_size = (x.size(2), x.size(3))642 for conv, block in zip(self.convs, self.blocks):643 x, input_size = conv(x, input_size)644 if self.enable_checkpoint:645 x, input_size = checkpoint.checkpoint(block, x, input_size)646 else:647 x, input_size = block(x, input_size)648 return x649 650 def forward_features(self, x):651 x = self.forward_features_unpool(x)652 653 # (batch_size, num_tokens, token_dim)654 x = self.avgpool(x.transpose(1, 2))655 # (batch_size, 1, num_tokens)656 x = torch.flatten(x, 1)657 x = self.norms(x)658 659 return x660 661 def forward(self, x):662 x = self.forward_features(x)663 x = self.head(x)664 return x665 666 @classmethod667 def from_config(cls, config):668 return cls(669 depths=config.depths,670 embed_dims=config.dim_embed,671 num_heads=config.num_heads,672 num_groups=config.num_groups,673 patch_size=config.patch_size,674 patch_stride=config.patch_stride,675 patch_padding=config.patch_padding,676 patch_prenorm=config.patch_prenorm,677 drop_path_rate=config.drop_path_rate,678 window_size=config.window_size,679 )680 681 682 683 684if is_flash_attn_2_available():685 from flash_attn import flash_attn_func, flash_attn_varlen_func686 from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input # noqa687 688# Copied from transformers.models.llama.modeling_llama._get_unpad_data689def _get_unpad_data(attention_mask):690 seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32)691 indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()692 max_seqlen_in_batch = seqlens_in_batch.max().item()693 cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0))694 return (695 indices,696 cu_seqlens,697 max_seqlen_in_batch,698 )699 700 701def shift_tokens_right(input_ids: torch.Tensor, pad_token_id: int, decoder_start_token_id: int):702 """703 Shift input ids one token to the right.704 """705 shifted_input_ids = input_ids.new_zeros(input_ids.shape)706 shifted_input_ids[:, 1:] = input_ids[:, :-1].clone()707 shifted_input_ids[:, 0] = decoder_start_token_id708 709 if pad_token_id is None:710 raise ValueError("self.model.config.pad_token_id has to be defined.")711 # replace possible -100 values in labels by `pad_token_id`712 shifted_input_ids.masked_fill_(shifted_input_ids == -100, pad_token_id)713 714 return shifted_input_ids715 716 717class Florence2LearnedPositionalEmbedding(nn.Embedding):718 """719 This module learns positional embeddings up to a fixed maximum size.720 """721 722 def __init__(self, num_embeddings: int, embedding_dim: int):723 # Florence2 is set up so that if padding_idx is specified then offset the embedding ids by 2724 # and adjust num_embeddings appropriately. Other models don't have this hack725 self.offset = 2726 super().__init__(num_embeddings + self.offset, embedding_dim)727 728 def forward(self, input_ids: torch.Tensor, past_key_values_length: int = 0):729 """`input_ids' shape is expected to be [bsz x seqlen]."""730 731 bsz, seq_len = input_ids.shape[:2]732 positions = torch.arange(733 past_key_values_length, past_key_values_length + seq_len, dtype=torch.long, device=self.weight.device734 ).expand(bsz, -1)735 736 return super().forward(positions + self.offset)737 738 739class Florence2ScaledWordEmbedding(nn.Embedding):740 """741 This module overrides nn.Embeddings' forward by multiplying with embeddings scale.742 """743 744 def __init__(self, num_embeddings: int, embedding_dim: int, padding_idx: int, embed_scale: Optional[float] = 1.0):745 super().__init__(num_embeddings, embedding_dim, padding_idx)746 self.embed_scale = embed_scale747 748 def forward(self, input_ids: torch.Tensor):749 return super().forward(input_ids) * self.embed_scale750 751 752class Florence2Attention(nn.Module):753 """Multi-headed attention from 'Attention Is All You Need' paper"""754 755 def __init__(756 self,757 embed_dim: int,758 num_heads: int,759 dropout: float = 0.0,760 is_decoder: bool = False,761 bias: bool = True,762 is_causal: bool = False,763 config: Optional[Florence2LanguageConfig] = None,764 ):765 super().__init__()766 self.embed_dim = embed_dim767 self.num_heads = num_heads768 self.dropout = dropout769 self.head_dim = embed_dim // num_heads770 self.config = config771 772 if (self.head_dim * num_heads) != self.embed_dim:773 raise ValueError(774 f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim}"775 f" and `num_heads`: {num_heads})."776 )777 self.scaling = self.head_dim**-0.5778 self.is_decoder = is_decoder779 self.is_causal = is_causal780 781 self.k_proj = nn.Linear(embed_dim, embed_dim, bias=bias)782 self.v_proj = nn.Linear(embed_dim, embed_dim, bias=bias)783 self.q_proj = nn.Linear(embed_dim, embed_dim, bias=bias)784 self.out_proj = nn.Linear(embed_dim, embed_dim, bias=bias)785 786 def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int):787 return tensor.view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2).contiguous()788 789 def forward(790 self,791 hidden_states: torch.Tensor,792 key_value_states: Optional[torch.Tensor] = None,793 past_key_value: Optional[Tuple[torch.Tensor]] = None,794 attention_mask: Optional[torch.Tensor] = None,795 layer_head_mask: Optional[torch.Tensor] = None,796 output_attentions: bool = False,797 ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:798 """Input shape: Batch x Time x Channel"""799 800 # if key_value_states are provided this layer is used as a cross-attention layer801 # for the decoder802 is_cross_attention = key_value_states is not None803 804 bsz, tgt_len, _ = hidden_states.size()805 806 # get query proj807 query_states = self.q_proj(hidden_states) * self.scaling808 # get key, value proj809 # `past_key_value[0].shape[2] == key_value_states.shape[1]`810 # is checking that the `sequence_length` of the `past_key_value` is the same as811 # the provided `key_value_states` to support prefix tuning812 if (813 is_cross_attention814 and past_key_value is not None815 and past_key_value[0].shape[2] == key_value_states.shape[1]816 ):817 # reuse k,v, cross_attentions818 key_states = past_key_value[0]819 value_states = past_key_value[1]820 elif is_cross_attention:821 # cross_attentions822 key_states = self._shape(self.k_proj(key_value_states), -1, bsz)823 value_states = self._shape(self.v_proj(key_value_states), -1, bsz)824 elif past_key_value is not None:825 # reuse k, v, self_attention826 key_states = self._shape(self.k_proj(hidden_states), -1, bsz)827 value_states = self._shape(self.v_proj(hidden_states), -1, bsz)828 key_states = torch.cat([past_key_value[0], key_states], dim=2)829 value_states = torch.cat([past_key_value[1], value_states], dim=2)830 else:831 # self_attention832 key_states = self._shape(self.k_proj(hidden_states), -1, bsz)833 value_states = self._shape(self.v_proj(hidden_states), -1, bsz)834 835 if self.is_decoder:836 # if cross_attention save Tuple(torch.Tensor, torch.Tensor) of all cross attention key/value_states.837 # Further calls to cross_attention layer can then reuse all cross-attention838 # key/value_states (first "if" case)839 # if uni-directional self-attention (decoder) save Tuple(torch.Tensor, torch.Tensor) of840 # all previous decoder key/value_states. Further calls to uni-directional self-attention841 # can concat previous decoder key/value_states to current projected key/value_states (third "elif" case)842 # if encoder bi-directional self-attention `past_key_value` is always `None`843 past_key_value = (key_states, value_states)844 845 proj_shape = (bsz * self.num_heads, -1, self.head_dim)846 query_states = self._shape(query_states, tgt_len, bsz).view(*proj_shape)847 key_states = key_states.reshape(*proj_shape)848 value_states = value_states.reshape(*proj_shape)849 850 src_len = key_states.size(1)851 attn_weights = torch.bmm(query_states, key_states.transpose(1, 2))852 853 if attn_weights.size() != (bsz * self.num_heads, tgt_len, src_len):854 raise ValueError(855 f"Attention weights should be of size {(bsz * self.num_heads, tgt_len, src_len)}, but is"856 f" {attn_weights.size()}"857 )858 859 if attention_mask is not None:860 if attention_mask.size() != (bsz, 1, tgt_len, src_len):861 raise ValueError(862 f"Attention mask should be of size {(bsz, 1, tgt_len, src_len)}, but is {attention_mask.size()}"863 )864 attn_weights = attn_weights.view(bsz, self.num_heads, tgt_len, src_len) + attention_mask865 attn_weights = attn_weights.view(bsz * self.num_heads, tgt_len, src_len)866 867 attn_weights = nn.functional.softmax(attn_weights, dim=-1)868 869 if layer_head_mask is not None:870 if layer_head_mask.size() != (self.num_heads,):871 raise ValueError(872 f"Head mask for a single layer should be of size {(self.num_heads,)}, but is"873 f" {layer_head_mask.size()}"874 )875 attn_weights = layer_head_mask.view(1, -1, 1, 1) * attn_weights.view(bsz, self.num_heads, tgt_len, src_len)876 attn_weights = attn_weights.view(bsz * self.num_heads, tgt_len, src_len)877 878 if output_attentions:879 # this operation is a bit awkward, but it's required to880 # make sure that attn_weights keeps its gradient.881 # In order to do so, attn_weights have to be reshaped882 # twice and have to be reused in the following883 attn_weights_reshaped = attn_weights.view(bsz, self.num_heads, tgt_len, src_len)884 attn_weights = attn_weights_reshaped.view(bsz * self.num_heads, tgt_len, src_len)885 else:886 attn_weights_reshaped = None887 888 attn_probs = nn.functional.dropout(attn_weights, p=self.dropout, training=self.training)889 890 attn_output = torch.bmm(attn_probs, value_states)891 892 if attn_output.size() != (bsz * self.num_heads, tgt_len, self.head_dim):893 raise ValueError(894 f"`attn_output` should be of size {(bsz * self.num_heads, tgt_len, self.head_dim)}, but is"895 f" {attn_output.size()}"896 )897 898 attn_output = attn_output.view(bsz, self.num_heads, tgt_len, self.head_dim)899 attn_output = attn_output.transpose(1, 2)900 901 # Use the `embed_dim` from the config (stored in the class) rather than `hidden_state` because `attn_output` can be902 # partitioned across GPUs when using tensor-parallelism.903 attn_output = attn_output.reshape(bsz, tgt_len, self.embed_dim)904 905 attn_output = self.out_proj(attn_output)906 907 return attn_output, attn_weights_reshaped, past_key_value908 909 910class Florence2FlashAttention2(Florence2Attention):911 """912 Florence2 flash attention module. This module inherits from `Florence2Attention` as the weights of the module stays913 untouched. The only required change would be on the forward pass where it needs to correctly call the public API of914 flash attention and deal with padding tokens in case the input contains any of them.915 """916 917 # Copied from transformers.models.llama.modeling_llama.LlamaFlashAttention2.__init__918 def __init__(self, *args, **kwargs):919 super().__init__(*args, **kwargs)920 921 # TODO: Should be removed once Flash Attention for RoCm is bumped to 2.1.922 # 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.923 # 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).924 self._flash_attn_uses_top_left_mask = not is_flash_attn_greater_or_equal_2_10()925 926 def _reshape(self, tensor: torch.Tensor, seq_len: int, bsz: int):927 return tensor.view(bsz, seq_len, self.num_heads, self.head_dim)928 929 def forward(930 self,931 hidden_states: torch.Tensor,932 key_value_states: Optional[torch.Tensor] = None,933 past_key_value: Optional[Tuple[torch.Tensor]] = None,934 attention_mask: Optional[torch.Tensor] = None,935 layer_head_mask: Optional[torch.Tensor] = None,936 output_attentions: bool = False,937 ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:938 # Florence2FlashAttention2 attention does not support output_attentions939 if output_attentions:940 raise ValueError("Florence2FlashAttention2 attention does not support output_attentions")941 942 # if key_value_states are provided this layer is used as a cross-attention layer943 # for the decoder944 is_cross_attention = key_value_states is not None945 946 bsz, q_len, _ = hidden_states.size()947 948 # get query proj949 query_states = self._reshape(self.q_proj(hidden_states), -1, bsz)950 # get key, value proj951 # `past_key_value[0].shape[2] == key_value_states.shape[1]`952 # is checking that the `sequence_length` of the `past_key_value` is the same as953 # the provided `key_value_states` to support prefix tuning954 if (955 is_cross_attention956 and past_key_value is not None957 and past_key_value[0].shape[2] == key_value_states.shape[1]958 ):959 # reuse k,v, cross_attentions960 key_states = past_key_value[0].transpose(1, 2)961 value_states = past_key_value[1].transpose(1, 2)962 elif is_cross_attention:963 # cross_attentions964 key_states = self._reshape(self.k_proj(key_value_states), -1, bsz)965 value_states = self._reshape(self.v_proj(key_value_states), -1, bsz)966 elif past_key_value is not None:967 # reuse k, v, self_attention968 key_states = self._reshape(self.k_proj(hidden_states), -1, bsz)969 value_states = self._reshape(self.v_proj(hidden_states), -1, bsz)970 key_states = torch.cat([past_key_value[0].transpose(1, 2), key_states], dim=1)971 value_states = torch.cat([past_key_value[1].transpose(1, 2), value_states], dim=1)972 else:973 # self_attention974 key_states = self._reshape(self.k_proj(hidden_states), -1, bsz)975 value_states = self._reshape(self.v_proj(hidden_states), -1, bsz)976 977 if self.is_decoder:978 # if cross_attention save Tuple(torch.Tensor, torch.Tensor) of all cross attention key/value_states.979 # Further calls to cross_attention layer can then reuse all cross-attention980 # key/value_states (first "if" case)981 # if uni-directional self-attention (decoder) save Tuple(torch.Tensor, torch.Tensor) of982 # all previous decoder key/value_states. Further calls to uni-directional self-attention983 # can concat previous decoder key/value_states to current projected key/value_states (third "elif" case)984 # if encoder bi-directional self-attention `past_key_value` is always `None`985 past_key_value = (key_states.transpose(1, 2), value_states.transpose(1, 2))986 987 kv_seq_len = key_states.shape[-2]988 if past_key_value is not None:989 kv_seq_len += past_key_value[0].shape[-2]990 991 # In PEFT, usually we cast the layer norms in float32 for training stability reasons992 # therefore the input hidden states gets silently casted in float32. Hence, we need993 # cast them back in the correct dtype just to be sure everything works as expected.994 # This might slowdown training & inference so it is recommended to not cast the LayerNorms995 # in fp32. (LlamaRMSNorm handles it correctly)996 997 input_dtype = query_states.dtype998 if input_dtype == torch.float32:999 if torch.is_autocast_enabled():1000 target_dtype = torch.get_autocast_gpu_dtype()1001 # Handle the case where the model is quantized1002 elif hasattr(self.config, "_pre_quantization_dtype"):1003 target_dtype = self.config._pre_quantization_dtype1004 else:1005 target_dtype = self.q_proj.weight.dtype1006 1007 logger.warning_once(1008 f"The input hidden states seems to be silently casted in float32, this might be related to"1009 f" the fact you have upcasted embedding or layer norm layers in float32. We will cast back the input in"1010 f" {target_dtype}."1011 )1012 1013 query_states = query_states.to(target_dtype)1014 key_states = key_states.to(target_dtype)1015 value_states = value_states.to(target_dtype)1016 1017 attn_output = self._flash_attention_forward(1018 query_states, key_states, value_states, attention_mask, q_len, dropout=self.dropout1019 )1020 1021 attn_output = attn_output.reshape(bsz, q_len, -1)1022 attn_output = self.out_proj(attn_output)1023 1024 if not output_attentions:1025 attn_weights = None1026 1027 return attn_output, attn_weights, past_key_value1028 1029 # Copied from transformers.models.llama.modeling_llama.LlamaFlashAttention2._flash_attention_forward1030 def _flash_attention_forward(1031 self, query_states, key_states, value_states, attention_mask, query_length, dropout=0.0, softmax_scale=None1032 ):1033 """1034 Calls the forward method of Flash Attention - if the input hidden states contain at least one padding token1035 first unpad the input, then computes the attention scores and pad the final attention scores.1036 1037 Args:1038 query_states (`torch.Tensor`):1039 Input query states to be passed to Flash Attention API1040 key_states (`torch.Tensor`):1041 Input key states to be passed to Flash Attention API1042 value_states (`torch.Tensor`):1043 Input value states to be passed to Flash Attention API1044 attention_mask (`torch.Tensor`):1045 The padding mask - corresponds to a tensor of size `(batch_size, seq_len)` where 0 stands for the1046 position of padding tokens and 1 for the position of non-padding tokens.1047 dropout (`float`):1048 Attention dropout1049 softmax_scale (`float`, *optional*):1050 The scaling of QK^T before applying softmax. Default to 1 / sqrt(head_dim)1051 """1052 if not self._flash_attn_uses_top_left_mask:1053 causal = self.is_causal1054 else:1055 # 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__.1056 causal = self.is_causal and query_length != 11057 1058 # Contains at least one padding token in the sequence1059 if attention_mask is not None:1060 batch_size = query_states.shape[0]1061 query_states, key_states, value_states, indices_q, cu_seq_lens, max_seq_lens = self._upad_input(1062 query_states, key_states, value_states, attention_mask, query_length1063 )1064 1065 cu_seqlens_q, cu_seqlens_k = cu_seq_lens1066 max_seqlen_in_batch_q, max_seqlen_in_batch_k = max_seq_lens1067 1068 attn_output_unpad = flash_attn_varlen_func(1069 query_states,1070 key_states,1071 value_states,1072 cu_seqlens_q=cu_seqlens_q,1073 cu_seqlens_k=cu_seqlens_k,1074 max_seqlen_q=max_seqlen_in_batch_q,1075 max_seqlen_k=max_seqlen_in_batch_k,1076 dropout_p=dropout,1077 softmax_scale=softmax_scale,1078 causal=causal,1079 )1080 1081 attn_output = pad_input(attn_output_unpad, indices_q, batch_size, query_length)1082 else:1083 attn_output = flash_attn_func(1084 query_states, key_states, value_states, dropout, softmax_scale=softmax_scale, causal=causal1085 )1086 1087 return attn_output1088 1089 # Copied from transformers.models.llama.modeling_llama.LlamaFlashAttention2._upad_input1090 def _upad_input(self, query_layer, key_layer, value_layer, attention_mask, query_length):1091 indices_k, cu_seqlens_k, max_seqlen_in_batch_k = _get_unpad_data(attention_mask)1092 batch_size, kv_seq_len, num_key_value_heads, head_dim = key_layer.shape1093 1094 key_layer = index_first_axis(1095 key_layer.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim), indices_k1096 )1097 value_layer = index_first_axis(1098 value_layer.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim), indices_k1099 )1100 if query_length == kv_seq_len:1101 query_layer = index_first_axis(1102 query_layer.reshape(batch_size * kv_seq_len, self.num_heads, head_dim), indices_k1103 )1104 cu_seqlens_q = cu_seqlens_k1105 max_seqlen_in_batch_q = max_seqlen_in_batch_k1106 indices_q = indices_k1107 elif query_length == 1:1108 max_seqlen_in_batch_q = 11109 cu_seqlens_q = torch.arange(1110 batch_size + 1, dtype=torch.int32, device=query_layer.device1111 ) # There is a memcpy here, that is very bad.1112 indices_q = cu_seqlens_q[:-1]1113 query_layer = query_layer.squeeze(1)1114 else:1115 # The -q_len: slice assumes left padding.1116 attention_mask = attention_mask[:, -query_length:]1117 query_layer, indices_q, cu_seqlens_q, max_seqlen_in_batch_q = unpad_input(query_layer, attention_mask)1118 1119 return (1120 query_layer,1121 key_layer,1122 value_layer,1123 indices_q,1124 (cu_seqlens_q, cu_seqlens_k),1125 (max_seqlen_in_batch_q, max_seqlen_in_batch_k),1126 )1127 1128 1129class Florence2SdpaAttention(Florence2Attention):1130 def forward(1131 self,1132 hidden_states: torch.Tensor,1133 key_value_states: Optional[torch.Tensor] = None,1134 past_key_value: Optional[Tuple[torch.Tensor]] = None,1135 attention_mask: Optional[torch.Tensor] = None,1136 layer_head_mask: Optional[torch.Tensor] = None,1137 output_attentions: bool = False,1138 ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:1139 """Input shape: Batch x Time x Channel"""1140 if output_attentions or layer_head_mask is not None:1141 # TODO: Improve this warning with e.g. `model.config._attn_implementation = "manual"` once this is implemented.1142 logger.warning_once(1143 "Florence2Model is using Florence2SdpaAttention, but `torch.nn.functional.scaled_dot_product_attention` does not support `output_attentions=True` or `layer_head_mask` not None. Falling back to the manual attention"1144 ' implementation, 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.'1145 )1146 return super().forward(1147 hidden_states,1148 key_value_states=key_value_states,1149 past_key_value=past_key_value,1150 attention_mask=attention_mask,1151 layer_head_mask=layer_head_mask,1152 output_attentions=output_attentions,1153 )1154 1155 # if key_value_states are provided this layer is used as a cross-attention layer1156 # for the decoder1157 is_cross_attention = key_value_states is not None1158 1159 bsz, tgt_len, _ = hidden_states.size()1160 1161 # get query proj1162 query_states = self.q_proj(hidden_states)1163 # get key, value proj1164 # `past_key_value[0].shape[2] == key_value_states.shape[1]`1165 # is checking that the `sequence_length` of the `past_key_value` is the same as1166 # the provided `key_value_states` to support prefix tuning1167 if (1168 is_cross_attention1169 and past_key_value is not None1170 and past_key_value[0].shape[2] == key_value_states.shape[1]1171 ):1172 # reuse k,v, cross_attentions1173 key_states = past_key_value[0]1174 value_states = past_key_value[1]1175 elif is_cross_attention:1176 # cross_attentions1177 key_states = self._shape(self.k_proj(key_value_states), -1, bsz)1178 value_states = self._shape(self.v_proj(key_value_states), -1, bsz)1179 elif past_key_value is not None:1180 # reuse k, v, self_attention1181 key_states = self._shape(self.k_proj(hidden_states), -1, bsz)1182 value_states = self._shape(self.v_proj(hidden_states), -1, bsz)1183 key_states = torch.cat([past_key_value[0], key_states], dim=2)1184 value_states = torch.cat([past_key_value[1], value_states], dim=2)1185 else:1186 # self_attention1187 key_states = self._shape(self.k_proj(hidden_states), -1, bsz)1188 value_states = self._shape(self.v_proj(hidden_states), -1, bsz)1189 1190 if self.is_decoder:1191 # if cross_attention save Tuple(torch.Tensor, torch.Tensor) of all cross attention key/value_states.1192 # Further calls to cross_attention layer can then reuse all cross-attention1193 # key/value_states (first "if" case)1194 # if uni-directional self-attention (decoder) save Tuple(torch.Tensor, torch.Tensor) of1195 # all previous decoder key/value_states. Further calls to uni-directional self-attention1196 # can concat previous decoder key/value_states to current projected key/value_states (third "elif" case)1197 # if encoder bi-directional self-attention `past_key_value` is always `None`1198 past_key_value = (key_states, value_states)1199 1200 query_states = self._shape(query_states, tgt_len, bsz)