Aluode/PerceptionLabPortable
0
1# coding=utf-82# Copyright 2023 The LAION-AI Team and The HuggingFace 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"""PyTorch CLAP model."""16 17import collections18import math19from dataclasses import dataclass20from typing import Any, Callable, Optional, Union21 22import torch23import torch.nn.functional as F24from torch import nn25 26from ...activations import ACT2FN27from ...modeling_layers import GradientCheckpointingLayer28from ...modeling_outputs import (29 BaseModelOutput,30 BaseModelOutputWithPooling,31 BaseModelOutputWithPoolingAndCrossAttentions,32)33from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel34from ...pytorch_utils import apply_chunking_to_forward, find_pruneable_heads_and_indices, meshgrid, prune_linear_layer35from ...utils import ModelOutput, auto_docstring, can_return_tuple, filter_out_non_signature_kwargs, logging, torch_int36from .configuration_clap import ClapAudioConfig, ClapConfig, ClapTextConfig37 38 39logger = logging.get_logger(__name__)40 41 42# Adapted from: https://github.com/LAION-AI/CLAP/blob/6ad05a971ba0622f6acee8c41993e0d02bbed639/src/open_clip/utils.py#L19143def interpolate(hidden_states, ratio):44 """45 Interpolate data in time domain. This is used to compensate the resolution reduction in downsampling of a CNN.46 47 Args:48 hidden_states (`torch.FloatTensor` of shape (batch_size, time_length, classes_num)):49 Input hidden states50 ratio (`int`):51 The ratio of the length of the output to the length of the input.52 """53 (batch_size, time_length, classes_num) = hidden_states.shape54 upsampled = hidden_states[:, :, None, :].repeat(1, 1, ratio, 1)55 upsampled = upsampled.reshape(batch_size, time_length * ratio, classes_num)56 return upsampled57 58 59# Adapted from https://github.com/LAION-AI/CLAP/blob/6ad05a971ba0622f6acee8c41993e0d02bbed639/src/open_clip/htsat.py#L24960def window_partition(hidden_states, window_size):61 """62 Returns the resized hidden states. The output shape should be `(batch_size * num_windows, window_size, window_size,63 num_channels)`64 65 Args:66 hidden_states (`torch.FloatTensor` of shape `(batch_size, height, width, num_channels)`):67 Input hidden states68 window_size (`int`):69 Window size70 """71 batch_size, height, width, num_channels = hidden_states.shape72 73 hidden_states = hidden_states.view(74 batch_size, height // window_size, window_size, width // window_size, window_size, num_channels75 )76 windows = hidden_states.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size, window_size, num_channels)77 return windows78 79 80# Adapted from https://github.com/LAION-AI/CLAP/blob/6ad05a971ba0622f6acee8c41993e0d02bbed639/src/open_clip/htsat.py#L26381def window_reverse(windows, window_size, height, width):82 """83 Merges windows to produce higher resolution features.84 Args:85 windows (`torch.FloatTensor` of shape `(num_windows * batch_size, window_size, window_size, num_channels)`):86 Input windows87 window_size (`int`):88 Window size89 height (`int`):90 Height of the resized audio91 width (`int`):92 Width of the resized audio93 """94 num_channels = windows.shape[-1]95 windows = windows.view(-1, height // window_size, width // window_size, window_size, window_size, num_channels)96 windows = windows.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, height, width, num_channels)97 return windows98 99 100# Copied from transformers.models.roberta.modeling_roberta.create_position_ids_from_input_ids101def create_position_ids_from_input_ids(input_ids, padding_idx, past_key_values_length=0):102 """103 Replace non-padding symbols with their position numbers. Position numbers begin at padding_idx+1. Padding symbols104 are ignored. This is modified from fairseq's `utils.make_positions`.105 106 Args:107 x: torch.Tensor x:108 109 Returns: torch.Tensor110 """111 # The series of casts and type-conversions here are carefully balanced to both work with ONNX export and XLA.112 mask = input_ids.ne(padding_idx).int()113 incremental_indices = (torch.cumsum(mask, dim=1).type_as(mask) + past_key_values_length) * mask114 return incremental_indices.long() + padding_idx115 116 117# contrastive loss function, adapted from118# https://sachinruk.github.io/blog/pytorch/pytorch%20lightning/loss%20function/gpu/2021/03/07/CLIP.html#CLIP-loss-function119def contrastive_loss(logits: torch.Tensor) -> torch.Tensor:120 labels = torch.arange(len(logits), device=logits.device)121 return nn.functional.cross_entropy(logits, labels)122 123 124@dataclass125@auto_docstring(126 custom_intro="""127 Base class for text model's outputs that also contains a pooling of the last hidden states.128 """129)130# Copied from transformers.models.clip.modeling_clip.CLIPTextModelOutput with CLIP->Clap131class ClapTextModelOutput(ModelOutput):132 r"""133 text_embeds (`torch.FloatTensor` of shape `(batch_size, output_dim)` *optional* returned when model is initialized with `with_projection=True`):134 The text embeddings obtained by applying the projection layer to the pooler_output.135 """136 137 text_embeds: Optional[torch.FloatTensor] = None138 last_hidden_state: Optional[torch.FloatTensor] = None139 hidden_states: Optional[tuple[torch.FloatTensor, ...]] = None140 attentions: Optional[tuple[torch.FloatTensor, ...]] = None141 142 143@dataclass144@auto_docstring(145 custom_intro="""146 ClapAudio model output to mimic the output of the original implementation.147 """148)149class ClapAudioModelOutput(ModelOutput):150 r"""151 audio_embeds (`torch.FloatTensor` of shape `(batch_size, hidden_size)`):152 The Audio embeddings obtained by applying the projection layer to the pooler_output.153 """154 155 audio_embeds: Optional[torch.FloatTensor] = None156 last_hidden_state: Optional[torch.FloatTensor] = None157 hidden_states: Optional[tuple[torch.FloatTensor, ...]] = None158 attentions: Optional[tuple[torch.FloatTensor, ...]] = None159 160 161@dataclass162@auto_docstring163# Copied from transformers.models.clip.modeling_clip.CLIPOutput with CLIP->Clap, vision->audio, Vision->Audio, image->audio164class ClapOutput(ModelOutput):165 r"""166 loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `return_loss` is `True`):167 Contrastive loss for audio-text similarity.168 logits_per_audio (`torch.FloatTensor` of shape `(audio_batch_size, text_batch_size)`):169 The scaled dot product scores between `audio_embeds` and `text_embeds`. This represents the audio-text170 similarity scores.171 logits_per_text (`torch.FloatTensor` of shape `(text_batch_size, audio_batch_size)`):172 The scaled dot product scores between `text_embeds` and `audio_embeds`. This represents the text-audio173 similarity scores.174 text_embeds (`torch.FloatTensor` of shape `(batch_size, output_dim`):175 The text embeddings obtained by applying the projection layer to the pooled output of [`ClapTextModel`].176 audio_embeds (`torch.FloatTensor` of shape `(batch_size, output_dim`):177 The audio embeddings obtained by applying the projection layer to the pooled output of [`ClapAudioModel`].178 text_model_output (`BaseModelOutputWithPooling`):179 The output of the [`ClapTextModel`].180 audio_model_output (`BaseModelOutputWithPooling`):181 The output of the [`ClapAudioModel`].182 """183 184 loss: Optional[torch.FloatTensor] = None185 logits_per_audio: Optional[torch.FloatTensor] = None186 logits_per_text: Optional[torch.FloatTensor] = None187 text_embeds: Optional[torch.FloatTensor] = None188 audio_embeds: Optional[torch.FloatTensor] = None189 text_model_output: BaseModelOutputWithPooling = None190 audio_model_output: BaseModelOutputWithPooling = None191 192 def to_tuple(self) -> tuple[Any]:193 return tuple(194 self[k] if k not in ["text_model_output", "audio_model_output"] else getattr(self, k).to_tuple()195 for k in self.keys()196 )197 198 199# Adapted from transformers.models.swin.modeling_swin.SwinDropPath200class ClapDropPath(nn.Module):201 """202 Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks). This is a slightly203 refactored version of the `SwinDropPath` implementation.204 """205 206 def __init__(self, drop_prob=None):207 super().__init__()208 self.drop_prob = drop_prob209 210 def forward(self, hidden_states):211 if self.drop_prob == 0.0 or not self.training:212 return hidden_states213 214 keep_prob = 1 - self.drop_prob215 # work with diff dim tensors, not just 2D ConvNets216 shape = (hidden_states.shape[0],) + (1,) * (hidden_states.ndim - 1)217 218 random_tensor = keep_prob + torch.rand(shape, dtype=hidden_states.dtype, device=hidden_states.device)219 random_tensor.floor_() # binarize220 output = hidden_states.div(keep_prob) * random_tensor221 return output222 223 224# Adapted from https://github.com/LAION-AI/CLAP/blob/6ad05a971ba0622f6acee8c41993e0d02bbed639/src/open_clip/feature_fusion.py#L133225class ClapAudioAFFBlock(nn.Module):226 r"""227 ATTENTIONAL FEATURE FUSION Block from CLAP, since in CLAP we are always in 2D mode, it is not needed to implement228 the 1D version.229 """230 231 def __init__(self, config: ClapAudioConfig):232 super().__init__()233 channels = config.patch_embeds_hidden_size234 downsize_ratio = config.aff_block_r235 inter_channels = int(channels // downsize_ratio)236 237 self.local_att = nn.Sequential(238 nn.Conv2d(channels, inter_channels, kernel_size=1, stride=1, padding=0),239 nn.BatchNorm2d(inter_channels),240 nn.ReLU(inplace=True),241 nn.Conv2d(inter_channels, channels, kernel_size=1, stride=1, padding=0),242 nn.BatchNorm2d(channels),243 )244 self.global_att = nn.Sequential(245 nn.AdaptiveAvgPool2d(1),246 nn.Conv2d(channels, inter_channels, kernel_size=1, stride=1, padding=0),247 nn.BatchNorm2d(inter_channels),248 nn.ReLU(inplace=True),249 nn.Conv2d(inter_channels, channels, kernel_size=1, stride=1, padding=0),250 nn.BatchNorm2d(channels),251 )252 253 self.sigmoid = nn.Sigmoid()254 255 def forward(self, hidden_states, residual):256 attention_input = hidden_states + residual257 258 fused_layer_output = self.local_att(attention_input) + self.global_att(attention_input)259 fused_layer_output = self.sigmoid(fused_layer_output)260 261 output = 2 * hidden_states * fused_layer_output + 2 * residual * (1 - fused_layer_output)262 return output263 264 265class ClapAudioPatchEmbed(nn.Module):266 """267 This module converts the hidden states reshaped as an image to patch embeddings ready to be passed to the268 Transformer block.269 """270 271 def __init__(self, config: ClapAudioConfig):272 super().__init__()273 img_size = (config.spec_size, config.spec_size) if isinstance(config.spec_size, int) else config.spec_size274 patch_size = (275 (config.patch_size, config.patch_size) if isinstance(config.patch_size, int) else config.patch_size276 )277 patch_stride = (278 (config.patch_stride, config.patch_stride) if isinstance(config.patch_stride, int) else config.patch_stride279 )280 281 self.img_size = img_size282 self.patch_stride = patch_stride283 284 self.grid_size = (img_size[0] // patch_stride[0], img_size[1] // patch_stride[1])285 self.num_patches = self.grid_size[0] * self.grid_size[1]286 287 self.flatten = config.flatten_patch_embeds288 self.enable_fusion = config.enable_fusion289 290 padding = ((patch_size[0] - patch_stride[0]) // 2, (patch_size[1] - patch_stride[1]) // 2)291 292 scale_factor = 4 if (self.enable_fusion) and (config.fusion_type == "channel_map") else 1293 294 self.proj = nn.Conv2d(295 config.patch_embed_input_channels * scale_factor,296 config.patch_embeds_hidden_size,297 kernel_size=patch_size,298 stride=patch_stride,299 padding=padding,300 )301 302 self.norm = nn.LayerNorm(config.patch_embeds_hidden_size) if config.enable_patch_layer_norm else nn.Identity()303 if self.enable_fusion:304 self.fusion_model = ClapAudioAFFBlock(config)305 self.mel_conv2d = nn.Conv2d(306 config.patch_embed_input_channels,307 config.patch_embeds_hidden_size,308 kernel_size=(patch_size[0], patch_size[1] * 3),309 stride=(patch_stride[0], patch_stride[1] * 3),310 padding=padding,311 )312 313 def forward(self, hidden_states, is_longer_idx=None):314 if self.enable_fusion:315 # retrieve the last mel as we have transposed the input316 global_hidden_states = hidden_states[:, 0:1, :, :]317 318 # global processing319 batch_size, num_channels, height, width = global_hidden_states.shape320 321 if height != self.img_size[0] or width != self.img_size[1]:322 raise ValueError(323 f"Input audio size ({height}*{width}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})."324 )325 326 global_hidden_states = self.proj(global_hidden_states)327 output_width = global_hidden_states.size(-1)328 if len(is_longer_idx) > 0:329 # local processing330 local_hidden_states = hidden_states[is_longer_idx, 1:, :, :].contiguous()331 batch_size, num_channels, height, width = local_hidden_states.shape332 local_hidden_states = local_hidden_states.view(batch_size * num_channels, 1, height, width)333 334 local_hidden_states = self.mel_conv2d(local_hidden_states)335 336 _, features, height, width = local_hidden_states.shape337 local_hidden_states = local_hidden_states.view(batch_size, num_channels, features, height, width)338 local_hidden_states = local_hidden_states.permute((0, 2, 3, 1, 4)).contiguous().flatten(3)339 340 local_width = local_hidden_states.size(-1)341 local_hidden_states = torch.nn.functional.pad(342 local_hidden_states, (0, output_width - local_width), "constant", 0343 )344 345 global_hidden_states[is_longer_idx] = self.fusion_model(346 global_hidden_states[is_longer_idx], local_hidden_states347 )348 hidden_states = global_hidden_states349 else:350 _, _, height, width = hidden_states.shape351 if height != self.img_size[0] or width != self.img_size[1]:352 raise ValueError(353 f"Input audio size ({height}*{width}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})."354 )355 hidden_states = self.proj(hidden_states)356 357 if self.flatten:358 hidden_states = hidden_states.flatten(2).transpose(1, 2)359 hidden_states = self.norm(hidden_states)360 return hidden_states361 362 363# Copied from transformers.models.swin.modeling_swin.SwinSelfAttention with Swin->ClapAudio364class ClapAudioSelfAttention(nn.Module):365 def __init__(self, config, dim, num_heads, window_size):366 super().__init__()367 if dim % num_heads != 0:368 raise ValueError(369 f"The hidden size ({dim}) is not a multiple of the number of attention heads ({num_heads})"370 )371 372 self.num_attention_heads = num_heads373 self.attention_head_size = int(dim / num_heads)374 self.all_head_size = self.num_attention_heads * self.attention_head_size375 self.window_size = (376 window_size if isinstance(window_size, collections.abc.Iterable) else (window_size, window_size)377 )378 379 self.relative_position_bias_table = nn.Parameter(380 torch.zeros((2 * self.window_size[0] - 1) * (2 * self.window_size[1] - 1), num_heads)381 )382 383 # get pair-wise relative position index for each token inside the window384 coords_h = torch.arange(self.window_size[0])385 coords_w = torch.arange(self.window_size[1])386 coords = torch.stack(meshgrid([coords_h, coords_w], indexing="ij"))387 coords_flatten = torch.flatten(coords, 1)388 relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :]389 relative_coords = relative_coords.permute(1, 2, 0).contiguous()390 relative_coords[:, :, 0] += self.window_size[0] - 1391 relative_coords[:, :, 1] += self.window_size[1] - 1392 relative_coords[:, :, 0] *= 2 * self.window_size[1] - 1393 relative_position_index = relative_coords.sum(-1)394 self.register_buffer("relative_position_index", relative_position_index)395 396 self.query = nn.Linear(self.all_head_size, self.all_head_size, bias=config.qkv_bias)397 self.key = nn.Linear(self.all_head_size, self.all_head_size, bias=config.qkv_bias)398 self.value = nn.Linear(self.all_head_size, self.all_head_size, bias=config.qkv_bias)399 400 self.dropout = nn.Dropout(config.attention_probs_dropout_prob)401 402 def forward(403 self,404 hidden_states: torch.Tensor,405 attention_mask: Optional[torch.FloatTensor] = None,406 head_mask: Optional[torch.FloatTensor] = None,407 output_attentions: Optional[bool] = False,408 ) -> tuple[torch.Tensor]:409 batch_size, dim, num_channels = hidden_states.shape410 hidden_shape = (batch_size, dim, -1, self.attention_head_size)411 412 query_layer = self.query(hidden_states).view(hidden_shape).transpose(1, 2)413 key_layer = self.key(hidden_states).view(hidden_shape).transpose(1, 2)414 value_layer = self.value(hidden_states).view(hidden_shape).transpose(1, 2)415 416 # Take the dot product between "query" and "key" to get the raw attention scores.417 attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2))418 419 attention_scores = attention_scores / math.sqrt(self.attention_head_size)420 421 relative_position_bias = self.relative_position_bias_table[self.relative_position_index.view(-1)]422 relative_position_bias = relative_position_bias.view(423 self.window_size[0] * self.window_size[1], self.window_size[0] * self.window_size[1], -1424 )425 426 relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous()427 attention_scores = attention_scores + relative_position_bias.unsqueeze(0)428 429 if attention_mask is not None:430 # Apply the attention mask is (precomputed for all layers in ClapAudioModel forward() function)431 mask_shape = attention_mask.shape[0]432 attention_scores = attention_scores.view(433 batch_size // mask_shape, mask_shape, self.num_attention_heads, dim, dim434 )435 attention_scores = attention_scores + attention_mask.unsqueeze(1).unsqueeze(0)436 attention_scores = attention_scores.view(-1, self.num_attention_heads, dim, dim)437 438 # Normalize the attention scores to probabilities.439 attention_probs = nn.functional.softmax(attention_scores, dim=-1)440 441 # This is actually dropping out entire tokens to attend to, which might442 # seem a bit unusual, but is taken from the original Transformer paper.443 attention_probs = self.dropout(attention_probs)444 445 # Mask heads if we want to446 if head_mask is not None:447 attention_probs = attention_probs * head_mask448 449 context_layer = torch.matmul(attention_probs, value_layer)450 context_layer = context_layer.permute(0, 2, 1, 3).contiguous()451 new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,)452 context_layer = context_layer.view(new_context_layer_shape)453 454 outputs = (context_layer, attention_probs) if output_attentions else (context_layer,)455 456 return outputs457 458 459# Copied from transformers.models.swin.modeling_swin.SwinSelfOutput with Swin->ClapAudio460class ClapAudioSelfOutput(nn.Module):461 def __init__(self, config, dim):462 super().__init__()463 self.dense = nn.Linear(dim, dim)464 self.dropout = nn.Dropout(config.attention_probs_dropout_prob)465 466 def forward(self, hidden_states: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor:467 hidden_states = self.dense(hidden_states)468 hidden_states = self.dropout(hidden_states)469 470 return hidden_states471 472 473# Copied from transformers.models.swin.modeling_swin.SwinAttention with Swin->ClapAudio474class ClapAudioAttention(nn.Module):475 def __init__(self, config, dim, num_heads, window_size):476 super().__init__()477 self.self = ClapAudioSelfAttention(config, dim, num_heads, window_size)478 self.output = ClapAudioSelfOutput(config, dim)479 self.pruned_heads = set()480 481 def prune_heads(self, heads):482 if len(heads) == 0:483 return484 heads, index = find_pruneable_heads_and_indices(485 heads, self.self.num_attention_heads, self.self.attention_head_size, self.pruned_heads486 )487 488 # Prune linear layers489 self.self.query = prune_linear_layer(self.self.query, index)490 self.self.key = prune_linear_layer(self.self.key, index)491 self.self.value = prune_linear_layer(self.self.value, index)492 self.output.dense = prune_linear_layer(self.output.dense, index, dim=1)493 494 # Update hyper params and store pruned heads495 self.self.num_attention_heads = self.self.num_attention_heads - len(heads)496 self.self.all_head_size = self.self.attention_head_size * self.self.num_attention_heads497 self.pruned_heads = self.pruned_heads.union(heads)498 499 def forward(500 self,501 hidden_states: torch.Tensor,502 attention_mask: Optional[torch.FloatTensor] = None,503 head_mask: Optional[torch.FloatTensor] = None,504 output_attentions: Optional[bool] = False,505 ) -> tuple[torch.Tensor]:506 self_outputs = self.self(hidden_states, attention_mask, head_mask, output_attentions)507 attention_output = self.output(self_outputs[0], hidden_states)508 outputs = (attention_output,) + self_outputs[1:] # add attentions if we output them509 return outputs510 511 512# Copied from transformers.models.swin.modeling_swin.SwinIntermediate with Swin->ClapAudio513class ClapAudioIntermediate(nn.Module):514 def __init__(self, config, dim):515 super().__init__()516 self.dense = nn.Linear(dim, int(config.mlp_ratio * dim))517 if isinstance(config.hidden_act, str):518 self.intermediate_act_fn = ACT2FN[config.hidden_act]519 else:520 self.intermediate_act_fn = config.hidden_act521 522 def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:523 hidden_states = self.dense(hidden_states)524 hidden_states = self.intermediate_act_fn(hidden_states)525 return hidden_states526 527 528# Copied from transformers.models.swin.modeling_swin.SwinOutput with Swin->ClapAudio529class ClapAudioOutput(nn.Module):530 def __init__(self, config, dim):531 super().__init__()532 self.dense = nn.Linear(int(config.mlp_ratio * dim), dim)533 self.dropout = nn.Dropout(config.hidden_dropout_prob)534 535 def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:536 hidden_states = self.dense(hidden_states)537 hidden_states = self.dropout(hidden_states)538 return hidden_states539 540 541# Copied from transformers.models.swin.modeling_swin.SwinLayer with SwinDropPath->ClapDropPath, Swin->ClapAudio542class ClapAudioLayer(nn.Module):543 def __init__(self, config, dim, input_resolution, num_heads, drop_path_rate=0.0, shift_size=0):544 super().__init__()545 self.chunk_size_feed_forward = config.chunk_size_feed_forward546 self.shift_size = shift_size547 self.window_size = config.window_size548 self.input_resolution = input_resolution549 self.layernorm_before = nn.LayerNorm(dim, eps=config.layer_norm_eps)550 self.attention = ClapAudioAttention(config, dim, num_heads, window_size=self.window_size)551 self.drop_path = ClapDropPath(drop_path_rate) if drop_path_rate > 0.0 else nn.Identity()552 self.layernorm_after = nn.LayerNorm(dim, eps=config.layer_norm_eps)553 self.intermediate = ClapAudioIntermediate(config, dim)554 self.output = ClapAudioOutput(config, dim)555 556 def set_shift_and_window_size(self, input_resolution):557 if min(input_resolution) <= self.window_size:558 # if window size is larger than input resolution, we don't partition windows559 self.shift_size = torch_int(0)560 self.window_size = (561 torch.min(torch.tensor(input_resolution)) if torch.jit.is_tracing() else min(input_resolution)562 )563 564 def get_attn_mask(self, height, width, dtype, device):565 if self.shift_size > 0:566 # calculate attention mask for SW-MSA567 img_mask = torch.zeros((1, height, width, 1), dtype=dtype, device=device)568 height_slices = (569 slice(0, -self.window_size),570 slice(-self.window_size, -self.shift_size),571 slice(-self.shift_size, None),572 )573 width_slices = (574 slice(0, -self.window_size),575 slice(-self.window_size, -self.shift_size),576 slice(-self.shift_size, None),577 )578 count = 0579 for height_slice in height_slices:580 for width_slice in width_slices:581 img_mask[:, height_slice, width_slice, :] = count582 count += 1583 584 mask_windows = window_partition(img_mask, self.window_size)585 mask_windows = mask_windows.view(-1, self.window_size * self.window_size)586 attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2)587 attn_mask = attn_mask.masked_fill(attn_mask != 0, -100.0).masked_fill(attn_mask == 0, 0.0)588 else:589 attn_mask = None590 return attn_mask591 592 def maybe_pad(self, hidden_states, height, width):593 pad_right = (self.window_size - width % self.window_size) % self.window_size594 pad_bottom = (self.window_size - height % self.window_size) % self.window_size595 pad_values = (0, 0, 0, pad_right, 0, pad_bottom)596 hidden_states = nn.functional.pad(hidden_states, pad_values)597 return hidden_states, pad_values598 599 def forward(600 self,601 hidden_states: torch.Tensor,602 input_dimensions: tuple[int, int],603 head_mask: Optional[torch.FloatTensor] = None,604 output_attentions: Optional[bool] = False,605 always_partition: Optional[bool] = False,606 ) -> tuple[torch.Tensor, torch.Tensor]:607 if not always_partition:608 self.set_shift_and_window_size(input_dimensions)609 else:610 pass611 height, width = input_dimensions612 batch_size, _, channels = hidden_states.size()613 shortcut = hidden_states614 615 hidden_states = self.layernorm_before(hidden_states)616 617 hidden_states = hidden_states.view(batch_size, height, width, channels)618 619 # pad hidden_states to multiples of window size620 hidden_states, pad_values = self.maybe_pad(hidden_states, height, width)621 622 _, height_pad, width_pad, _ = hidden_states.shape623 # cyclic shift624 if self.shift_size > 0:625 shifted_hidden_states = torch.roll(hidden_states, shifts=(-self.shift_size, -self.shift_size), dims=(1, 2))626 else:627 shifted_hidden_states = hidden_states628 629 # partition windows630 hidden_states_windows = window_partition(shifted_hidden_states, self.window_size)631 hidden_states_windows = hidden_states_windows.view(-1, self.window_size * self.window_size, channels)632 attn_mask = self.get_attn_mask(633 height_pad, width_pad, dtype=hidden_states.dtype, device=hidden_states_windows.device634 )635 636 attention_outputs = self.attention(637 hidden_states_windows, attn_mask, head_mask, output_attentions=output_attentions638 )639 640 attention_output = attention_outputs[0]641 642 attention_windows = attention_output.view(-1, self.window_size, self.window_size, channels)643 shifted_windows = window_reverse(attention_windows, self.window_size, height_pad, width_pad)644 645 # reverse cyclic shift646 if self.shift_size > 0:647 attention_windows = torch.roll(shifted_windows, shifts=(self.shift_size, self.shift_size), dims=(1, 2))648 else:649 attention_windows = shifted_windows650 651 was_padded = pad_values[3] > 0 or pad_values[5] > 0652 if was_padded:653 attention_windows = attention_windows[:, :height, :width, :].contiguous()654 655 attention_windows = attention_windows.view(batch_size, height * width, channels)656 657 hidden_states = shortcut + self.drop_path(attention_windows)658 659 layer_output = self.layernorm_after(hidden_states)660 layer_output = self.intermediate(layer_output)661 layer_output = hidden_states + self.output(layer_output)662 663 layer_outputs = (layer_output, attention_outputs[1]) if output_attentions else (layer_output,)664 return layer_outputs665 666 667# Copied from transformers.models.swin.modeling_swin.SwinStage with Swin->ClapAudio668class ClapAudioStage(GradientCheckpointingLayer):669 def __init__(self, config, dim, input_resolution, depth, num_heads, drop_path, downsample):670 super().__init__()671 self.config = config672 self.dim = dim673 self.blocks = nn.ModuleList(674 [675 ClapAudioLayer(676 config=config,677 dim=dim,678 input_resolution=input_resolution,679 num_heads=num_heads,680 drop_path_rate=drop_path[i],681 shift_size=0 if (i % 2 == 0) else config.window_size // 2,682 )683 for i in range(depth)684 ]685 )686 687 # patch merging layer688 if downsample is not None:689 self.downsample = downsample(input_resolution, dim=dim, norm_layer=nn.LayerNorm)690 else:691 self.downsample = None692 693 self.pointing = False694 695 def forward(696 self,697 hidden_states: torch.Tensor,698 input_dimensions: tuple[int, int],699 head_mask: Optional[torch.FloatTensor] = None,700 output_attentions: Optional[bool] = False,701 always_partition: Optional[bool] = False,702 ) -> tuple[torch.Tensor]:703 height, width = input_dimensions704 for i, layer_module in enumerate(self.blocks):705 layer_head_mask = head_mask[i] if head_mask is not None else None706 707 layer_outputs = layer_module(708 hidden_states, input_dimensions, layer_head_mask, output_attentions, always_partition709 )710 711 hidden_states = layer_outputs[0]712 713 hidden_states_before_downsampling = hidden_states714 if self.downsample is not None:715 height_downsampled, width_downsampled = (height + 1) // 2, (width + 1) // 2716 output_dimensions = (height, width, height_downsampled, width_downsampled)717 hidden_states = self.downsample(hidden_states_before_downsampling, input_dimensions)718 else:719 output_dimensions = (height, width, height, width)720 721 stage_outputs = (hidden_states, hidden_states_before_downsampling, output_dimensions)722 723 if output_attentions:724 stage_outputs += layer_outputs[1:]725 return stage_outputs726 727 728# Copied from transformers.models.swin.modeling_swin.SwinPatchMerging with Swin->ClapAudio729class ClapAudioPatchMerging(nn.Module):730 """731 Patch Merging Layer.732 733 Args:734 input_resolution (`tuple[int]`):735 Resolution of input feature.736 dim (`int`):737 Number of input channels.738 norm_layer (`nn.Module`, *optional*, defaults to `nn.LayerNorm`):739 Normalization layer class.740 """741 742 def __init__(self, input_resolution: tuple[int], dim: int, norm_layer: nn.Module = nn.LayerNorm) -> None:743 super().__init__()744 self.input_resolution = input_resolution745 self.dim = dim746 self.reduction = nn.Linear(4 * dim, 2 * dim, bias=False)747 self.norm = norm_layer(4 * dim)748 749 def maybe_pad(self, input_feature, height, width):750 should_pad = (height % 2 == 1) or (width % 2 == 1)751 if should_pad:752 pad_values = (0, 0, 0, width % 2, 0, height % 2)753 input_feature = nn.functional.pad(input_feature, pad_values)754 755 return input_feature756 757 def forward(self, input_feature: torch.Tensor, input_dimensions: tuple[int, int]) -> torch.Tensor:758 height, width = input_dimensions759 # `dim` is height * width760 batch_size, dim, num_channels = input_feature.shape761 762 input_feature = input_feature.view(batch_size, height, width, num_channels)763 # pad input to be divisible by width and height, if needed764 input_feature = self.maybe_pad(input_feature, height, width)765 # [batch_size, height/2, width/2, num_channels]766 input_feature_0 = input_feature[:, 0::2, 0::2, :]767 # [batch_size, height/2, width/2, num_channels]768 input_feature_1 = input_feature[:, 1::2, 0::2, :]769 # [batch_size, height/2, width/2, num_channels]770 input_feature_2 = input_feature[:, 0::2, 1::2, :]771 # [batch_size, height/2, width/2, num_channels]772 input_feature_3 = input_feature[:, 1::2, 1::2, :]773 # batch_size height/2 width/2 4*num_channels774 input_feature = torch.cat([input_feature_0, input_feature_1, input_feature_2, input_feature_3], -1)775 input_feature = input_feature.view(batch_size, -1, 4 * num_channels) # batch_size height/2*width/2 4*C776 777 input_feature = self.norm(input_feature)778 input_feature = self.reduction(input_feature)779 780 return input_feature781 782 783class ClapAudioEncoder(nn.Module):784 def __init__(self, config):785 super().__init__()786 self.num_layers = len(config.depths)787 788 self.config = config789 self.patch_embed = ClapAudioPatchEmbed(config)790 self.enable_fusion = config.enable_fusion791 self.patch_stride = self.patch_embed.patch_stride792 self.spec_size = config.spec_size793 self.freq_ratio = config.spec_size // config.num_mel_bins794 795 self.num_features = int(config.patch_embeds_hidden_size * 2 ** (self.num_layers - 1))796 797 drop_path_rate = [x.item() for x in torch.linspace(0, config.drop_path_rate, sum(config.depths), device="cpu")]798 799 grid_size = self.patch_embed.grid_size800 self.input_resolutions = [(grid_size[0] // (2**i), grid_size[1] // (2**i)) for i in range(self.num_layers)]801 802 self.layers = nn.ModuleList(803 [804 ClapAudioStage(805 config=config,806 dim=int(config.patch_embeds_hidden_size * 2**i_layer),807 input_resolution=self.input_resolutions[i_layer],808 depth=config.depths[i_layer],809 num_heads=config.num_attention_heads[i_layer],810 drop_path=drop_path_rate[sum(config.depths[:i_layer]) : sum(config.depths[: i_layer + 1])],811 downsample=ClapAudioPatchMerging if (i_layer < self.num_layers - 1) else None,812 )813 for i_layer in range(self.num_layers)814 ]815 )816 817 self.gradient_checkpointing = False818 819 self.batch_norm = nn.BatchNorm2d(config.num_mel_bins)820 self.norm = nn.LayerNorm(self.num_features)821 self.depths = config.depths822 self.avgpool = nn.AdaptiveAvgPool1d(1)823 824 def reshape_mel2img(self, normalized_input_features):825 """826 The input is 4 normalized log mel spectrograms. It is reshape to the common shape of images. Each channel827 should represent 1 of the 4 crops of the spectrogram. For more details, refer to the [`ClapFeatureExtractor`].828 """829 _, _, time_length, freq_length = normalized_input_features.shape830 831 spec_width = int(self.spec_size * self.freq_ratio)832 spec_height = self.spec_size // self.freq_ratio833 834 if time_length > spec_width or freq_length > spec_height:835 raise ValueError("the wav size should be less than or equal to the swin input size")836 837 # to avoid bicubic zero error838 if time_length < spec_width:839 normalized_input_features = nn.functional.interpolate(840 normalized_input_features, (spec_width, freq_length), mode="bicubic", align_corners=True841 )842 if freq_length < spec_height:843 normalized_input_features = nn.functional.interpolate(844 normalized_input_features, (time_length, spec_height), mode="bicubic", align_corners=True845 )846 847 batch, channels, time, freq = normalized_input_features.shape848 849 # batch_size, channels, spec_width, spec_height --> batch_size, channels, spec_height * freq_ratio, spec_width // freq_ratio850 normalized_input_features = normalized_input_features.reshape(851 batch, channels * self.freq_ratio, time // self.freq_ratio, freq852 )853 normalized_input_features = normalized_input_features.permute(0, 1, 3, 2).contiguous()854 normalized_input_features = normalized_input_features.reshape(855 batch, channels, freq * self.freq_ratio, time // self.freq_ratio856 )857 858 return normalized_input_features859 860 def forward(861 self,862 input_features,863 is_longer: Optional[torch.FloatTensor] = None,864 head_mask: Optional[torch.FloatTensor] = None,865 output_attentions: Optional[bool] = False,866 output_hidden_states: Optional[bool] = False,867 output_hidden_states_before_downsampling: Optional[bool] = False,868 always_partition: Optional[bool] = False,869 return_dict: Optional[bool] = True,870 ) -> Union[tuple, ClapAudioModelOutput]:871 input_features = input_features.transpose(1, 3)872 normalized_input_features = self.batch_norm(input_features)873 normalized_input_features = normalized_input_features.transpose(1, 3)874 875 is_longer_list_idx = None876 if self.enable_fusion:877 is_longer_list = is_longer.to(input_features.device)878 is_longer_list_idx = torch.where(is_longer_list == 1)[0]879 880 hidden_states = self.reshape_mel2img(normalized_input_features)881 882 frames_num = hidden_states.shape[2]883 884 hidden_states = self.patch_embed(hidden_states, is_longer_list_idx)885 886 all_hidden_states = () if output_hidden_states else None887 all_reshaped_hidden_states = () if output_hidden_states else None888 all_self_attentions = () if output_attentions else None889 890 input_dimensions = self.input_resolutions[0]891 892 if output_hidden_states:893 batch_size, _, hidden_size = hidden_states.shape894 # rearrange batch_size (height width) channels -> batch_size channel height width895 reshaped_hidden_state = hidden_states.view(batch_size, *input_dimensions, hidden_size)896 reshaped_hidden_state = reshaped_hidden_state.permute(0, 3, 1, 2)897 all_hidden_states += (hidden_states,)898 all_reshaped_hidden_states += (reshaped_hidden_state,)899 900 for i, layer_module in enumerate(self.layers):901 layer_head_mask = head_mask[i] if head_mask is not None else None902 903 input_dimensions = self.input_resolutions[i]904 905 layer_outputs = layer_module(906 hidden_states, input_dimensions, layer_head_mask, output_attentions, always_partition907 )908 909 hidden_states = layer_outputs[0]910 911 hidden_states_before_downsampling = layer_outputs[1]912 output_dimensions = layer_outputs[2]913 914 input_dimensions = (output_dimensions[-2], output_dimensions[-1])915 916 if output_hidden_states and output_hidden_states_before_downsampling:917 batch_size, _, hidden_size = hidden_states_before_downsampling.shape918 # rearrange batch_size (height width) channels -> batch_size channel height width919 # here we use the original (not downsampled) height and width920 reshaped_hidden_state = hidden_states_before_downsampling.view(921 batch_size, *(output_dimensions[0], output_dimensions[1]), hidden_size922 )923 reshaped_hidden_state = reshaped_hidden_state.permute(0, 3, 1, 2)924 all_hidden_states += (hidden_states_before_downsampling,)925 all_reshaped_hidden_states += (reshaped_hidden_state,)926 elif output_hidden_states and not output_hidden_states_before_downsampling:927 batch_size, _, hidden_size = hidden_states.shape928 # rearrange batch_size (height width) channels -> batch_size channel height width929 reshaped_hidden_state = hidden_states.view(batch_size, *input_dimensions, hidden_size)930 reshaped_hidden_state = reshaped_hidden_state.permute(0, 3, 1, 2)931 all_hidden_states += (hidden_states,)932 all_reshaped_hidden_states += (reshaped_hidden_state,)933 934 if output_attentions:935 all_self_attentions += layer_outputs[3:]936 937 last_hidden_state = self.norm(hidden_states)938 939 batch_size, _, n_channels = last_hidden_state.shape940 941 freq_shape = frames_num // (2 ** (len(self.depths) - 1)) // self.patch_stride[0]942 temporal_shape = frames_num // (2 ** (len(self.depths) - 1)) // self.patch_stride[1]943 944 last_hidden_state = (945 last_hidden_state.permute(0, 2, 1).contiguous().reshape(batch_size, n_channels, freq_shape, temporal_shape)946 )947 948 batch_size, n_channels, n_frequencies, n_temp = last_hidden_state.shape949 # group 2D CNN950 c_freq_bin = n_frequencies // self.freq_ratio951 last_hidden_state = last_hidden_state.reshape(952 batch_size, n_channels, n_frequencies // c_freq_bin, c_freq_bin, n_temp953 )954 last_hidden_state = (955 last_hidden_state.permute(0, 1, 3, 2, 4).contiguous().reshape(batch_size, n_channels, c_freq_bin, -1)956 )957 latent_output = self.avgpool(torch.flatten(last_hidden_state, 2))958 latent_output = torch.flatten(latent_output, 1)959 960 if not return_dict:961 return tuple(962 v963 for v in [964 last_hidden_state,965 latent_output,966 all_reshaped_hidden_states,967 all_self_attentions,968 ]969 if v is not None970 )971 972 return BaseModelOutputWithPooling(973 last_hidden_state=last_hidden_state,974 pooler_output=latent_output,975 hidden_states=all_reshaped_hidden_states,976 attentions=all_self_attentions,977 )978 979 980class ClapProjectionLayer(nn.Module):981 def __init__(self, config: Union[ClapAudioConfig, ClapTextConfig]):982 super().__init__()983 self.config = config984 hidden_size = config.hidden_size985 projection_dim = config.projection_dim986 987 self.linear1 = nn.Linear(hidden_size, projection_dim)988 self.activation = ACT2FN[config.projection_hidden_act]989 self.linear2 = nn.Linear(projection_dim, projection_dim)990 991 def forward(self, hidden_states):992 hidden_states = self.linear1(hidden_states)993 hidden_states = self.activation(hidden_states)994 hidden_states = self.linear2(hidden_states)995 return hidden_states996 997 998# Copied from transformers.models.roberta.modeling_roberta.RobertaEmbeddings with Roberta->ClapText, persistent=False->persistent=True999class ClapTextEmbeddings(nn.Module):1000 """1001 Same as BertEmbeddings with a tiny tweak for positional embeddings indexing.1002 """1003 1004 # Copied from transformers.models.bert.modeling_bert.BertEmbeddings.__init__1005 def __init__(self, config):1006 super().__init__()1007 self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id)1008 self.position_embeddings = nn.Embedding(config.max_position_embeddings, config.hidden_size)1009 self.token_type_embeddings = nn.Embedding(config.type_vocab_size, config.hidden_size)1010 1011 # self.LayerNorm is not snake-cased to stick with TensorFlow model variable name and be able to load1012 # any TensorFlow checkpoint file1013 self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)1014 self.dropout = nn.Dropout(config.hidden_dropout_prob)1015 # position_ids (1, len position emb) is contiguous in memory and exported when serialized1016 self.position_embedding_type = getattr(config, "position_embedding_type", "absolute")1017 self.register_buffer(1018 "position_ids", torch.arange(config.max_position_embeddings).expand((1, -1)), persistent=True1019 )1020 self.register_buffer(1021 "token_type_ids", torch.zeros(self.position_ids.size(), dtype=torch.long), persistent=True1022 )1023 1024 # End copy1025 self.padding_idx = config.pad_token_id1026 self.position_embeddings = nn.Embedding(1027 config.max_position_embeddings, config.hidden_size, padding_idx=self.padding_idx1028 )1029 1030 def forward(1031 self, input_ids=None, token_type_ids=None, position_ids=None, inputs_embeds=None, past_key_values_length=01032 ):1033 if position_ids is None:1034 if input_ids is not None:1035 # Create the position ids from the input token ids. Any padded tokens remain padded.1036 position_ids = create_position_ids_from_input_ids(input_ids, self.padding_idx, past_key_values_length)1037 else:1038 position_ids = self.create_position_ids_from_inputs_embeds(inputs_embeds)1039 1040 if input_ids is not None:1041 input_shape = input_ids.size()1042 else:1043 input_shape = inputs_embeds.size()[:-1]1044 1045 seq_length = input_shape[1]1046 1047 # Setting the token_type_ids to the registered buffer in constructor where it is all zeros, which usually occurs1048 # when its auto-generated, registered buffer helps users when tracing the model without passing token_type_ids, solves1049 # issue #56641050 if token_type_ids is None:1051 if hasattr(self, "token_type_ids"):1052 buffered_token_type_ids = self.token_type_ids[:, :seq_length]1053 buffered_token_type_ids_expanded = buffered_token_type_ids.expand(input_shape[0], seq_length)1054 token_type_ids = buffered_token_type_ids_expanded1055 else:1056 token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=self.position_ids.device)1057 1058 if inputs_embeds is None:1059 inputs_embeds = self.word_embeddings(input_ids)1060 token_type_embeddings = self.token_type_embeddings(token_type_ids)1061 1062 embeddings = inputs_embeds + token_type_embeddings1063 if self.position_embedding_type == "absolute":1064 position_embeddings = self.position_embeddings(position_ids)1065 embeddings += position_embeddings1066 embeddings = self.LayerNorm(embeddings)1067 embeddings = self.dropout(embeddings)1068 return embeddings1069 1070 def create_position_ids_from_inputs_embeds(self, inputs_embeds):1071 """1072 We are provided embeddings directly. We cannot infer which are padded so just generate sequential position ids.1073 1074 Args:1075 inputs_embeds: torch.Tensor1076 1077 Returns: torch.Tensor1078 """1079 input_shape = inputs_embeds.size()[:-1]1080 sequence_length = input_shape[1]1081 1082 position_ids = torch.arange(1083 self.padding_idx + 1, sequence_length + self.padding_idx + 1, dtype=torch.long, device=inputs_embeds.device1084 )1085 return position_ids.unsqueeze(0).expand(input_shape)1086 1087 1088# Copied from transformers.models.align.modeling_align.eager_attention_forward1089def eager_attention_forward(1090 module: nn.Module,1091 query: torch.Tensor,1092 key: torch.Tensor,1093 value: torch.Tensor,1094 attention_mask: Optional[torch.Tensor],1095 scaling: float,1096 dropout: float = 0.0,1097 head_mask: Optional[torch.Tensor] = None,1098 **kwargs,1099):1100 attn_weights = torch.matmul(query, key.transpose(2, 3)) * scaling1101 if attention_mask is not None:1102 causal_mask = attention_mask[:, :, :, : key.shape[-2]]1103 attn_weights = attn_weights + causal_mask1104 1105 attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)1106 attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)1107 1108 if head_mask is not None:1109 attn_weights = attn_weights * head_mask.view(1, -1, 1, 1)1110 1111 attn_output = torch.matmul(attn_weights, value)1112 attn_output = attn_output.transpose(1, 2).contiguous()1113 return attn_output, attn_weights1114 1115 1116# Copied from transformers.models.align.modeling_align.AlignTextSelfAttention with Align->Clap1117class ClapTextSelfAttention(nn.Module):1118 def __init__(self, config):1119 super().__init__()1120 if config.hidden_size % config.num_attention_heads != 0 and not hasattr(config, "embedding_size"):1121 raise ValueError(1122 f"The hidden size ({config.hidden_size}) is not a multiple of the number of attention "1123 f"heads ({config.num_attention_heads})"1124 )1125 1126 self.config = config1127 self.num_attention_heads = config.num_attention_heads1128 self.attention_head_size = int(config.hidden_size / config.num_attention_heads)1129 self.all_head_size = self.num_attention_heads * self.attention_head_size1130 1131 self.query = nn.Linear(config.hidden_size, self.all_head_size)1132 self.key = nn.Linear(config.hidden_size, self.all_head_size)1133 self.value = nn.Linear(config.hidden_size, self.all_head_size)1134 1135 self.dropout = nn.Dropout(config.attention_probs_dropout_prob)1136 self.attention_dropout = config.attention_probs_dropout_prob1137 self.scaling = self.attention_head_size**-0.51138 1139 def forward(1140 self,1141 hidden_states: torch.Tensor,1142 attention_mask: Optional[torch.FloatTensor] = None,1143 head_mask: Optional[torch.FloatTensor] = None,1144 output_attentions: Optional[bool] = False,1145 **kwargs,1146 ) -> tuple[torch.Tensor]:1147 input_shape = hidden_states.shape[:-1]1148 hidden_shape = (*input_shape, -1, self.attention_head_size)1149 1150 query_states = self.query(hidden_states).view(hidden_shape).transpose(1, 2)1151 key_states = self.key(hidden_states).view(hidden_shape).transpose(1, 2)1152 value_states = self.value(hidden_states).view(hidden_shape).transpose(1, 2)1153 1154 attention_interface: Callable = eager_attention_forward1155 if self.config._attn_implementation != "eager":1156 attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]1157 1158 attn_output, attn_weights = attention_interface(1159 self,1160 query_states,1161 key_states,1162 value_states,1163 attention_mask,1164 dropout=0.0 if not self.training else self.attention_dropout,1165 scaling=self.scaling,1166 head_mask=head_mask,1167 **kwargs,1168 )1169 1170 attn_output = attn_output.reshape(*input_shape, -1).contiguous()1171 outputs = (attn_output, attn_weights) if output_attentions else (attn_output,)1172 return outputs1173 1174 1175# Copied from transformers.models.bert.modeling_bert.BertSelfOutput1176class ClapTextSelfOutput(nn.Module):1177 def __init__(self, config):1178 super().__init__()1179 self.dense = nn.Linear(config.hidden_size, config.hidden_size)1180 self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)1181 self.dropout = nn.Dropout(config.hidden_dropout_prob)1182 1183 def forward(self, hidden_states: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor:1184 hidden_states = self.dense(hidden_states)1185 hidden_states = self.dropout(hidden_states)1186 hidden_states = self.LayerNorm(hidden_states + input_tensor)1187 return hidden_states1188 1189 1190# Copied from transformers.models.align.modeling_align.AlignTextAttention with Align->Clap1191class ClapTextAttention(nn.Module):1192 def __init__(self, config):1193 super().__init__()1194 self.self = ClapTextSelfAttention(config)1195 self.output = ClapTextSelfOutput(config)1196 self.pruned_heads = set()1197 1198 def prune_heads(self, heads):1199 if len(heads) == 0:1200 return