microsoft/Florence-2-base
4013.1m
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.layers import DropPath, trunc_normal_30 31from transformers.modeling_utils import PreTrainedModel32from transformers.generation.utils import GenerationMixin33from transformers.utils import (34 ModelOutput,35 add_start_docstrings,36 add_start_docstrings_to_model_forward,37 is_flash_attn_2_available,38 logging,39 replace_return_docstrings,40 is_flash_attn_2_available,41 is_flash_attn_greater_or_equal_2_10,42)43from .configuration_florence2 import Florence2Config 44from .configuration_florence2 import Florence2LanguageConfig45from .configuration_florence2 import Florence2VisionConfig46 47 48from transformers.activations import ACT2FN49from transformers.modeling_attn_mask_utils import (50 _prepare_4d_attention_mask,51 _prepare_4d_attention_mask_for_sdpa,52 _prepare_4d_causal_attention_mask,53 _prepare_4d_causal_attention_mask_for_sdpa,54)55from transformers.modeling_outputs import (56 BaseModelOutput,57 BaseModelOutputWithPastAndCrossAttentions,58 Seq2SeqLMOutput,59 Seq2SeqModelOutput,60)61 62 63if is_flash_attn_2_available():64 from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input # noqa65 66logger = logging.get_logger(__name__)67 68_CONFIG_FOR_DOC = "Florence2Config"69 70class LearnedAbsolutePositionEmbedding2D(nn.Module):71 """72 This module learns positional embeddings up to a fixed maximum size.73 """74 75 def __init__(self, embedding_dim=256, num_pos=50):76 super().__init__()77 self.row_embeddings = nn.Embedding(num_pos, embedding_dim // 2)78 self.column_embeddings = nn.Embedding(num_pos, embedding_dim - (embedding_dim // 2))79 80 def forward(self, pixel_values):81 """82 pixel_values: (batch_size, height, width, num_channels) 83 returns: (batch_size, height, width, embedding_dim * 2)84 """85 if len(pixel_values.shape) != 4:86 raise ValueError('pixel_values must be a 4D tensor')87 height, width = pixel_values.shape[1:3]88 width_values = torch.arange(width, device=pixel_values.device)89 height_values = torch.arange(height, device=pixel_values.device)90 x_emb = self.column_embeddings(width_values)91 y_emb = self.row_embeddings(height_values)92 # (height, width, embedding_dim * 2)93 pos = torch.cat([x_emb.unsqueeze(0).repeat(height, 1, 1), y_emb.unsqueeze(1).repeat(1, width, 1)], dim=-1)94 # (embedding_dim * 2, height, width)95 pos = pos.permute(2, 0, 1)96 pos = pos.unsqueeze(0)97 # (batch_size, embedding_dim * 2, height, width)98 pos = pos.repeat(pixel_values.shape[0], 1, 1, 1)99 # (batch_size, height, width, embedding_dim * 2)100 pos = pos.permute(0, 2, 3, 1)101 return pos102 103class PositionalEmbeddingCosine1D(nn.Module):104 """105 This class implements a very simple positional encoding. It follows closely106 the encoder from the link below:107 https://pytorch.org/tutorials/beginner/translation_transformer.html108 109 Args:110 embed_dim: The dimension of the embeddings.111 dropout_prob: The dropout probability.112 max_seq_len: The maximum length to precompute the positional encodings.113 """114 def __init__(115 self,116 embed_dim: int = 512,117 max_seq_len: int = 1024) -> None:118 super(PositionalEmbeddingCosine1D, self).__init__()119 self.embed_dim = embed_dim120 self.max_seq_len = max_seq_len121 # Generate the sinusoidal arrays.122 factor = math.log(10000)123 denominator = torch.exp(124 -factor * torch.arange(0, self.embed_dim, 2) / self.embed_dim)125 # Matrix where rows correspond to a positional embedding as a function126 # of the position index (i.e., the row index).127 frequencies = \128 torch.arange(0, self.max_seq_len) \129 .reshape(self.max_seq_len, 1) * denominator130 pos_idx_to_embed = torch.zeros((self.max_seq_len, self.embed_dim))131 # Populate uneven entries.132 pos_idx_to_embed[:, 0::2] = torch.sin(frequencies)133 pos_idx_to_embed[:, 1::2] = torch.cos(frequencies)134 # Save the positional embeddings in a constant buffer.135 self.register_buffer("pos_idx_to_embed", pos_idx_to_embed)136 137 def forward(self, seq_embeds: torch.Tensor) -> torch.Tensor:138 """139 Args:140 seq_embeds: The sequence embeddings in order. Allowed size:141 1. [T, D], where T is the length of the sequence, and D is the142 frame embedding dimension.143 2. [B, T, D], where B is the batch size and T and D are the144 same as above.145 146 Returns a tensor of with the same dimensions as the input: i.e.,147 [1, T, D] or [T, D].148 """149 shape_len = len(seq_embeds.shape)150 assert 2 <= shape_len <= 3151 len_seq = seq_embeds.size(-2)152 assert len_seq <= self.max_seq_len153 pos_embeds = self.pos_idx_to_embed[0:seq_embeds.size(-2), :]154 # Adapt pre-computed positional embeddings to the input.155 if shape_len == 3:156 pos_embeds = pos_embeds.view(157 (1, pos_embeds.size(0), pos_embeds.size(1)))158 return pos_embeds159 160 161class LearnedAbsolutePositionEmbedding1D(nn.Module):162 """163 Learnable absolute positional embeddings for 1D sequences.164 165 Args:166 embed_dim: The dimension of the embeddings.167 max_seq_len: The maximum length to precompute the positional encodings.168 """169 def __init__(170 self,171 embedding_dim: int = 512,172 num_pos: int = 1024) -> None:173 super(LearnedAbsolutePositionEmbedding1D, self).__init__()174 self.embeddings = nn.Embedding(num_pos, embedding_dim)175 self.num_pos = num_pos176 177 def forward(self, seq_embeds: torch.Tensor) -> torch.Tensor:178 """179 Args:180 seq_embeds: The sequence embeddings in order. Allowed size:181 1. [T, D], where T is the length of the sequence, and D is the182 frame embedding dimension.183 2. [B, T, D], where B is the batch size and T and D are the184 same as above.185 186 Returns a tensor of with the same dimensions as the input: i.e.,187 [1, T, D] or [T, D].188 """189 shape_len = len(seq_embeds.shape)190 assert 2 <= shape_len <= 3191 len_seq = seq_embeds.size(-2)192 assert len_seq <= self.num_pos193 # [T, D]194 pos_embeds = self.embeddings(torch.arange(len_seq).to(seq_embeds.device))195 # Adapt pre-computed positional embeddings to the input.196 if shape_len == 3:197 pos_embeds = pos_embeds.view(198 (1, pos_embeds.size(0), pos_embeds.size(1)))199 return pos_embeds200 201 202 203class MySequential(nn.Sequential):204 def forward(self, *inputs):205 for module in self._modules.values():206 if type(inputs) == tuple:207 inputs = module(*inputs)208 else:209 inputs = module(inputs)210 return inputs211 212 213class PreNorm(nn.Module):214 def __init__(self, norm, fn, drop_path=None):215 super().__init__()216 self.norm = norm217 self.fn = fn218 self.drop_path = drop_path219 220 def forward(self, x, *args, **kwargs):221 shortcut = x222 if self.norm != None:223 x, size = self.fn(self.norm(x), *args, **kwargs)224 else:225 x, size = self.fn(x, *args, **kwargs)226 227 if self.drop_path:228 x = self.drop_path(x)229 230 x = shortcut + x231 232 return x, size233 234 235class Mlp(nn.Module):236 def __init__(237 self,238 in_features,239 hidden_features=None,240 out_features=None,241 act_layer=nn.GELU,242 ):243 super().__init__()244 out_features = out_features or in_features245 hidden_features = hidden_features or in_features246 self.net = nn.Sequential(OrderedDict([247 ("fc1", nn.Linear(in_features, hidden_features)),248 ("act", act_layer()),249 ("fc2", nn.Linear(hidden_features, out_features))250 ]))251 252 def forward(self, x, size):253 return self.net(x), size254 255 256class DepthWiseConv2d(nn.Module):257 def __init__(258 self,259 dim_in,260 kernel_size,261 padding,262 stride,263 bias=True,264 ):265 super().__init__()266 self.dw = nn.Conv2d(267 dim_in, dim_in,268 kernel_size=kernel_size,269 padding=padding,270 groups=dim_in,271 stride=stride,272 bias=bias273 )274 275 def forward(self, x, size):276 B, N, C = x.shape277 H, W = size278 assert N == H * W279 280 x = self.dw(x.transpose(1, 2).view(B, C, H, W))281 size = (x.size(-2), x.size(-1))282 x = x.flatten(2).transpose(1, 2)283 return x, size284 285 286class ConvEmbed(nn.Module):287 """ Image to Patch Embedding288 """289 290 def __init__(291 self,292 patch_size=7,293 in_chans=3,294 embed_dim=64,295 stride=4,296 padding=2,297 norm_layer=None,298 pre_norm=True299 ):300 super().__init__()301 self.patch_size = patch_size302 303 self.proj = nn.Conv2d(304 in_chans, embed_dim,305 kernel_size=patch_size,306 stride=stride,307 padding=padding308 )309 310 dim_norm = in_chans if pre_norm else embed_dim311 self.norm = norm_layer(dim_norm) if norm_layer else None312 313 self.pre_norm = pre_norm314 315 def forward(self, x, size):316 H, W = size317 if len(x.size()) == 3:318 if self.norm and self.pre_norm:319 x = self.norm(x)320 x = rearrange(321 x, 'b (h w) c -> b c h w',322 h=H, w=W323 )324 325 x = self.proj(x)326 327 _, _, H, W = x.shape328 x = rearrange(x, 'b c h w -> b (h w) c')329 if self.norm and not self.pre_norm:330 x = self.norm(x)331 332 return x, (H, W)333 334 335class ChannelAttention(nn.Module):336 337 def __init__(self, dim, groups=8, qkv_bias=True):338 super().__init__()339 340 self.groups = groups341 self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)342 self.proj = nn.Linear(dim, dim)343 344 def forward(self, x, size):345 B, N, C = x.shape346 347 qkv = self.qkv(x).reshape(B, N, 3, self.groups, C // self.groups).permute(2, 0, 3, 1, 4)348 q, k, v = qkv[0], qkv[1], qkv[2]349 350 q = q * (float(N) ** -0.5)351 attention = q.transpose(-1, -2) @ k352 attention = attention.softmax(dim=-1)353 x = (attention @ v.transpose(-1, -2)).transpose(-1, -2)354 x = x.transpose(1, 2).reshape(B, N, C)355 x = self.proj(x)356 return x, size357 358 359class ChannelBlock(nn.Module):360 361 def __init__(self, dim, groups, mlp_ratio=4., qkv_bias=True,362 drop_path_rate=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm,363 conv_at_attn=True, conv_at_ffn=True):364 super().__init__()365 366 drop_path = DropPath(drop_path_rate) if drop_path_rate > 0. else nn.Identity()367 368 self.conv1 = PreNorm(None, DepthWiseConv2d(dim, 3, 1, 1)) if conv_at_attn else None369 self.channel_attn = PreNorm(370 norm_layer(dim),371 ChannelAttention(dim, groups=groups, qkv_bias=qkv_bias),372 drop_path373 )374 self.conv2 = PreNorm(None, DepthWiseConv2d(dim, 3, 1, 1)) if conv_at_ffn else None375 self.ffn = PreNorm(376 norm_layer(dim),377 Mlp(in_features=dim, hidden_features=int(dim*mlp_ratio), act_layer=act_layer),378 drop_path379 )380 381 def forward(self, x, size):382 if self.conv1:383 x, size = self.conv1(x, size)384 x, size = self.channel_attn(x, size)385 386 if self.conv2:387 x, size = self.conv2(x, size)388 x, size = self.ffn(x, size)389 390 return x, size391 392 393def window_partition(x, window_size: int):394 B, H, W, C = x.shape395 x = x.view(B, H // window_size, window_size, W // window_size, window_size, C)396 windows = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size, window_size, C)397 return windows398 399 400def window_reverse(windows, batch_size: int, window_size: int, H: int, W: int):401 B = batch_size 402 # this will cause onnx conversion failed for dynamic axis, because treated as constant403 # int(windows.shape[0] / (H * W / window_size / window_size)) 404 x = windows.view(B, H // window_size, W // window_size, window_size, window_size, -1)405 x = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, H, W, -1)406 return x407 408 409class WindowAttention(nn.Module):410 def __init__(self, dim, num_heads, window_size, qkv_bias=True):411 412 super().__init__()413 self.dim = dim414 self.window_size = window_size415 self.num_heads = num_heads416 head_dim = dim // num_heads417 self.scale = float(head_dim) ** -0.5418 419 self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)420 self.proj = nn.Linear(dim, dim)421 422 self.softmax = nn.Softmax(dim=-1)423 424 def forward(self, x, size):425 426 H, W = size427 B, L, C = x.shape428 assert L == H * W, "input feature has wrong size"429 430 x = x.view(B, H, W, C)431 432 pad_l = pad_t = 0433 pad_r = (self.window_size - W % self.window_size) % self.window_size434 pad_b = (self.window_size - H % self.window_size) % self.window_size435 x = F.pad(x, (0, 0, pad_l, pad_r, pad_t, pad_b))436 _, Hp, Wp, _ = x.shape437 438 x = window_partition(x, self.window_size)439 x = x.view(-1, self.window_size * self.window_size, C)440 441 # W-MSA/SW-MSA442 # attn_windows = self.attn(x_windows)443 444 B_, N, C = x.shape445 qkv = self.qkv(x).reshape(B_, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)446 q, k, v = qkv[0], qkv[1], qkv[2]447 448 q = q * self.scale449 attn = (q @ k.transpose(-2, -1))450 attn = self.softmax(attn)451 452 x = (attn @ v).transpose(1, 2).reshape(B_, N, C)453 x = self.proj(x)454 455 # merge windows456 x = x.view(457 -1, self.window_size, self.window_size, C458 )459 x = window_reverse(x, B, self.window_size, Hp, Wp)460 461 if pad_r > 0 or pad_b > 0:462 x = x[:, :H, :W, :].contiguous()463 464 x = x.view(B, H * W, C)465 466 return x, size467 468 469class SpatialBlock(nn.Module):470 471 def __init__(self, dim, num_heads, window_size,472 mlp_ratio=4., qkv_bias=True, drop_path_rate=0., act_layer=nn.GELU,473 norm_layer=nn.LayerNorm, conv_at_attn=True, conv_at_ffn=True):474 super().__init__()475 476 drop_path = DropPath(drop_path_rate) if drop_path_rate > 0. else nn.Identity()477 478 self.conv1 = PreNorm(None, DepthWiseConv2d(dim, 3, 1, 1)) if conv_at_attn else None479 self.window_attn = PreNorm(480 norm_layer(dim),481 WindowAttention(dim, num_heads, window_size, qkv_bias=qkv_bias),482 drop_path483 )484 self.conv2 = PreNorm(None, DepthWiseConv2d(dim, 3, 1, 1)) if conv_at_ffn else None485 self.ffn = PreNorm(486 norm_layer(dim),487 Mlp(in_features=dim, hidden_features=int(dim*mlp_ratio), act_layer=act_layer),488 drop_path489 )490 491 def forward(self, x, size):492 if self.conv1:493 x, size = self.conv1(x, size)494 x, size = self.window_attn(x, size)495 496 if self.conv2:497 x, size = self.conv2(x, size)498 x, size = self.ffn(x, size)499 return x, size500 501 502class DaViT(nn.Module):503 """ DaViT: Dual-Attention Transformer504 505 Args:506 in_chans (int): Number of input image channels. Default: 3.507 num_classes (int): Number of classes for classification head. Default: 1000.508 patch_size (tuple(int)): Patch size of convolution in different stages. Default: (7, 2, 2, 2).509 patch_stride (tuple(int)): Patch stride of convolution in different stages. Default: (4, 2, 2, 2).510 patch_padding (tuple(int)): Patch padding of convolution in different stages. Default: (3, 0, 0, 0).511 patch_prenorm (tuple(bool)): If True, perform norm before convlution layer. Default: (True, False, False, False).512 embed_dims (tuple(int)): Patch embedding dimension in different stages. Default: (64, 128, 192, 256).513 num_heads (tuple(int)): Number of spatial attention heads in different stages. Default: (4, 8, 12, 16).514 num_groups (tuple(int)): Number of channel groups in different stages. Default: (4, 8, 12, 16).515 window_size (int): Window size. Default: 7.516 mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. Default: 4.517 qkv_bias (bool): If True, add a learnable bias to query, key, value. Default: True.518 drop_path_rate (float): Stochastic depth rate. Default: 0.1.519 norm_layer (nn.Module): Normalization layer. Default: nn.LayerNorm.520 enable_checkpoint (bool): If True, enable checkpointing. Default: False.521 conv_at_attn (bool): If True, performe depthwise convolution before attention layer. Default: True.522 conv_at_ffn (bool): If True, performe depthwise convolution before ffn layer. Default: True.523 """524 525 def __init__(526 self,527 in_chans=3,528 num_classes=1000,529 depths=(1, 1, 3, 1),530 patch_size=(7, 2, 2, 2),531 patch_stride=(4, 2, 2, 2),532 patch_padding=(3, 0, 0, 0),533 patch_prenorm=(False, False, False, False),534 embed_dims=(64, 128, 192, 256),535 num_heads=(3, 6, 12, 24),536 num_groups=(3, 6, 12, 24),537 window_size=7,538 mlp_ratio=4.,539 qkv_bias=True,540 drop_path_rate=0.1,541 norm_layer=nn.LayerNorm,542 enable_checkpoint=False,543 conv_at_attn=True,544 conv_at_ffn=True,545 ):546 super().__init__()547 548 self.num_classes = num_classes549 self.embed_dims = embed_dims550 self.num_heads = num_heads551 self.num_groups = num_groups552 self.num_stages = len(self.embed_dims)553 self.enable_checkpoint = enable_checkpoint554 assert self.num_stages == len(self.num_heads) == len(self.num_groups)555 556 num_stages = len(embed_dims)557 dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths)*2)]558 559 depth_offset = 0560 convs = []561 blocks = []562 for i in range(num_stages):563 conv_embed = ConvEmbed(564 patch_size=patch_size[i],565 stride=patch_stride[i],566 padding=patch_padding[i],567 in_chans=in_chans if i == 0 else self.embed_dims[i - 1],568 embed_dim=self.embed_dims[i],569 norm_layer=norm_layer,570 pre_norm=patch_prenorm[i]571 )572 convs.append(conv_embed)573 574 block = MySequential(575 *[576 MySequential(OrderedDict([577 (578 'spatial_block', SpatialBlock(579 embed_dims[i],580 num_heads[i],581 window_size,582 drop_path_rate=dpr[depth_offset+j*2],583 qkv_bias=qkv_bias,584 mlp_ratio=mlp_ratio,585 conv_at_attn=conv_at_attn,586 conv_at_ffn=conv_at_ffn,587 )588 ),589 (590 'channel_block', ChannelBlock(591 embed_dims[i],592 num_groups[i],593 drop_path_rate=dpr[depth_offset+j*2+1],594 qkv_bias=qkv_bias,595 mlp_ratio=mlp_ratio,596 conv_at_attn=conv_at_attn,597 conv_at_ffn=conv_at_ffn,598 )599 )600 ])) for j in range(depths[i])601 ]602 )603 blocks.append(block)604 depth_offset += depths[i]*2605 606 self.convs = nn.ModuleList(convs)607 self.blocks = nn.ModuleList(blocks)608 609 self.norms = norm_layer(self.embed_dims[-1])610 self.avgpool = nn.AdaptiveAvgPool1d(1)611 self.head = nn.Linear(self.embed_dims[-1], num_classes) if num_classes > 0 else nn.Identity()612 613 @property614 def dim_out(self):615 return self.embed_dims[-1]616 617 def forward_features_unpool(self, x):618 """619 forward until avg pooling 620 Args:621 x (_type_): input image tensor622 """623 input_size = (x.size(2), x.size(3))624 for conv, block in zip(self.convs, self.blocks):625 x, input_size = conv(x, input_size)626 if self.enable_checkpoint:627 x, input_size = checkpoint.checkpoint(block, x, input_size)628 else:629 x, input_size = block(x, input_size)630 return x631 632 def forward_features(self, x):633 x = self.forward_features_unpool(x)634 635 # (batch_size, num_tokens, token_dim)636 x = self.avgpool(x.transpose(1, 2))637 # (batch_size, 1, num_tokens)638 x = torch.flatten(x, 1)639 x = self.norms(x)640 641 return x642 643 def forward(self, x):644 x = self.forward_features(x)645 x = self.head(x)646 return x647 648 @classmethod649 def from_config(cls, config):650 return cls(651 depths=config.depths,652 embed_dims=config.dim_embed,653 num_heads=config.num_heads,654 num_groups=config.num_groups,655 patch_size=config.patch_size,656 patch_stride=config.patch_stride,657 patch_padding=config.patch_padding,658 patch_prenorm=config.patch_prenorm,659 drop_path_rate=config.drop_path_rate,660 window_size=config.window_size,661 )662 663 664 665 666if is_flash_attn_2_available():667 from flash_attn import flash_attn_func, flash_attn_varlen_func668 from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input # noqa669 670# Copied from transformers.models.llama.modeling_llama._get_unpad_data671def _get_unpad_data(attention_mask):672 seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32)673 indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()674 max_seqlen_in_batch = seqlens_in_batch.max().item()675 cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0))676 return (677 indices,678 cu_seqlens,679 max_seqlen_in_batch,680 )681 682 683def shift_tokens_right(input_ids: torch.Tensor, pad_token_id: int, decoder_start_token_id: int):684 """685 Shift input ids one token to the right.686 """687 shifted_input_ids = input_ids.new_zeros(input_ids.shape)688 shifted_input_ids[:, 1:] = input_ids[:, :-1].clone()689 shifted_input_ids[:, 0] = decoder_start_token_id690 691 if pad_token_id is None:692 raise ValueError("self.model.config.pad_token_id has to be defined.")693 # replace possible -100 values in labels by `pad_token_id`694 shifted_input_ids.masked_fill_(shifted_input_ids == -100, pad_token_id)695 696 return shifted_input_ids697 698 699class Florence2LearnedPositionalEmbedding(nn.Embedding):700 """701 This module learns positional embeddings up to a fixed maximum size.702 """703 704 def __init__(self, num_embeddings: int, embedding_dim: int):705 # Florence2 is set up so that if padding_idx is specified then offset the embedding ids by 2706 # and adjust num_embeddings appropriately. Other models don't have this hack707 self.offset = 2708 super().__init__(num_embeddings + self.offset, embedding_dim)709 710 def forward(self, input_ids: torch.Tensor, past_key_values_length: int = 0):711 """`input_ids' shape is expected to be [bsz x seqlen]."""712 713 bsz, seq_len = input_ids.shape[:2]714 positions = torch.arange(715 past_key_values_length, past_key_values_length + seq_len, dtype=torch.long, device=self.weight.device716 ).expand(bsz, -1)717 718 return super().forward(positions + self.offset)719 720 721class Florence2ScaledWordEmbedding(nn.Embedding):722 """723 This module overrides nn.Embeddings' forward by multiplying with embeddings scale.724 """725 726 def __init__(self, num_embeddings: int, embedding_dim: int, padding_idx: int, embed_scale: Optional[float] = 1.0):727 super().__init__(num_embeddings, embedding_dim, padding_idx)728 self.embed_scale = embed_scale729 730 def forward(self, input_ids: torch.Tensor):731 return super().forward(input_ids) * self.embed_scale732 733 734class Florence2Attention(nn.Module):735 """Multi-headed attention from 'Attention Is All You Need' paper"""736 737 def __init__(738 self,739 embed_dim: int,740 num_heads: int,741 dropout: float = 0.0,742 is_decoder: bool = False,743 bias: bool = True,744 is_causal: bool = False,745 config: Optional[Florence2LanguageConfig] = None,746 ):747 super().__init__()748 self.embed_dim = embed_dim749 self.num_heads = num_heads750 self.dropout = dropout751 self.head_dim = embed_dim // num_heads752 self.config = config753 754 if (self.head_dim * num_heads) != self.embed_dim:755 raise ValueError(756 f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim}"757 f" and `num_heads`: {num_heads})."758 )759 self.scaling = self.head_dim**-0.5760 self.is_decoder = is_decoder761 self.is_causal = is_causal762 763 self.k_proj = nn.Linear(embed_dim, embed_dim, bias=bias)764 self.v_proj = nn.Linear(embed_dim, embed_dim, bias=bias)765 self.q_proj = nn.Linear(embed_dim, embed_dim, bias=bias)766 self.out_proj = nn.Linear(embed_dim, embed_dim, bias=bias)767 768 def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int):769 return tensor.view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2).contiguous()770 771 def forward(772 self,773 hidden_states: torch.Tensor,774 key_value_states: Optional[torch.Tensor] = None,775 past_key_value: Optional[Tuple[torch.Tensor]] = None,776 attention_mask: Optional[torch.Tensor] = None,777 layer_head_mask: Optional[torch.Tensor] = None,778 output_attentions: bool = False,779 ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:780 """Input shape: Batch x Time x Channel"""781 782 # if key_value_states are provided this layer is used as a cross-attention layer783 # for the decoder784 is_cross_attention = key_value_states is not None785 786 bsz, tgt_len, _ = hidden_states.size()787 788 # get query proj789 query_states = self.q_proj(hidden_states) * self.scaling790 # get key, value proj791 # `past_key_value[0].shape[2] == key_value_states.shape[1]`792 # is checking that the `sequence_length` of the `past_key_value` is the same as793 # the provided `key_value_states` to support prefix tuning794 if (795 is_cross_attention796 and past_key_value is not None797 and past_key_value[0].shape[2] == key_value_states.shape[1]798 ):799 # reuse k,v, cross_attentions800 key_states = past_key_value[0]801 value_states = past_key_value[1]802 elif is_cross_attention:803 # cross_attentions804 key_states = self._shape(self.k_proj(key_value_states), -1, bsz)805 value_states = self._shape(self.v_proj(key_value_states), -1, bsz)806 elif past_key_value is not None:807 # reuse k, v, self_attention808 key_states = self._shape(self.k_proj(hidden_states), -1, bsz)809 value_states = self._shape(self.v_proj(hidden_states), -1, bsz)810 key_states = torch.cat([past_key_value[0], key_states], dim=2)811 value_states = torch.cat([past_key_value[1], value_states], dim=2)812 else:813 # self_attention814 key_states = self._shape(self.k_proj(hidden_states), -1, bsz)815 value_states = self._shape(self.v_proj(hidden_states), -1, bsz)816 817 if self.is_decoder:818 # if cross_attention save Tuple(torch.Tensor, torch.Tensor) of all cross attention key/value_states.819 # Further calls to cross_attention layer can then reuse all cross-attention820 # key/value_states (first "if" case)821 # if uni-directional self-attention (decoder) save Tuple(torch.Tensor, torch.Tensor) of822 # all previous decoder key/value_states. Further calls to uni-directional self-attention823 # can concat previous decoder key/value_states to current projected key/value_states (third "elif" case)824 # if encoder bi-directional self-attention `past_key_value` is always `None`825 past_key_value = (key_states, value_states)826 827 proj_shape = (bsz * self.num_heads, -1, self.head_dim)828 query_states = self._shape(query_states, tgt_len, bsz).view(*proj_shape)829 key_states = key_states.reshape(*proj_shape)830 value_states = value_states.reshape(*proj_shape)831 832 src_len = key_states.size(1)833 attn_weights = torch.bmm(query_states, key_states.transpose(1, 2))834 835 if attn_weights.size() != (bsz * self.num_heads, tgt_len, src_len):836 raise ValueError(837 f"Attention weights should be of size {(bsz * self.num_heads, tgt_len, src_len)}, but is"838 f" {attn_weights.size()}"839 )840 841 if attention_mask is not None:842 if attention_mask.size() != (bsz, 1, tgt_len, src_len):843 raise ValueError(844 f"Attention mask should be of size {(bsz, 1, tgt_len, src_len)}, but is {attention_mask.size()}"845 )846 attn_weights = attn_weights.view(bsz, self.num_heads, tgt_len, src_len) + attention_mask847 attn_weights = attn_weights.view(bsz * self.num_heads, tgt_len, src_len)848 849 attn_weights = nn.functional.softmax(attn_weights, dim=-1)850 851 if layer_head_mask is not None:852 if layer_head_mask.size() != (self.num_heads,):853 raise ValueError(854 f"Head mask for a single layer should be of size {(self.num_heads,)}, but is"855 f" {layer_head_mask.size()}"856 )857 attn_weights = layer_head_mask.view(1, -1, 1, 1) * attn_weights.view(bsz, self.num_heads, tgt_len, src_len)858 attn_weights = attn_weights.view(bsz * self.num_heads, tgt_len, src_len)859 860 if output_attentions:861 # this operation is a bit awkward, but it's required to862 # make sure that attn_weights keeps its gradient.863 # In order to do so, attn_weights have to be reshaped864 # twice and have to be reused in the following865 attn_weights_reshaped = attn_weights.view(bsz, self.num_heads, tgt_len, src_len)866 attn_weights = attn_weights_reshaped.view(bsz * self.num_heads, tgt_len, src_len)867 else:868 attn_weights_reshaped = None869 870 attn_probs = nn.functional.dropout(attn_weights, p=self.dropout, training=self.training)871 872 attn_output = torch.bmm(attn_probs, value_states)873 874 if attn_output.size() != (bsz * self.num_heads, tgt_len, self.head_dim):875 raise ValueError(876 f"`attn_output` should be of size {(bsz * self.num_heads, tgt_len, self.head_dim)}, but is"877 f" {attn_output.size()}"878 )879 880 attn_output = attn_output.view(bsz, self.num_heads, tgt_len, self.head_dim)881 attn_output = attn_output.transpose(1, 2)882 883 # Use the `embed_dim` from the config (stored in the class) rather than `hidden_state` because `attn_output` can be884 # partitioned across GPUs when using tensor-parallelism.885 attn_output = attn_output.reshape(bsz, tgt_len, self.embed_dim)886 887 attn_output = self.out_proj(attn_output)888 889 return attn_output, attn_weights_reshaped, past_key_value890 891 892class Florence2FlashAttention2(Florence2Attention):893 """894 Florence2 flash attention module. This module inherits from `Florence2Attention` as the weights of the module stays895 untouched. The only required change would be on the forward pass where it needs to correctly call the public API of896 flash attention and deal with padding tokens in case the input contains any of them.897 """898 899 # Copied from transformers.models.llama.modeling_llama.LlamaFlashAttention2.__init__900 def __init__(self, *args, **kwargs):901 super().__init__(*args, **kwargs)902 903 # TODO: Should be removed once Flash Attention for RoCm is bumped to 2.1.904 # 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.905 # 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).906 self._flash_attn_uses_top_left_mask = not is_flash_attn_greater_or_equal_2_10()907 908 def _reshape(self, tensor: torch.Tensor, seq_len: int, bsz: int):909 return tensor.view(bsz, seq_len, self.num_heads, self.head_dim)910 911 def forward(912 self,913 hidden_states: torch.Tensor,914 key_value_states: Optional[torch.Tensor] = None,915 past_key_value: Optional[Tuple[torch.Tensor]] = None,916 attention_mask: Optional[torch.Tensor] = None,917 layer_head_mask: Optional[torch.Tensor] = None,918 output_attentions: bool = False,919 ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:920 # Florence2FlashAttention2 attention does not support output_attentions921 if output_attentions:922 raise ValueError("Florence2FlashAttention2 attention does not support output_attentions")923 924 # if key_value_states are provided this layer is used as a cross-attention layer925 # for the decoder926 is_cross_attention = key_value_states is not None927 928 bsz, q_len, _ = hidden_states.size()929 930 # get query proj931 query_states = self._reshape(self.q_proj(hidden_states), -1, bsz)932 # get key, value proj933 # `past_key_value[0].shape[2] == key_value_states.shape[1]`934 # is checking that the `sequence_length` of the `past_key_value` is the same as935 # the provided `key_value_states` to support prefix tuning936 if (937 is_cross_attention938 and past_key_value is not None939 and past_key_value[0].shape[2] == key_value_states.shape[1]940 ):941 # reuse k,v, cross_attentions942 key_states = past_key_value[0].transpose(1, 2)943 value_states = past_key_value[1].transpose(1, 2)944 elif is_cross_attention:945 # cross_attentions946 key_states = self._reshape(self.k_proj(key_value_states), -1, bsz)947 value_states = self._reshape(self.v_proj(key_value_states), -1, bsz)948 elif past_key_value is not None:949 # reuse k, v, self_attention950 key_states = self._reshape(self.k_proj(hidden_states), -1, bsz)951 value_states = self._reshape(self.v_proj(hidden_states), -1, bsz)952 key_states = torch.cat([past_key_value[0].transpose(1, 2), key_states], dim=1)953 value_states = torch.cat([past_key_value[1].transpose(1, 2), value_states], dim=1)954 else:955 # self_attention956 key_states = self._reshape(self.k_proj(hidden_states), -1, bsz)957 value_states = self._reshape(self.v_proj(hidden_states), -1, bsz)958 959 if self.is_decoder:960 # if cross_attention save Tuple(torch.Tensor, torch.Tensor) of all cross attention key/value_states.961 # Further calls to cross_attention layer can then reuse all cross-attention962 # key/value_states (first "if" case)963 # if uni-directional self-attention (decoder) save Tuple(torch.Tensor, torch.Tensor) of964 # all previous decoder key/value_states. Further calls to uni-directional self-attention965 # can concat previous decoder key/value_states to current projected key/value_states (third "elif" case)966 # if encoder bi-directional self-attention `past_key_value` is always `None`967 past_key_value = (key_states.transpose(1, 2), value_states.transpose(1, 2))968 969 kv_seq_len = key_states.shape[-2]970 if past_key_value is not None:971 kv_seq_len += past_key_value[0].shape[-2]972 973 # In PEFT, usually we cast the layer norms in float32 for training stability reasons974 # therefore the input hidden states gets silently casted in float32. Hence, we need975 # cast them back in the correct dtype just to be sure everything works as expected.976 # This might slowdown training & inference so it is recommended to not cast the LayerNorms977 # in fp32. (LlamaRMSNorm handles it correctly)978 979 input_dtype = query_states.dtype980 if input_dtype == torch.float32:981 if torch.is_autocast_enabled():982 target_dtype = torch.get_autocast_gpu_dtype()983 # Handle the case where the model is quantized984 elif hasattr(self.config, "_pre_quantization_dtype"):985 target_dtype = self.config._pre_quantization_dtype986 else:987 target_dtype = self.q_proj.weight.dtype988 989 logger.warning_once(990 f"The input hidden states seems to be silently casted in float32, this might be related to"991 f" the fact you have upcasted embedding or layer norm layers in float32. We will cast back the input in"992 f" {target_dtype}."993 )994 995 query_states = query_states.to(target_dtype)996 key_states = key_states.to(target_dtype)997 value_states = value_states.to(target_dtype)998 999 attn_output = self._flash_attention_forward(1000 query_states, key_states, value_states, attention_mask, q_len, dropout=self.dropout1001 )1002 1003 attn_output = attn_output.reshape(bsz, q_len, -1)1004 attn_output = self.out_proj(attn_output)1005 1006 if not output_attentions:1007 attn_weights = None1008 1009 return attn_output, attn_weights, past_key_value1010 1011 # Copied from transformers.models.llama.modeling_llama.LlamaFlashAttention2._flash_attention_forward1012 def _flash_attention_forward(1013 self, query_states, key_states, value_states, attention_mask, query_length, dropout=0.0, softmax_scale=None1014 ):1015 """1016 Calls the forward method of Flash Attention - if the input hidden states contain at least one padding token1017 first unpad the input, then computes the attention scores and pad the final attention scores.1018 1019 Args:1020 query_states (`torch.Tensor`):1021 Input query states to be passed to Flash Attention API1022 key_states (`torch.Tensor`):1023 Input key states to be passed to Flash Attention API1024 value_states (`torch.Tensor`):1025 Input value states to be passed to Flash Attention API1026 attention_mask (`torch.Tensor`):1027 The padding mask - corresponds to a tensor of size `(batch_size, seq_len)` where 0 stands for the1028 position of padding tokens and 1 for the position of non-padding tokens.1029 dropout (`float`):1030 Attention dropout1031 softmax_scale (`float`, *optional*):1032 The scaling of QK^T before applying softmax. Default to 1 / sqrt(head_dim)1033 """1034 if not self._flash_attn_uses_top_left_mask:1035 causal = self.is_causal1036 else:1037 # 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__.1038 causal = self.is_causal and query_length != 11039 1040 # Contains at least one padding token in the sequence1041 if attention_mask is not None:1042 batch_size = query_states.shape[0]1043 query_states, key_states, value_states, indices_q, cu_seq_lens, max_seq_lens = self._upad_input(1044 query_states, key_states, value_states, attention_mask, query_length1045 )1046 1047 cu_seqlens_q, cu_seqlens_k = cu_seq_lens1048 max_seqlen_in_batch_q, max_seqlen_in_batch_k = max_seq_lens1049 1050 attn_output_unpad = flash_attn_varlen_func(1051 query_states,1052 key_states,1053 value_states,1054 cu_seqlens_q=cu_seqlens_q,1055 cu_seqlens_k=cu_seqlens_k,1056 max_seqlen_q=max_seqlen_in_batch_q,1057 max_seqlen_k=max_seqlen_in_batch_k,1058 dropout_p=dropout,1059 softmax_scale=softmax_scale,1060 causal=causal,1061 )1062 1063 attn_output = pad_input(attn_output_unpad, indices_q, batch_size, query_length)1064 else:1065 attn_output = flash_attn_func(1066 query_states, key_states, value_states, dropout, softmax_scale=softmax_scale, causal=causal1067 )1068 1069 return attn_output1070 1071 # Copied from transformers.models.llama.modeling_llama.LlamaFlashAttention2._upad_input1072 def _upad_input(self, query_layer, key_layer, value_layer, attention_mask, query_length):1073 indices_k, cu_seqlens_k, max_seqlen_in_batch_k = _get_unpad_data(attention_mask)1074 batch_size, kv_seq_len, num_key_value_heads, head_dim = key_layer.shape1075 1076 key_layer = index_first_axis(1077 key_layer.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim), indices_k1078 )1079 value_layer = index_first_axis(1080 value_layer.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim), indices_k1081 )1082 if query_length == kv_seq_len:1083 query_layer = index_first_axis(1084 query_layer.reshape(batch_size * kv_seq_len, self.num_heads, head_dim), indices_k1085 )1086 cu_seqlens_q = cu_seqlens_k1087 max_seqlen_in_batch_q = max_seqlen_in_batch_k1088 indices_q = indices_k1089 elif query_length == 1:1090 max_seqlen_in_batch_q = 11091 cu_seqlens_q = torch.arange(1092 batch_size + 1, dtype=torch.int32, device=query_layer.device1093 ) # There is a memcpy here, that is very bad.1094 indices_q = cu_seqlens_q[:-1]1095 query_layer = query_layer.squeeze(1)1096 else:1097 # The -q_len: slice assumes left padding.1098 attention_mask = attention_mask[:, -query_length:]1099 query_layer, indices_q, cu_seqlens_q, max_seqlen_in_batch_q = unpad_input(query_layer, attention_mask)1100 1101 return (1102 query_layer,1103 key_layer,1104 value_layer,1105 indices_q,1106 (cu_seqlens_q, cu_seqlens_k),1107 (max_seqlen_in_batch_q, max_seqlen_in_batch_k),1108 )1109 1110 1111class Florence2SdpaAttention(Florence2Attention):1112 def forward(1113 self,1114 hidden_states: torch.Tensor,1115 key_value_states: Optional[torch.Tensor] = None,1116 past_key_value: Optional[Tuple[torch.Tensor]] = None,1117 attention_mask: Optional[torch.Tensor] = None,1118 layer_head_mask: Optional[torch.Tensor] = None,1119 output_attentions: bool = False,1120 ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:1121 """Input shape: Batch x Time x Channel"""1122 if output_attentions or layer_head_mask is not None:1123 # TODO: Improve this warning with e.g. `model.config._attn_implementation = "manual"` once this is implemented.1124 logger.warning_once(1125 "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"1126 ' 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.'1127 )1128 return super().forward(1129 hidden_states,1130 key_value_states=key_value_states,1131 past_key_value=past_key_value,1132 attention_mask=attention_mask,1133 layer_head_mask=layer_head_mask,1134 output_attentions=output_attentions,1135 )1136 1137 # if key_value_states are provided this layer is used as a cross-attention layer1138 # for the decoder1139 is_cross_attention = key_value_states is not None1140 1141 bsz, tgt_len, _ = hidden_states.size()1142 1143 # get query proj1144 query_states = self.q_proj(hidden_states)1145 # get key, value proj1146 # `past_key_value[0].shape[2] == key_value_states.shape[1]`1147 # is checking that the `sequence_length` of the `past_key_value` is the same as1148 # the provided `key_value_states` to support prefix tuning1149 if (1150 is_cross_attention1151 and past_key_value is not None1152 and past_key_value[0].shape[2] == key_value_states.shape[1]1153 ):1154 # reuse k,v, cross_attentions1155 key_states = past_key_value[0]1156 value_states = past_key_value[1]1157 elif is_cross_attention:1158 # cross_attentions1159 key_states = self._shape(self.k_proj(key_value_states), -1, bsz)1160 value_states = self._shape(self.v_proj(key_value_states), -1, bsz)1161 elif past_key_value is not None:1162 # reuse k, v, self_attention1163 key_states = self._shape(self.k_proj(hidden_states), -1, bsz)1164 value_states = self._shape(self.v_proj(hidden_states), -1, bsz)1165 key_states = torch.cat([past_key_value[0], key_states], dim=2)1166 value_states = torch.cat([past_key_value[1], value_states], dim=2)1167 else:1168 # self_attention1169 key_states = self._shape(self.k_proj(hidden_states), -1, bsz)1170 value_states = self._shape(self.v_proj(hidden_states), -1, bsz)1171 1172 if self.is_decoder:1173 # if cross_attention save Tuple(torch.Tensor, torch.Tensor) of all cross attention key/value_states.1174 # Further calls to cross_attention layer can then reuse all cross-attention1175 # key/value_states (first "if" case)1176 # if uni-directional self-attention (decoder) save Tuple(torch.Tensor, torch.Tensor) of1177 # all previous decoder key/value_states. Further calls to uni-directional self-attention1178 # can concat previous decoder key/value_states to current projected key/value_states (third "elif" case)1179 # if encoder bi-directional self-attention `past_key_value` is always `None`1180 past_key_value = (key_states, value_states)1181 1182 query_states = self._shape(query_states, tgt_len, bsz)1183 1184 # We dispatch to SDPA's Flash Attention or Efficient kernels via this `is_causal` if statement instead of an inline conditional assignment1185 # in SDPA to support both torch.compile's dynamic shapes and full graph options. An inline conditional prevents dynamic shapes from compiling.1186 # The tgt_len > 1 is necessary to match with AttentionMaskConverter.to_causal_4d that does not create a causal mask in case tgt_len == 1.1187 is_causal = True if self.is_causal and attention_mask is None and tgt_len > 1 else False1188 1189 # NOTE: SDPA with memory-efficient backend is currently (torch==2.1.2) bugged when using non-contiguous inputs and a custom attn_mask,1190 # but we are fine here as `_shape` do call `.contiguous()`. Reference: https://github.com/pytorch/pytorch/issues/1125771191 attn_output = torch.nn.functional.scaled_dot_product_attention(1192 query_states,1193 key_states,1194 value_states,1195 attn_mask=attention_mask,1196 dropout_p=self.dropout if self.training else 0.0,1197 is_causal=is_causal,1198 )1199 1200 if attn_output.size() != (bsz, self.num_heads, tgt_len, self.head_dim):