Aluode/PerceptionLabPortable
0
1# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨2# This file was automatically generated from src/transformers/models/metaclip_2/modular_metaclip_2.py.3# Do NOT edit this file manually as any edits will be overwritten by the generation of4# the file from the modular. If any change should be done, please apply the change to the5# modular_metaclip_2.py file directly. One of our CI enforces this.6# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨7from dataclasses import dataclass8from typing import Any, Callable, Optional, Union9 10import torch11from torch import nn12 13from ...activations import ACT2FN14from ...modeling_attn_mask_utils import _create_4d_causal_attention_mask, _prepare_4d_attention_mask15from ...modeling_layers import GradientCheckpointingLayer16from ...modeling_outputs import BaseModelOutput, BaseModelOutputWithPooling, ImageClassifierOutput17from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel18from ...processing_utils import Unpack19from ...utils import (20 ModelOutput,21 TransformersKwargs,22 auto_docstring,23 can_return_tuple,24 filter_out_non_signature_kwargs,25 torch_int,26)27from ...utils.generic import check_model_inputs28from .configuration_metaclip_2 import MetaClip2Config, MetaClip2TextConfig, MetaClip2VisionConfig29 30 31class MetaClip2TextEmbeddings(nn.Module):32 def __init__(self, config: MetaClip2TextConfig):33 super().__init__()34 embed_dim = config.hidden_size35 36 self.token_embedding = nn.Embedding(config.vocab_size, embed_dim)37 self.position_embedding = nn.Embedding(config.max_position_embeddings, embed_dim)38 39 # position_ids (1, len position emb) is contiguous in memory and exported when serialized40 self.register_buffer(41 "position_ids", torch.arange(config.max_position_embeddings).expand((1, -1)), persistent=False42 )43 44 def forward(45 self,46 input_ids: Optional[torch.LongTensor] = None,47 position_ids: Optional[torch.LongTensor] = None,48 inputs_embeds: Optional[torch.FloatTensor] = None,49 ) -> torch.Tensor:50 seq_length = input_ids.shape[-1] if input_ids is not None else inputs_embeds.shape[-2]51 max_position_embedding = self.position_embedding.weight.shape[0]52 53 if seq_length > max_position_embedding:54 raise ValueError(55 f"Sequence length must be less than max_position_embeddings (got `sequence length`: "56 f"{seq_length} and max_position_embeddings: {max_position_embedding}"57 )58 59 if position_ids is None:60 position_ids = self.position_ids[:, :seq_length]61 62 if inputs_embeds is None:63 inputs_embeds = self.token_embedding(input_ids)64 65 position_embeddings = self.position_embedding(position_ids)66 embeddings = inputs_embeds + position_embeddings67 68 return embeddings69 70 71class MetaClip2VisionEmbeddings(nn.Module):72 def __init__(self, config: MetaClip2VisionConfig):73 super().__init__()74 self.config = config75 self.embed_dim = config.hidden_size76 self.image_size = config.image_size77 self.patch_size = config.patch_size78 79 self.class_embedding = nn.Parameter(torch.randn(self.embed_dim))80 81 self.patch_embedding = nn.Conv2d(82 in_channels=config.num_channels,83 out_channels=self.embed_dim,84 kernel_size=self.patch_size,85 stride=self.patch_size,86 bias=False,87 )88 89 self.num_patches = (self.image_size // self.patch_size) ** 290 self.num_positions = self.num_patches + 191 self.position_embedding = nn.Embedding(self.num_positions, self.embed_dim)92 self.register_buffer("position_ids", torch.arange(self.num_positions).expand((1, -1)), persistent=False)93 94 def interpolate_pos_encoding(self, embeddings: torch.Tensor, height: int, width: int) -> torch.Tensor:95 """96 This method allows to interpolate the pre-trained position encodings, to be able to use the model on higher resolution97 images. This method is also adapted to support torch.jit tracing.98 99 Adapted from:100 - https://github.com/facebookresearch/dino/blob/de9ee3df6cf39fac952ab558447af1fa1365362a/vision_transformer.py#L174-L194, and101 - https://github.com/facebookresearch/dinov2/blob/e1277af2ba9496fbadf7aec6eba56e8d882d1e35/dinov2/models/vision_transformer.py#L179-L211102 """103 104 num_patches = embeddings.shape[1] - 1105 position_embedding = self.position_embedding.weight.unsqueeze(0)106 num_positions = position_embedding.shape[1] - 1107 108 # always interpolate when tracing to ensure the exported model works for dynamic input shapes109 if not torch.jit.is_tracing() and num_patches == num_positions and height == width:110 return self.position_embedding(self.position_ids)111 112 class_pos_embed = position_embedding[:, :1]113 patch_pos_embed = position_embedding[:, 1:]114 115 dim = embeddings.shape[-1]116 117 new_height = height // self.patch_size118 new_width = width // self.patch_size119 120 sqrt_num_positions = torch_int(num_positions**0.5)121 patch_pos_embed = patch_pos_embed.reshape(1, sqrt_num_positions, sqrt_num_positions, dim)122 patch_pos_embed = patch_pos_embed.permute(0, 3, 1, 2)123 124 patch_pos_embed = nn.functional.interpolate(125 patch_pos_embed,126 size=(new_height, new_width),127 mode="bicubic",128 align_corners=False,129 )130 131 patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim)132 133 return torch.cat((class_pos_embed, patch_pos_embed), dim=1)134 135 def forward(self, pixel_values: torch.FloatTensor, interpolate_pos_encoding=False) -> torch.Tensor:136 batch_size, _, height, width = pixel_values.shape137 if not interpolate_pos_encoding and (height != self.image_size or width != self.image_size):138 raise ValueError(139 f"Input image size ({height}*{width}) doesn't match model ({self.image_size}*{self.image_size})."140 )141 target_dtype = self.patch_embedding.weight.dtype142 patch_embeds = self.patch_embedding(pixel_values.to(dtype=target_dtype)) # shape = [*, width, grid, grid]143 patch_embeds = patch_embeds.flatten(2).transpose(1, 2)144 145 class_embeds = self.class_embedding.expand(batch_size, 1, -1)146 embeddings = torch.cat([class_embeds, patch_embeds], dim=1)147 if interpolate_pos_encoding:148 embeddings = embeddings + self.interpolate_pos_encoding(embeddings, height, width)149 else:150 embeddings = embeddings + self.position_embedding(self.position_ids)151 return embeddings152 153 154def eager_attention_forward(155 module: nn.Module,156 query: torch.Tensor,157 key: torch.Tensor,158 value: torch.Tensor,159 attention_mask: Optional[torch.Tensor],160 scaling: float,161 dropout: float = 0.0,162 output_attentions: bool = True,163 **kwargs,164):165 attn_weights = torch.matmul(query, key.transpose(-1, -2)) * scaling166 if attention_mask is not None:167 attn_weights = attn_weights + attention_mask168 attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)169 attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)170 171 attn_output = torch.matmul(attn_weights, value)172 attn_output = attn_output.transpose(1, 2).contiguous()173 if not output_attentions:174 attn_weights = None175 return attn_output, attn_weights176 177 178class MetaClip2Attention(nn.Module):179 """Multi-headed attention from 'Attention Is All You Need' paper"""180 181 def __init__(self, config: Union[MetaClip2VisionConfig, MetaClip2TextConfig]):182 super().__init__()183 self.config = config184 self.embed_dim = config.hidden_size185 self.num_heads = config.num_attention_heads186 self.head_dim = self.embed_dim // self.num_heads187 if self.head_dim * self.num_heads != self.embed_dim:188 raise ValueError(189 f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim} and `num_heads`:"190 f" {self.num_heads})."191 )192 self.scale = self.head_dim**-0.5193 self.dropout = config.attention_dropout194 self.is_causal = False195 196 self.k_proj = nn.Linear(self.embed_dim, self.embed_dim)197 self.v_proj = nn.Linear(self.embed_dim, self.embed_dim)198 self.q_proj = nn.Linear(self.embed_dim, self.embed_dim)199 self.out_proj = nn.Linear(self.embed_dim, self.embed_dim)200 201 def forward(202 self,203 hidden_states: torch.Tensor,204 attention_mask: Optional[torch.Tensor] = None,205 causal_attention_mask: Optional[torch.Tensor] = None,206 output_attentions: Optional[bool] = False,207 ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:208 """Input shape: Batch x Time x Channel"""209 210 batch_size, seq_length, embed_dim = hidden_states.shape211 212 queries = self.q_proj(hidden_states)213 keys = self.k_proj(hidden_states)214 values = self.v_proj(hidden_states)215 216 queries = queries.view(batch_size, seq_length, -1, self.head_dim).transpose(1, 2)217 keys = keys.view(batch_size, seq_length, -1, self.head_dim).transpose(1, 2)218 values = values.view(batch_size, seq_length, -1, self.head_dim).transpose(1, 2)219 # METACLIP_2 text model uses both `causal_attention_mask` and `attention_mask`220 # in case FA2 kernel is called, `is_causal` should be inferred from `causal_attention_mask`221 if self.config._attn_implementation == "flash_attention_2":222 self.is_causal = causal_attention_mask is not None223 else:224 if attention_mask is not None and causal_attention_mask is not None:225 attention_mask = attention_mask + causal_attention_mask226 elif causal_attention_mask is not None:227 attention_mask = causal_attention_mask228 229 attention_interface: Callable = eager_attention_forward230 if self.config._attn_implementation != "eager":231 attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]232 233 attn_output, attn_weights = attention_interface(234 self,235 queries,236 keys,237 values,238 attention_mask,239 is_causal=self.is_causal,240 scaling=self.scale,241 dropout=0.0 if not self.training else self.dropout,242 output_attentions=output_attentions,243 )244 245 attn_output = attn_output.reshape(batch_size, seq_length, embed_dim).contiguous()246 attn_output = self.out_proj(attn_output)247 248 if not output_attentions:249 attn_weights = None250 return attn_output, attn_weights251 252 253class MetaClip2MLP(nn.Module):254 def __init__(self, config):255 super().__init__()256 self.config = config257 self.activation_fn = ACT2FN[config.hidden_act]258 self.fc1 = nn.Linear(config.hidden_size, config.intermediate_size)259 self.fc2 = nn.Linear(config.intermediate_size, config.hidden_size)260 261 def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:262 hidden_states = self.fc1(hidden_states)263 hidden_states = self.activation_fn(hidden_states)264 hidden_states = self.fc2(hidden_states)265 return hidden_states266 267 268@auto_docstring269class MetaClip2PreTrainedModel(PreTrainedModel):270 config: MetaClip2Config271 base_model_prefix = "metaclip_2"272 supports_gradient_checkpointing = True273 _supports_sdpa = True274 _supports_flash_attn = True275 _supports_flex_attn = True276 _supports_attention_backend = True277 278 def _init_weights(self, module):279 """Initialize the weights"""280 factor = self.config.initializer_factor281 if isinstance(module, MetaClip2TextEmbeddings):282 module.token_embedding.weight.data.normal_(mean=0.0, std=factor * 0.02)283 module.position_embedding.weight.data.normal_(mean=0.0, std=factor * 0.02)284 elif isinstance(module, MetaClip2VisionEmbeddings):285 factor = self.config.initializer_factor286 nn.init.normal_(module.class_embedding, mean=0.0, std=module.embed_dim**-0.5 * factor)287 nn.init.normal_(module.patch_embedding.weight, std=module.config.initializer_range * factor)288 nn.init.normal_(module.position_embedding.weight, std=module.config.initializer_range * factor)289 elif isinstance(module, MetaClip2Attention):290 factor = self.config.initializer_factor291 in_proj_std = (module.embed_dim**-0.5) * ((2 * module.config.num_hidden_layers) ** -0.5) * factor292 out_proj_std = (module.embed_dim**-0.5) * factor293 nn.init.normal_(module.q_proj.weight, std=in_proj_std)294 nn.init.normal_(module.k_proj.weight, std=in_proj_std)295 nn.init.normal_(module.v_proj.weight, std=in_proj_std)296 nn.init.normal_(module.out_proj.weight, std=out_proj_std)297 elif isinstance(module, MetaClip2MLP):298 factor = self.config.initializer_factor299 in_proj_std = (module.config.hidden_size**-0.5) * ((2 * module.config.num_hidden_layers) ** -0.5) * factor300 fc_std = (2 * module.config.hidden_size) ** -0.5 * factor301 nn.init.normal_(module.fc1.weight, std=fc_std)302 nn.init.normal_(module.fc2.weight, std=in_proj_std)303 elif isinstance(module, MetaClip2Model):304 nn.init.normal_(305 module.text_projection.weight,306 std=module.text_embed_dim**-0.5 * self.config.initializer_factor,307 )308 nn.init.normal_(309 module.visual_projection.weight,310 std=module.vision_embed_dim**-0.5 * self.config.initializer_factor,311 )312 elif isinstance(module, MetaClip2VisionModelWithProjection):313 nn.init.normal_(314 module.visual_projection.weight,315 std=self.config.hidden_size**-0.5 * self.config.initializer_factor,316 )317 elif isinstance(module, MetaClip2TextModelWithProjection):318 nn.init.normal_(319 module.text_projection.weight,320 std=self.config.hidden_size**-0.5 * self.config.initializer_factor,321 )322 elif isinstance(module, MetaClip2ForImageClassification):323 nn.init.normal_(324 module.classifier.weight,325 std=self.config.vision_config.hidden_size**-0.5 * self.config.initializer_factor,326 )327 328 if isinstance(module, nn.LayerNorm):329 module.bias.data.zero_()330 module.weight.data.fill_(1.0)331 if isinstance(module, nn.Linear) and module.bias is not None:332 module.bias.data.zero_()333 334 335class MetaClip2EncoderLayer(GradientCheckpointingLayer):336 def __init__(self, config: Union[MetaClip2VisionConfig, MetaClip2TextConfig]):337 super().__init__()338 self.embed_dim = config.hidden_size339 self.self_attn = MetaClip2Attention(config)340 self.layer_norm1 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps)341 self.mlp = MetaClip2MLP(config)342 self.layer_norm2 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps)343 344 def forward(345 self,346 hidden_states: torch.Tensor,347 attention_mask: torch.Tensor,348 causal_attention_mask: torch.Tensor,349 output_attentions: Optional[bool] = False,350 ) -> tuple[torch.FloatTensor]:351 """352 Args:353 hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`354 attention_mask (`torch.FloatTensor`): attention mask of size355 `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values.356 `(config.encoder_attention_heads,)`.357 output_attentions (`bool`, *optional*):358 Whether or not to return the attentions tensors of all attention layers. See `attentions` under359 returned tensors for more detail.360 """361 residual = hidden_states362 363 hidden_states = self.layer_norm1(hidden_states)364 hidden_states, attn_weights = self.self_attn(365 hidden_states=hidden_states,366 attention_mask=attention_mask,367 causal_attention_mask=causal_attention_mask,368 output_attentions=output_attentions,369 )370 hidden_states = residual + hidden_states371 372 residual = hidden_states373 hidden_states = self.layer_norm2(hidden_states)374 hidden_states = self.mlp(hidden_states)375 hidden_states = residual + hidden_states376 377 outputs = (hidden_states,)378 379 if output_attentions:380 outputs += (attn_weights,)381 382 return outputs383 384 385class MetaClip2Encoder(nn.Module):386 """387 Transformer encoder consisting of `config.num_hidden_layers` self attention layers. Each layer is a388 [`MetaClip2EncoderLayer`].389 390 Args:391 config: MetaClip2Config392 """393 394 def __init__(self, config: MetaClip2Config):395 super().__init__()396 self.config = config397 self.layers = nn.ModuleList([MetaClip2EncoderLayer(config) for _ in range(config.num_hidden_layers)])398 self.gradient_checkpointing = False399 400 def forward(401 self,402 inputs_embeds,403 attention_mask: Optional[torch.Tensor] = None,404 causal_attention_mask: Optional[torch.Tensor] = None,405 output_attentions: Optional[bool] = None,406 output_hidden_states: Optional[bool] = None,407 ) -> BaseModelOutput:408 r"""409 Args:410 inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):411 Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation.412 This is useful if you want more control over how to convert `input_ids` indices into associated vectors413 than the model's internal embedding lookup matrix.414 attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):415 Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:416 417 - 1 for tokens that are **not masked**,418 - 0 for tokens that are **masked**.419 420 [What are attention masks?](../glossary#attention-mask)421 causal_attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):422 Causal mask for the text model. Mask values selected in `[0, 1]`:423 424 - 1 for tokens that are **not masked**,425 - 0 for tokens that are **masked**.426 427 [What are attention masks?](../glossary#attention-mask)428 output_attentions (`bool`, *optional*):429 Whether or not to return the attentions tensors of all attention layers. See `attentions` under430 returned tensors for more detail.431 output_hidden_states (`bool`, *optional*):432 Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors433 for more detail.434 return_dict (`bool`, *optional*):435 Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.436 """437 output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions438 output_hidden_states = (439 output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states440 )441 442 encoder_states = () if output_hidden_states else None443 all_attentions = () if output_attentions else None444 445 hidden_states = inputs_embeds446 for idx, encoder_layer in enumerate(self.layers):447 if output_hidden_states:448 encoder_states = encoder_states + (hidden_states,)449 layer_outputs = encoder_layer(450 hidden_states,451 attention_mask,452 causal_attention_mask,453 output_attentions=output_attentions,454 )455 456 hidden_states = layer_outputs[0]457 458 if output_attentions:459 all_attentions = all_attentions + (layer_outputs[1],)460 461 if output_hidden_states:462 encoder_states = encoder_states + (hidden_states,)463 464 return BaseModelOutput(465 last_hidden_state=hidden_states,466 hidden_states=encoder_states,467 attentions=all_attentions,468 )469 470 471class MetaClip2TextTransformer(nn.Module):472 def __init__(self, config: MetaClip2TextConfig):473 super().__init__()474 self.config = config475 embed_dim = config.hidden_size476 self.embeddings = MetaClip2TextEmbeddings(config)477 self.encoder = MetaClip2Encoder(config)478 self.final_layer_norm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps)479 480 # For `pooled_output` computation481 self.eos_token_id = config.eos_token_id482 483 @check_model_inputs(tie_last_hidden_states=False)484 @auto_docstring485 def forward(486 self,487 input_ids,488 attention_mask: Optional[torch.Tensor] = None,489 position_ids: Optional[torch.Tensor] = None,490 use_cache: Optional[bool] = None,491 **kwargs: Unpack[TransformersKwargs],492 ) -> BaseModelOutputWithPooling:493 input_shape = input_ids.size()494 input_ids = input_ids.view(-1, input_shape[-1])495 496 hidden_states = self.embeddings(input_ids=input_ids, position_ids=position_ids)497 498 # CLIP's text model uses causal mask, prepare it here.499 # https://github.com/openai/CLIP/blob/cfcffb90e69f37bf2ff1e988237a0fbe41f33c04/clip/model.py#L324500 causal_attention_mask = _create_4d_causal_attention_mask(501 input_shape, hidden_states.dtype, device=hidden_states.device502 )503 504 # expand attention_mask505 if attention_mask is not None and self.config._attn_implementation != "flash_attention_2":506 # [batch_size, seq_len] -> [batch_size, 1, tgt_seq_len, src_seq_len]507 attention_mask = _prepare_4d_attention_mask(attention_mask, hidden_states.dtype)508 509 encoder_outputs: BaseModelOutput = self.encoder(510 inputs_embeds=hidden_states,511 attention_mask=attention_mask,512 causal_attention_mask=causal_attention_mask,513 **kwargs,514 )515 516 last_hidden_state = encoder_outputs.last_hidden_state517 last_hidden_state = self.final_layer_norm(last_hidden_state)518 519 # Use robust pooling like CLIP - finds the first EOS token position per sequence520 pooled_output = last_hidden_state[521 torch.arange(last_hidden_state.shape[0], device=last_hidden_state.device),522 (input_ids.to(dtype=torch.int, device=last_hidden_state.device) == self.eos_token_id).int().argmax(dim=-1),523 ]524 525 return BaseModelOutputWithPooling(526 last_hidden_state=last_hidden_state,527 pooler_output=pooled_output,528 hidden_states=encoder_outputs.hidden_states,529 attentions=encoder_outputs.attentions,530 )531 532 533@auto_docstring(534 custom_intro="""535 The text model from METACLIP_2 without any head or projection on top.536 """537)538class MetaClip2TextModel(MetaClip2PreTrainedModel):539 """540 The text model from MetaClip2 without any head or projection on top.541 This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the542 library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads543 etc.)544 545 This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.546 Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage547 and behavior.548 549 Args:550 config ([`MetaClip2TextConfig`]): Model configuration class with all the parameters of the model.551 Initializing with a config file does not load the weights associated with the model, only the552 configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.553 554 Examples:555 556 ```python557 >>> from transformers import AutoTokenizer, MetaClip2TextModel558 559 >>> model = MetaClip2TextModel.from_pretrained("facebook/metaclip-2-worldwide-huge-quickgelu")560 >>> tokenizer = AutoTokenizer.from_pretrained("facebook/metaclip-2-worldwide-huge-quickgelu")561 562 >>> inputs = tokenizer(["a photo of a cat", "a photo of a dog"], padding=True, return_tensors="pt")563 564 >>> outputs = model(**inputs)565 >>> last_hidden_state = outputs.last_hidden_state566 >>> pooled_output = outputs.pooler_output # pooled (EOS token) states567 ```"""568 569 config: MetaClip2TextConfig570 571 _no_split_modules = ["MetaClip2TextEmbeddings", "MetaClip2EncoderLayer"]572 _supports_flash_attn = False # mask creation only accounts for sdpa/eager573 574 def __init__(self, config: MetaClip2TextConfig):575 super().__init__(config)576 self.text_model = MetaClip2TextTransformer(config)577 # Initialize weights and apply final processing578 self.post_init()579 580 def get_input_embeddings(self) -> nn.Module:581 return self.text_model.embeddings.token_embedding582 583 def set_input_embeddings(self, value):584 self.text_model.embeddings.token_embedding = value585 586 @can_return_tuple587 @auto_docstring588 def forward(589 self,590 input_ids: Optional[torch.Tensor] = None,591 attention_mask: Optional[torch.Tensor] = None,592 position_ids: Optional[torch.Tensor] = None,593 output_attentions: Optional[bool] = None,594 output_hidden_states: Optional[bool] = None,595 ) -> BaseModelOutputWithPooling:596 r"""597 Examples:598 599 ```python600 >>> from transformers import AutoTokenizer, MetaClip2TextModel601 602 >>> model = MetaClip2TextModel.from_pretrained("facebook/metaclip-2-worldwide-huge-quickgelu")603 >>> tokenizer = AutoTokenizer.from_pretrained("facebook/metaclip-2-worldwide-huge-quickgelu")604 605 >>> inputs = tokenizer(["a photo of a cat", "a photo of a dog"], padding=True, return_tensors="pt")606 607 >>> outputs = model(**inputs)608 >>> last_hidden_state = outputs.last_hidden_state609 >>> pooled_output = outputs.pooler_output # pooled (EOS token) states610 ```"""611 612 return self.text_model(613 input_ids=input_ids,614 attention_mask=attention_mask,615 position_ids=position_ids,616 output_attentions=output_attentions,617 output_hidden_states=output_hidden_states,618 )619 620 621@dataclass622@auto_docstring(623 custom_intro="""624 Base class for text model's outputs that also contains a pooling of the last hidden states.625 """626)627class MetaClip2TextModelOutput(ModelOutput):628 r"""629 text_embeds (`torch.FloatTensor` of shape `(batch_size, output_dim)` *optional* returned when model is initialized with `with_projection=True`):630 The text embeddings obtained by applying the projection layer to the pooler_output.631 """632 633 text_embeds: Optional[torch.FloatTensor] = None634 last_hidden_state: Optional[torch.FloatTensor] = None635 hidden_states: Optional[tuple[torch.FloatTensor, ...]] = None636 attentions: Optional[tuple[torch.FloatTensor, ...]] = None637 638 639@auto_docstring640class MetaClip2TextModelWithProjection(MetaClip2PreTrainedModel):641 """642 MetaClip2 text model with a projection layer on top (a linear layer on top of the pooled output).643 644 This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the645 library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads646 etc.)647 648 This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.649 Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage650 and behavior.651 652 Args:653 config ([`MetaClip2TextConfig`]): Model configuration class with all the parameters of the model.654 Initializing with a config file does not load the weights associated with the model, only the655 configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.656 657 Examples:658 659 ```python660 >>> from transformers import AutoTokenizer, MetaClip2TextModelWithProjection661 662 >>> model = MetaClip2TextModelWithProjection.from_pretrained("facebook/metaclip-2-worldwide-huge-quickgelu")663 >>> tokenizer = AutoTokenizer.from_pretrained("facebook/metaclip-2-worldwide-huge-quickgelu")664 665 >>> inputs = tokenizer(["a photo of a cat", "a photo of a dog"], padding=True, return_tensors="pt")666 667 >>> outputs = model(**inputs)668 >>> text_embeds = outputs.text_embeds669 ```"""670 671 config: MetaClip2TextConfig672 673 _supports_flash_attn = False674 _no_split_modules = ["MetaClip2TextEmbeddings", "MetaClip2EncoderLayer"]675 676 def __init__(self, config: MetaClip2TextConfig):677 super().__init__(config)678 679 text_model = MetaClip2TextModel._from_config(config)680 self.text_model = text_model.text_model681 682 self.text_projection = nn.Linear(config.hidden_size, config.projection_dim, bias=False)683 684 # Initialize weights and apply final processing685 self.post_init()686 687 def get_input_embeddings(self) -> nn.Module:688 return self.text_model.embeddings.token_embedding689 690 def set_input_embeddings(self, value):691 self.text_model.embeddings.token_embedding = value692 693 @can_return_tuple694 @auto_docstring695 def forward(696 self,697 input_ids: Optional[torch.Tensor] = None,698 attention_mask: Optional[torch.Tensor] = None,699 position_ids: Optional[torch.Tensor] = None,700 output_attentions: Optional[bool] = None,701 output_hidden_states: Optional[bool] = None,702 ) -> MetaClip2TextModelOutput:703 r"""704 Examples:705 706 ```python707 >>> from transformers import AutoTokenizer, MetaClip2TextModelWithProjection708 709 >>> model = MetaClip2TextModelWithProjection.from_pretrained("facebook/metaclip-2-worldwide-huge-quickgelu")710 >>> tokenizer = AutoTokenizer.from_pretrained("facebook/metaclip-2-worldwide-huge-quickgelu")711 712 >>> inputs = tokenizer(["a photo of a cat", "a photo of a dog"], padding=True, return_tensors="pt")713 714 >>> outputs = model(**inputs)715 >>> text_embeds = outputs.text_embeds716 ```"""717 718 text_outputs: BaseModelOutputWithPooling = self.text_model(719 input_ids=input_ids,720 attention_mask=attention_mask,721 position_ids=position_ids,722 output_attentions=output_attentions,723 output_hidden_states=output_hidden_states,724 )725 pooled_output = text_outputs.pooler_output726 text_embeds = self.text_projection(pooled_output)727 728 return MetaClip2TextModelOutput(729 text_embeds=text_embeds,730 last_hidden_state=text_outputs.last_hidden_state,731 hidden_states=text_outputs.hidden_states,732 attentions=text_outputs.attentions,733 )734 735 736@dataclass737@auto_docstring738class MetaClip2Output(ModelOutput):739 r"""740 loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `return_loss` is `True`):741 Contrastive loss for image-text similarity.742 logits_per_image (`torch.FloatTensor` of shape `(image_batch_size, text_batch_size)`):743 The scaled dot product scores between `image_embeds` and `text_embeds`. This represents the image-text744 similarity scores.745 logits_per_text (`torch.FloatTensor` of shape `(text_batch_size, image_batch_size)`):746 The scaled dot product scores between `text_embeds` and `image_embeds`. This represents the text-image747 similarity scores.748 text_embeds (`torch.FloatTensor` of shape `(batch_size, output_dim`):749 The text embeddings obtained by applying the projection layer to the pooled output of [`MetaClip2TextModel`].750 image_embeds (`torch.FloatTensor` of shape `(batch_size, output_dim`):751 The image embeddings obtained by applying the projection layer to the pooled output of [`MetaClip2VisionModel`].752 text_model_output (`BaseModelOutputWithPooling`):753 The output of the [`MetaClip2TextModel`].754 vision_model_output (`BaseModelOutputWithPooling`):755 The output of the [`MetaClip2VisionModel`].756 """757 758 loss: Optional[torch.FloatTensor] = None759 logits_per_image: Optional[torch.FloatTensor] = None760 logits_per_text: Optional[torch.FloatTensor] = None761 text_embeds: Optional[torch.FloatTensor] = None762 image_embeds: Optional[torch.FloatTensor] = None763 text_model_output: BaseModelOutputWithPooling = None764 vision_model_output: BaseModelOutputWithPooling = None765 766 def to_tuple(self) -> tuple[Any]:767 return tuple(768 self[k] if k not in ["text_model_output", "vision_model_output"] else getattr(self, k).to_tuple()769 for k in self.keys()770 )771 772 773# contrastive loss function, adapted from774# https://sachinruk.github.io/blog/2021-03-07-metaclip_2.html775def contrastive_loss(logits: torch.Tensor) -> torch.Tensor:776 return nn.functional.cross_entropy(logits, torch.arange(len(logits), device=logits.device))777 778 779def metaclip_2_loss(similarity: torch.Tensor) -> torch.Tensor:780 caption_loss = contrastive_loss(similarity)781 image_loss = contrastive_loss(similarity.t())782 return (caption_loss + image_loss) / 2.0783 784 785def _get_vector_norm(tensor: torch.Tensor) -> torch.Tensor:786 """787 This method is equivalent to tensor.norm(p=2, dim=-1, keepdim=True) and used to make788 model `executorch` exportable. See issue https://github.com/pytorch/executorch/issues/3566789 """790 square_tensor = torch.pow(tensor, 2)791 sum_tensor = torch.sum(square_tensor, dim=-1, keepdim=True)792 normed_tensor = torch.pow(sum_tensor, 0.5)793 return normed_tensor794 795 796@auto_docstring797class MetaClip2Model(MetaClip2PreTrainedModel):798 """799 This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the800 library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads801 etc.)802 803 This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.804 Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage805 and behavior.806 807 Args:808 config ([`MetaClip2Config`]): Model configuration class with all the parameters of the model.809 Initializing with a config file does not load the weights associated with the model, only the810 configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.811 812 Examples:813 814 ```python815 >>> from PIL import Image816 >>> import requests817 >>> from transformers import AutoProcessor, MetaClip2Model818 819 >>> model = MetaClip2Model.from_pretrained("facebook/metaclip-2-worldwide-huge-quickgelu")820 >>> processor = AutoProcessor.from_pretrained("facebook/metaclip-2-worldwide-huge-quickgelu")821 822 >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"823 >>> image = Image.open(requests.get(url, stream=True).raw)824 825 >>> inputs = processor(826 ... text=["a photo of a cat", "a photo of a dog"], images=image, return_tensors="pt", padding=True827 ... )828 829 >>> outputs = model(**inputs)830 >>> logits_per_image = outputs.logits_per_image # this is the image-text similarity score831 >>> probs = logits_per_image.softmax(dim=1) # we can take the softmax to get the label probabilities832 ```"""833 834 config: MetaClip2Config835 _no_split_modules = ["MetaClip2TextEmbeddings", "MetaClip2EncoderLayer", "MetaClip2VisionEmbeddings"]836 _supports_flash_attn = False # mask creation only accounts for sdpa/eager837 838 def __init__(self, config: MetaClip2Config):839 super().__init__(config)840 841 if not isinstance(config.text_config, MetaClip2TextConfig):842 raise TypeError(843 "config.text_config is expected to be of type MetaClip2TextConfig but is of type"844 f" {type(config.text_config)}."845 )846 847 if not isinstance(config.vision_config, MetaClip2VisionConfig):848 raise TypeError(849 "config.vision_config is expected to be of type MetaClip2VisionConfig but is of type"850 f" {type(config.vision_config)}."851 )852 853 text_config = config.text_config854 vision_config = config.vision_config855 856 self.projection_dim = config.projection_dim857 self.text_embed_dim = text_config.hidden_size858 self.vision_embed_dim = vision_config.hidden_size859 860 text_model = MetaClip2TextModel._from_config(text_config)861 self.text_model = text_model.text_model862 863 vision_model = MetaClip2VisionModel._from_config(vision_config)864 self.vision_model = vision_model.vision_model865 866 self.visual_projection = nn.Linear(self.vision_embed_dim, self.projection_dim, bias=False)867 self.text_projection = nn.Linear(self.text_embed_dim, self.projection_dim, bias=False)868 self.logit_scale = nn.Parameter(torch.tensor(self.config.logit_scale_init_value))869 870 # Initialize weights and apply final processing871 self.post_init()872 873 @filter_out_non_signature_kwargs()874 @auto_docstring875 def get_text_features(876 self,877 input_ids: Optional[torch.Tensor] = None,878 attention_mask: Optional[torch.Tensor] = None,879 position_ids: Optional[torch.Tensor] = None,880 output_attentions: Optional[bool] = None,881 output_hidden_states: Optional[bool] = None,882 ) -> torch.FloatTensor:883 r"""884 Returns:885 text_features (`torch.FloatTensor` of shape `(batch_size, output_dim`): The text embeddings obtained by886 applying the projection layer to the pooled output of [`MetaClip2TextModel`].887 888 Examples:889 890 ```python891 >>> from transformers import AutoTokenizer, MetaClip2Model892 893 >>> model = MetaClip2Model.from_pretrained("facebook/metaclip-2-worldwide-huge-quickgelu")894 >>> tokenizer = AutoTokenizer.from_pretrained("facebook/metaclip-2-worldwide-huge-quickgelu")895 896 >>> inputs = tokenizer(["a photo of a cat", "a photo of a dog"], padding=True, return_tensors="pt")897 >>> text_features = model.get_text_features(**inputs)898 ```"""899 text_outputs: BaseModelOutputWithPooling = self.text_model(900 input_ids=input_ids,901 attention_mask=attention_mask,902 position_ids=position_ids,903 )904 pooled_output = text_outputs.pooler_output905 text_features = self.text_projection(pooled_output)906 907 return text_features908 909 @filter_out_non_signature_kwargs()910 @auto_docstring911 def get_image_features(912 self,913 pixel_values: Optional[torch.FloatTensor] = None,914 output_attentions: Optional[bool] = None,915 output_hidden_states: Optional[bool] = None,916 interpolate_pos_encoding: bool = False,917 ) -> torch.FloatTensor:918 r"""919 Returns:920 image_features (`torch.FloatTensor` of shape `(batch_size, output_dim`): The image embeddings obtained by921 applying the projection layer to the pooled output of [`MetaClip2VisionModel`].922 923 Examples:924 925 ```python926 >>> from PIL import Image927 >>> import requests928 >>> from transformers import AutoProcessor, MetaClip2Model929 930 >>> model = MetaClip2Model.from_pretrained("facebook/metaclip-2-worldwide-huge-quickgelu")931 >>> processor = AutoProcessor.from_pretrained("facebook/metaclip-2-worldwide-huge-quickgelu")932 933 >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"934 >>> image = Image.open(requests.get(url, stream=True).raw)935 936 >>> inputs = processor(images=image, return_tensors="pt")937 938 >>> image_features = model.get_image_features(**inputs)939 ```"""940 vision_outputs: BaseModelOutputWithPooling = self.vision_model(941 pixel_values=pixel_values,942 interpolate_pos_encoding=interpolate_pos_encoding,943 )944 pooled_output = vision_outputs.pooler_output945 image_features = self.visual_projection(pooled_output)946 947 return image_features948 949 @can_return_tuple950 @auto_docstring951 def forward(952 self,953 input_ids: Optional[torch.LongTensor] = None,954 pixel_values: Optional[torch.FloatTensor] = None,955 attention_mask: Optional[torch.Tensor] = None,956 position_ids: Optional[torch.LongTensor] = None,957 return_loss: Optional[bool] = None,958 output_attentions: Optional[bool] = None,959 output_hidden_states: Optional[bool] = None,960 interpolate_pos_encoding: bool = False,961 ) -> MetaClip2Output:962 r"""963 return_loss (`bool`, *optional*):964 Whether or not to return the contrastive loss.965 966 Examples:967 968 ```python969 >>> from PIL import Image970 >>> import requests971 >>> from transformers import AutoProcessor, MetaClip2Model972 973 >>> model = MetaClip2Model.from_pretrained("facebook/metaclip-2-worldwide-huge-quickgelu")974 >>> processor = AutoProcessor.from_pretrained("facebook/metaclip-2-worldwide-huge-quickgelu")975 976 >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"977 >>> image = Image.open(requests.get(url, stream=True).raw)978 979 >>> inputs = processor(980 ... text=["a photo of a cat", "a photo of a dog"], images=image, return_tensors="pt", padding=True981 ... )982 983 >>> outputs = model(**inputs)984 >>> logits_per_image = outputs.logits_per_image # this is the image-text similarity score985 >>> probs = logits_per_image.softmax(dim=1) # we can take the softmax to get the label probabilities986 ```"""987 # Use METACLIP_2 model's config for some fields (if specified) instead of those of vision & text components.988 output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions989 output_hidden_states = (990 output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states991 )992 993 vision_outputs: BaseModelOutputWithPooling = self.vision_model(994 pixel_values=pixel_values,995 output_attentions=output_attentions,996 output_hidden_states=output_hidden_states,997 interpolate_pos_encoding=interpolate_pos_encoding,998 )999 1000 text_outputs: BaseModelOutputWithPooling = self.text_model(1001 input_ids=input_ids,1002 attention_mask=attention_mask,1003 position_ids=position_ids,1004 output_attentions=output_attentions,1005 output_hidden_states=output_hidden_states,1006 )1007 1008 image_embeds = vision_outputs.pooler_output1009 image_embeds = self.visual_projection(image_embeds)1010 1011 text_embeds = text_outputs.pooler_output1012 text_embeds = self.text_projection(text_embeds)1013 1014 # normalized features1015 image_embeds = image_embeds / _get_vector_norm(image_embeds)1016 text_embeds = text_embeds / _get_vector_norm(text_embeds)1017 1018 # cosine similarity as logits1019 logits_per_text = torch.matmul(text_embeds, image_embeds.t().to(text_embeds.device))1020 logits_per_text = logits_per_text * self.logit_scale.exp().to(text_embeds.device)1021 1022 logits_per_image = logits_per_text.t()1023 1024 loss = None1025 if return_loss:1026 loss = metaclip_2_loss(logits_per_text)1027 1028 return MetaClip2Output(1029 loss=loss,1030 logits_per_image=logits_per_image,1031 logits_per_text=logits_per_text,1032 text_embeds=text_embeds,1033 image_embeds=image_embeds,1034 text_model_output=text_outputs,1035 vision_model_output=vision_outputs,1036 )1037 1038 1039class MetaClip2VisionTransformer(nn.Module):1040 def __init__(self, config: MetaClip2VisionConfig):1041 super().__init__()1042 self.config = config1043 embed_dim = config.hidden_size1044 1045 self.embeddings = MetaClip2VisionEmbeddings(config)1046 self.pre_layrnorm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps)1047 self.encoder = MetaClip2Encoder(config)1048 self.post_layernorm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps)1049 1050 @auto_docstring1051 def forward(1052 self,1053 pixel_values: Optional[torch.FloatTensor] = None,1054 output_attentions: Optional[bool] = None,1055 output_hidden_states: Optional[bool] = None,1056 interpolate_pos_encoding: Optional[bool] = False,1057 ) -> BaseModelOutputWithPooling:1058 output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions1059 output_hidden_states = (1060 output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states1061 )1062 1063 if pixel_values is None:1064 raise ValueError("You have to specify pixel_values")1065 1066 hidden_states = self.embeddings(pixel_values, interpolate_pos_encoding=interpolate_pos_encoding)1067 hidden_states = self.pre_layrnorm(hidden_states)1068 1069 encoder_outputs: BaseModelOutput = self.encoder(1070 inputs_embeds=hidden_states,1071 output_attentions=output_attentions,1072 output_hidden_states=output_hidden_states,1073 )1074 1075 last_hidden_state = encoder_outputs.last_hidden_state1076 pooled_output = last_hidden_state[:, 0, :]1077 pooled_output = self.post_layernorm(pooled_output)1078 1079 return BaseModelOutputWithPooling(1080 last_hidden_state=last_hidden_state,1081 pooler_output=pooled_output,1082 hidden_states=encoder_outputs.hidden_states,1083 attentions=encoder_outputs.attentions,1084 )1085 1086 1087@auto_docstring(1088 custom_intro="""1089 The vision model from METACLIP_2 without any head or projection on top.1090 """1091)1092class MetaClip2VisionModel(MetaClip2PreTrainedModel):1093 """1094 The vision model from MetaClip2 without any head or projection on top.1095 1096 This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the1097 library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads1098 etc.)1099 1100 This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.1101 Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage1102 and behavior.1103 1104 Args:1105 config ([`MetaClip2VisionConfig`]): Model configuration class with all the parameters of the model.1106 Initializing with a config file does not load the weights associated with the model, only the1107 configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.1108 1109 Examples:1110 1111 ```python1112 >>> from PIL import Image1113 >>> import requests1114 >>> from transformers import AutoProcessor, MetaClip2VisionModel1115 1116 >>> model = MetaClip2VisionModel.from_pretrained("facebook/metaclip-2-worldwide-huge-quickgelu")1117 >>> processor = AutoProcessor.from_pretrained("facebook/metaclip-2-worldwide-huge-quickgelu")1118 1119 >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"1120 >>> image = Image.open(requests.get(url, stream=True).raw)1121 1122 >>> inputs = processor(images=image, return_tensors="pt")1123 1124 >>> outputs = model(**inputs)1125 >>> last_hidden_state = outputs.last_hidden_state1126 >>> pooled_output = outputs.pooler_output # pooled CLS states1127 ```"""1128 1129 config: MetaClip2VisionConfig1130 main_input_name = "pixel_values"1131 _no_split_modules = ["MetaClip2EncoderLayer"]1132 1133 def __init__(self, config: MetaClip2VisionConfig):1134 super().__init__(config)1135 self.vision_model = MetaClip2VisionTransformer(config)1136 # Initialize weights and apply final processing1137 self.post_init()1138 1139 def get_input_embeddings(self) -> nn.Module:1140 return self.vision_model.embeddings.patch_embedding1141 1142 @can_return_tuple1143 @auto_docstring1144 def forward(1145 self,1146 pixel_values: Optional[torch.FloatTensor] = None,1147 output_attentions: Optional[bool] = None,1148 output_hidden_states: Optional[bool] = None,1149 interpolate_pos_encoding: bool = False,1150 ) -> BaseModelOutputWithPooling:1151 r"""1152 Examples:1153 1154 ```python1155 >>> from PIL import Image1156 >>> import requests1157 >>> from transformers import AutoProcessor, MetaClip2VisionModel1158 1159 >>> model = MetaClip2VisionModel.from_pretrained("facebook/metaclip-2-worldwide-huge-quickgelu")1160 >>> processor = AutoProcessor.from_pretrained("facebook/metaclip-2-worldwide-huge-quickgelu")1161 1162 >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"1163 >>> image = Image.open(requests.get(url, stream=True).raw)1164 1165 >>> inputs = processor(images=image, return_tensors="pt")1166 1167 >>> outputs = model(**inputs)1168 >>> last_hidden_state = outputs.last_hidden_state1169 >>> pooled_output = outputs.pooler_output # pooled CLS states1170 ```"""1171 1172 return self.vision_model(1173 pixel_values=pixel_values,1174 output_attentions=output_attentions,1175 output_hidden_states=output_hidden_states,1176 interpolate_pos_encoding=interpolate_pos_encoding,1177 )1178 1179 1180@dataclass1181@auto_docstring(1182 custom_intro="""1183 Base class for vision model's outputs that also contains image embeddings of the pooling of the last hidden states.1184 """1185)1186class MetaClip2VisionModelOutput(ModelOutput):1187 r"""1188 image_embeds (`torch.FloatTensor` of shape `(batch_size, output_dim)` *optional* returned when model is initialized with `with_projection=True`):1189 The image embeddings obtained by applying the projection layer to the pooler_output.1190 """1191 1192 image_embeds: Optional[torch.FloatTensor] = None1193 last_hidden_state: Optional[torch.FloatTensor] = None1194 hidden_states: Optional[tuple[torch.FloatTensor, ...]] = None1195 attentions: Optional[tuple[torch.FloatTensor, ...]] = None1196 1197 1198@auto_docstring1199class MetaClip2VisionModelWithProjection(MetaClip2PreTrainedModel):1200 """