content: # coding=utf-8
  # Copyright 2021 The OpenAI Team Authors and The HuggingFace Team. All rights reserved.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  """ PyTorch CLIP model."""
  
  
  from dataclasses import dataclass
  from typing import Any, Optional, Tuple, Union
  
  import torch
  import torch.utils.checkpoint
  from torch import nn
  from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
  
  from transformers.activations import ACT2FN
  from transformers.modeling_attn_mask_utils import _create_4d_causal_attention_mask, _prepare_4d_attention_mask
  from transformers.modeling_outputs import BaseModelOutput, BaseModelOutputWithPooling, ImageClassifierOutput
  from transformers.modeling_utils import PreTrainedModel
  from transformers.utils import (
      ModelOutput,
      add_code_sample_docstrings,
      add_start_docstrings,
      add_start_docstrings_to_model_forward,
      logging,
      replace_return_docstrings,
  )
  from transformers.models.clip.configuration_clip import CLIPConfig, CLIPTextConfig, CLIPVisionConfig
  
  
  logger = logging.get_logger(__name__)
  
  # General docstring
  _CONFIG_FOR_DOC = "CLIPConfig"
  _CHECKPOINT_FOR_DOC = "openai/clip-vit-base-patch32"
  
  # Image classification docstring
  _IMAGE_CLASS_CHECKPOINT = "openai/clip-vit-base-patch32"
  _IMAGE_CLASS_EXPECTED_OUTPUT = "LABEL_0"
  
  
  # contrastive loss function, adapted from
  # https://sachinruk.github.io/blog/2021-03-07-clip.html
  def contrastive_loss(logits: torch.Tensor) -> torch.Tensor:
      return nn.functional.cross_entropy(logits, torch.arange(len(logits), device=logits.device))
  
  
  def clip_loss(similarity: torch.Tensor) -> torch.Tensor:
      caption_loss = contrastive_loss(similarity)
      image_loss = contrastive_loss(similarity.t())
      return (caption_loss + image_loss) / 2.0
  
  
  @dataclass
  class CLIPVisionModelOutput(ModelOutput):
      """
      Base class for vision model's outputs that also contains image embeddings of the pooling of the last hidden states.
  
      Args:
          image_embeds (`torch.FloatTensor` of shape `(batch_size, output_dim)` *optional* returned when model is initialized with `with_projection=True`):
              The image embeddings obtained by applying the projection layer to the pooler_output.
          last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
              Sequence of hidden-states at the output of the last layer of the model.
          hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
              Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, +
              one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`.
  
              Hidden-states of the model at the output of each layer plus the optional initial embedding outputs.
          attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
              Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
              sequence_length)`.
  
              Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
              heads.
      """
  
      image_embeds: Optional[torch.FloatTensor] = None
      last_hidden_state: torch.FloatTensor = None
      hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None
      attentions: Optional[Tuple[torch.FloatTensor, ...]] = None
  
  
  @dataclass
  class CLIPTextModelOutput(ModelOutput):
      """
      Base class for text model's outputs that also contains a pooling of the last hidden states.
  
      Args:
          text_embeds (`torch.FloatTensor` of shape `(batch_size, output_dim)` *optional* returned when model is initialized with `with_projection=True`):
              The text embeddings obtained by applying the projection layer to the pooler_output.
          last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
              Sequence of hidden-states at the output of the last layer of the model.
          hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
              Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, +
              one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`.
  
              Hidden-states of the model at the output of each layer plus the optional initial embedding outputs.
          attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
              Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
              sequence_length)`.
  
              Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
              heads.
      """
  
      text_embeds: Optional[torch.FloatTensor] = None
      last_hidden_state: torch.FloatTensor = None
      hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None
      attentions: Optional[Tuple[torch.FloatTensor, ...]] = None
  
  
  @dataclass
  class CLIPOutput(ModelOutput):
      """
      Args:
          loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `return_loss` is `True`):
              Contrastive loss for image-text similarity.
          logits_per_image:(`torch.FloatTensor` of shape `(image_batch_size, text_batch_size)`):
              The scaled dot product scores between `image_embeds` and `text_embeds`. This represents the image-text
              similarity scores.
          logits_per_text:(`torch.FloatTensor` of shape `(text_batch_size, image_batch_size)`):
              The scaled dot product scores between `text_embeds` and `image_embeds`. This represents the text-image
              similarity scores.
          text_embeds(`torch.FloatTensor` of shape `(batch_size, output_dim`):
              The text embeddings obtained by applying the projection layer to the pooled output of [`CLIPTextModel`].
          image_embeds(`torch.FloatTensor` of shape `(batch_size, output_dim`):
              The image embeddings obtained by applying the projection layer to the pooled output of [`CLIPVisionModel`].
          text_model_output(`BaseModelOutputWithPooling`):
              The output of the [`CLIPTextModel`].
          vision_model_output(`BaseModelOutputWithPooling`):
              The output of the [`CLIPVisionModel`].
      """
  
      loss: Optional[torch.FloatTensor] = None
      logits_per_image: torch.FloatTensor = None
      logits_per_text: torch.FloatTensor = None
      text_embeds: torch.FloatTensor = None
      image_embeds: torch.FloatTensor = None
      text_model_output: BaseModelOutputWithPooling = None
      vision_model_output: BaseModelOutputWithPooling = None
  
      def to_tuple(self) -> Tuple[Any]:
          return tuple(
              self[k] if k not in ["text_model_output", "vision_model_output"] else getattr(self, k).to_tuple()
              for k in self.keys()
          )
  
  
  class CLIPVisionEmbeddings(nn.Module):
      def __init__(self, config: CLIPVisionConfig):
          super().__init__()
          self.config = config
          self.embed_dim = config.hidden_size
          self.image_size = config.image_size
          self.patch_size = config.patch_size
  
          self.class_embedding = nn.Parameter(torch.randn(self.embed_dim))
  
          self.patch_embedding = nn.Conv2d(
              in_channels=config.num_channels,
              out_channels=self.embed_dim,
              kernel_size=self.patch_size,
              stride=self.patch_size,
              bias=False,
          )
  
          self.num_patches = (self.image_size // self.patch_size) ** 2
          self.num_positions = self.num_patches + 1
          self.position_embedding = nn.Embedding(self.num_positions, self.embed_dim)
          self.register_buffer("position_ids", torch.arange(self.num_positions).expand((1, -1)), persistent=False)
  
      def forward(self, pixel_values: torch.FloatTensor) -> torch.Tensor:
          batch_size = pixel_values.shape[0]
          target_dtype = self.patch_embedding.weight.dtype
          patch_embeds = self.patch_embedding(pixel_values.to(dtype=target_dtype))  # shape = [*, width, grid, grid]
          patch_embeds = patch_embeds.flatten(2).transpose(1, 2)
  
          class_embeds = self.class_embedding.expand(batch_size, 1, -1)
          embeddings = torch.cat([class_embeds, patch_embeds], dim=1)
          embeddings = embeddings + self.position_embedding(self.position_ids)
          return embeddings
  
  
  class CLIPTextEmbeddings(nn.Module):
      def __init__(self, config: CLIPTextConfig):
          super().__init__()
          embed_dim = config.hidden_size
  
          self.token_embedding = nn.Embedding(config.vocab_size, embed_dim)
          self.position_embedding = nn.Embedding(config.max_position_embeddings, embed_dim)
  
          # position_ids (1, len position emb) is contiguous in memory and exported when serialized
          self.register_buffer(
              "position_ids", torch.arange(config.max_position_embeddings).expand((1, -1)), persistent=False
          )
  
      def forward(
          self,
          input_ids: Optional[torch.LongTensor] = None,
          position_ids: Optional[torch.LongTensor] = None,
          inputs_embeds: Optional[torch.FloatTensor] = None,
      ) -> torch.Tensor:
          seq_length = input_ids.shape[-1] if input_ids is not None else inputs_embeds.shape[-2]
  
          if position_ids is None:
              position_ids = self.position_ids[:, :seq_length]
  
          if inputs_embeds is None:
              inputs_embeds = self.token_embedding(input_ids)
  
          position_embeddings = self.position_embedding(position_ids)
          embeddings = inputs_embeds + position_embeddings
  
          return embeddings
  
  
  class CLIPAttention(nn.Module):
      """Multi-headed attention from 'Attention Is All You Need' paper"""
  
      def __init__(self, config):
          super().__init__()
          self.config = config
          self.embed_dim = config.hidden_size
          self.num_heads = config.num_attention_heads
          self.head_dim = self.embed_dim // self.num_heads
          if self.head_dim * self.num_heads != self.embed_dim:
              raise ValueError(
                  f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim} and `num_heads`:"
                  f" {self.num_heads})."
              )
          self.scale = self.head_dim**-0.5
          self.dropout = config.attention_dropout
  
          self.k_proj = nn.Linear(self.embed_dim, self.embed_dim)
          self.v_proj = nn.Linear(self.embed_dim, self.embed_dim)
          self.q_proj = nn.Linear(self.embed_dim, self.embed_dim)
          self.out_proj = nn.Linear(self.embed_dim, self.embed_dim)
  
      def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int):
          return tensor.view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2).contiguous()
  
      def forward(
          self,
          hidden_states: torch.Tensor,
          attention_mask: Optional[torch.Tensor] = None,
          causal_attention_mask: Optional[torch.Tensor] = None,
          output_attentions: Optional[bool] = False,
      ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
          """Input shape: Batch x Time x Channel"""
  
          bsz, tgt_len, embed_dim = hidden_states.size()
  
          # get query proj
          query_states = self.q_proj(hidden_states) * self.scale
          key_states = self._shape(self.k_proj(hidden_states), -1, bsz)
          value_states = self._shape(self.v_proj(hidden_states), -1, bsz)
  
          proj_shape = (bsz * self.num_heads, -1, self.head_dim)
          query_states = self._shape(query_states, tgt_len, bsz).view(*proj_shape)
          key_states = key_states.view(*proj_shape)
          value_states = value_states.view(*proj_shape)
  
          src_len = key_states.size(1)
          attn_weights = torch.bmm(query_states, key_states.transpose(1, 2))
  
          if attn_weights.size() != (bsz * self.num_heads, tgt_len, src_len):
              raise ValueError(
                  f"Attention weights should be of size {(bsz * self.num_heads, tgt_len, src_len)}, but is"
                  f" {attn_weights.size()}"
              )
  
          # apply the causal_attention_mask first
          if causal_attention_mask is not None:
              if causal_attention_mask.size() != (bsz, 1, tgt_len, src_len):
                  raise ValueError(
                      f"Attention mask should be of size {(bsz, 1, tgt_len, src_len)}, but is"
                      f" {causal_attention_mask.size()}"
                  )
              attn_weights = attn_weights.view(bsz, self.num_heads, tgt_len, src_len) + causal_attention_mask
              attn_weights = attn_weights.view(bsz * self.num_heads, tgt_len, src_len)
  
          if attention_mask is not None:
              if attention_mask.size() != (bsz, 1, tgt_len, src_len):
                  raise ValueError(
                      f"Attention mask should be of size {(bsz, 1, tgt_len, src_len)}, but is {attention_mask.size()}"
                  )
              attn_weights = attn_weights.view(bsz, self.num_heads, tgt_len, src_len) + attention_mask
              attn_weights = attn_weights.view(bsz * self.num_heads, tgt_len, src_len)
  
          attn_weights = nn.functional.softmax(attn_weights, dim=-1)
  
          if output_attentions:
              # this operation is a bit akward, but it's required to
              # make sure that attn_weights keeps its gradient.
              # In order to do so, attn_weights have to reshaped
              # twice and have to be reused in the following
              attn_weights_reshaped = attn_weights.view(bsz, self.num_heads, tgt_len, src_len)
              attn_weights = attn_weights_reshaped.view(bsz * self.num_heads, tgt_len, src_len)
          else:
              attn_weights_reshaped = None
  
          attn_probs = nn.functional.dropout(attn_weights, p=self.dropout, training=self.training)
  
          attn_output = torch.bmm(attn_probs, value_states)
  
          if attn_output.size() != (bsz * self.num_heads, tgt_len, self.head_dim):
              raise ValueError(
                  f"`attn_output` should be of size {(bsz, self.num_heads, tgt_len, self.head_dim)}, but is"
                  f" {attn_output.size()}"
              )
  
          attn_output = attn_output.view(bsz, self.num_heads, tgt_len, self.head_dim)
          attn_output = attn_output.transpose(1, 2)
          attn_output = attn_output.reshape(bsz, tgt_len, embed_dim)
  
          attn_output = self.out_proj(attn_output)
  
          return attn_output, attn_weights_reshaped
  
  
  class CLIPMLP(nn.Module):
      def __init__(self, config):
          super().__init__()
          self.config = config
          self.activation_fn = ACT2FN[config.hidden_act]
          self.fc1 = nn.Linear(config.hidden_size, config.intermediate_size)
          self.fc2 = nn.Linear(config.intermediate_size, config.hidden_size)
  
      def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
          hidden_states = self.fc1(hidden_states)
          hidden_states = self.activation_fn(hidden_states)
          hidden_states = self.fc2(hidden_states)
          return hidden_states
  
  
  class CLIPEncoderLayer(nn.Module):
      def __init__(self, config: CLIPConfig):
          super().__init__()
          self.embed_dim = config.hidden_size
          self.self_attn = CLIPAttention(config)
          self.layer_norm1 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps)
          self.mlp = CLIPMLP(config)
          self.layer_norm2 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps)
  
      def forward(
          self,
          hidden_states: torch.Tensor,
          attention_mask: torch.Tensor,
          causal_attention_mask: torch.Tensor,
          output_attentions: Optional[bool] = False,
      ) -> Tuple[torch.FloatTensor]:
          """
          Args:
              hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`
              attention_mask (`torch.FloatTensor`): attention mask of size
                  `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values.
                  `(config.encoder_attention_heads,)`.
              output_attentions (`bool`, *optional*):
                  Whether or not to return the attentions tensors of all attention layers. See `attentions` under
                  returned tensors for more detail.
          """
          residual = hidden_states
  
          hidden_states = self.layer_norm1(hidden_states)
          hidden_states, attn_weights = self.self_attn(
              hidden_states=hidden_states,
              attention_mask=attention_mask,
              causal_attention_mask=causal_attention_mask,
              output_attentions=output_attentions,
          )
          hidden_states = residual + hidden_states
  
          residual = hidden_states
          hidden_states = self.layer_norm2(hidden_states)
          hidden_states = self.mlp(hidden_states)
          hidden_states = residual + hidden_states
  
          outputs = (hidden_states,)
  
          if output_attentions:
              outputs += (attn_weights,)
  
          return outputs
  
  
  class CLIPPreTrainedModel(PreTrainedModel):
      """
      An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
      models.
      """
  
      config_class = CLIPConfig
      base_model_prefix = "clip"
      supports_gradient_checkpointing = True
  
      def _init_weights(self, module):
          """Initialize the weights"""
          factor = self.config.initializer_factor
          if isinstance(module, CLIPTextEmbeddings):
              module.token_embedding.weight.data.normal_(mean=0.0, std=factor * 0.02)
              module.position_embedding.weight.data.normal_(mean=0.0, std=factor * 0.02)
          elif isinstance(module, CLIPVisionEmbeddings):
              factor = self.config.initializer_factor
              nn.init.normal_(module.class_embedding, mean=0.0, std=module.embed_dim**-0.5 * factor)
              nn.init.normal_(module.patch_embedding.weight, std=module.config.initializer_range * factor)
              nn.init.normal_(module.position_embedding.weight, std=module.config.initializer_range * factor)
          elif isinstance(module, CLIPAttention):
              factor = self.config.initializer_factor
              in_proj_std = (module.embed_dim**-0.5) * ((2 * module.config.num_hidden_layers) ** -0.5) * factor
              out_proj_std = (module.embed_dim**-0.5) * factor
              nn.init.normal_(module.q_proj.weight, std=in_proj_std)
              nn.init.normal_(module.k_proj.weight, std=in_proj_std)
              nn.init.normal_(module.v_proj.weight, std=in_proj_std)
              nn.init.normal_(module.out_proj.weight, std=out_proj_std)
          elif isinstance(module, CLIPMLP):
              factor = self.config.initializer_factor
              in_proj_std = (module.config.hidden_size**-0.5) * ((2 * module.config.num_hidden_layers) ** -0.5) * factor
              fc_std = (2 * module.config.hidden_size) ** -0.5 * factor
              nn.init.normal_(module.fc1.weight, std=fc_std)
              nn.init.normal_(module.fc2.weight, std=in_proj_std)
          elif isinstance(module, CLIPModel):
              nn.init.normal_(
                  module.text_projection.weight,
                  std=module.text_embed_dim**-0.5 * self.config.initializer_factor,
              )
              nn.init.normal_(
                  module.visual_projection.weight,
                  std=module.vision_embed_dim**-0.5 * self.config.initializer_factor,
              )
          elif isinstance(module, CLIPVisionModelWithProjection):
              nn.init.normal_(
                  module.visual_projection.weight,
                  std=self.config.hidden_size**-0.5 * self.config.initializer_factor,
              )
          elif isinstance(module, CLIPTextModelWithProjection):
              nn.init.normal_(
                  module.text_projection.weight,
                  std=self.config.hidden_size**-0.5 * self.config.initializer_factor,
              )
          elif isinstance(module, CLIPForImageClassification):
              nn.init.normal_(
                  module.classifier.weight,
                  std=self.config.vision_config.hidden_size**-0.5 * self.config.initializer_factor,
              )
  
          if isinstance(module, nn.LayerNorm):
              module.bias.data.zero_()
              module.weight.data.fill_(1.0)
          if isinstance(module, nn.Linear) and module.bias is not None:
              module.bias.data.zero_()
  
  
  CLIP_START_DOCSTRING = r"""
      This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the
      library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
      etc.)
  
      This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.
      Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage
      and behavior.
  
      Parameters:
          config ([`CLIPConfig`]): Model configuration class with all the parameters of the model.
              Initializing with a config file does not load the weights associated with the model, only the
              configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
  """
  
  CLIP_TEXT_INPUTS_DOCSTRING = r"""
      Args:
          input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
              Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide
              it.
  
              Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
              [`PreTrainedTokenizer.__call__`] for details.
  
              [What are input IDs?](../glossary#input-ids)
          attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
              Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
  
              - 1 for tokens that are **not masked**,
              - 0 for tokens that are **masked**.
  
              [What are attention masks?](../glossary#attention-mask)
          position_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
              Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,
              config.max_position_embeddings - 1]`.
  
              [What are position IDs?](../glossary#position-ids)
          output_attentions (`bool`, *optional*):
              Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
              tensors for more detail.
          output_hidden_states (`bool`, *optional*):
              Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
              more detail.
          return_dict (`bool`, *optional*):
              Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
  """
  
  CLIP_VISION_INPUTS_DOCSTRING = r"""
      Args:
          pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):
              Pixel values. Padding will be ignored by default should you provide it. Pixel values can be obtained using
              [`AutoImageProcessor`]. See [`CLIPImageProcessor.__call__`] for details.
          output_attentions (`bool`, *optional*):
              Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
              tensors for more detail.
          output_hidden_states (`bool`, *optional*):
              Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
              more detail.
          return_dict (`bool`, *optional*):
              Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
  """
  
  CLIP_INPUTS_DOCSTRING = r"""
      Args:
          input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
              Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide
              it.
  
              Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
              [`PreTrainedTokenizer.__call__`] for details.
  
              [What are input IDs?](../glossary#input-ids)
          attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
              Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
  
              - 1 for tokens that are **not masked**,
              - 0 for tokens that are **masked**.
  
              [What are attention masks?](../glossary#attention-mask)
          position_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
              Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,
              config.max_position_embeddings - 1]`.
  
              [What are position IDs?](../glossary#position-ids)
          pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):
              Pixel values. Padding will be ignored by default should you provide it. Pixel values can be obtained using
              [`AutoImageProcessor`]. See [`CLIPImageProcessor.__call__`] for details.
          return_loss (`bool`, *optional*):
              Whether or not to return the contrastive loss.
          output_attentions (`bool`, *optional*):
              Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
              tensors for more detail.
          output_hidden_states (`bool`, *optional*):
              Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
              more detail.
          return_dict (`bool`, *optional*):
              Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
  """
  
  
  class CLIPEncoder(nn.Module):
      """
      Transformer encoder consisting of `config.num_hidden_layers` self attention layers. Each layer is a
      [`CLIPEncoderLayer`].
  
      Args:
          config: CLIPConfig
      """
  
      def __init__(self, config: CLIPConfig):
          super().__init__()
          self.config = config
          self.layers = nn.ModuleList([CLIPEncoderLayer(config) for _ in range(config.num_hidden_layers)])
          self.gradient_checkpointing = False
  
      def forward(
          self,
          inputs_embeds,
          attention_mask: Optional[torch.Tensor] = None,
          causal_attention_mask: Optional[torch.Tensor] = None,
          output_attentions: Optional[bool] = None,
          output_hidden_states: Optional[bool] = None,
          return_dict: Optional[bool] = None,
      ) -> Union[Tuple, BaseModelOutput]:
          r"""
          Args:
              inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
                  Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation.
                  This is useful if you want more control over how to convert `input_ids` indices into associated vectors
                  than the model's internal embedding lookup matrix.
              attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
                  Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
  
                  - 1 for tokens that are **not masked**,
                  - 0 for tokens that are **masked**.
  
                  [What are attention masks?](../glossary#attention-mask)
              causal_attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
                  Causal mask for the text model. Mask values selected in `[0, 1]`:
  
                  - 1 for tokens that are **not masked**,
                  - 0 for tokens that are **masked**.
  
                  [What are attention masks?](../glossary#attention-mask)
              output_attentions (`bool`, *optional*):
                  Whether or not to return the attentions tensors of all attention layers. See `attentions` under
                  returned tensors for more detail.
              output_hidden_states (`bool`, *optional*):
                  Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors
                  for more detail.
              return_dict (`bool`, *optional*):
                  Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
          """
          output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
          output_hidden_states = (
              output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
          )
          return_dict = return_dict if return_dict is not None else self.config.use_return_dict
  
          encoder_states = () if output_hidden_states else None
          all_attentions = () if output_attentions else None
  
          hidden_states = inputs_embeds
          for idx, encoder_layer in enumerate(self.layers):
              if output_hidden_states:
                  encoder_states = encoder_states + (hidden_states,)
              if self.gradient_checkpointing and self.training:
                  layer_outputs = self._gradient_checkpointing_func(
                      encoder_layer.__call__,
                      hidden_states,
                      attention_mask,
                      causal_attention_mask,
                      output_attentions,
                  )
              else:
                  layer_outputs = encoder_layer(
                      hidden_states,
                      attention_mask,
                      causal_attention_mask,
                      output_attentions=output_attentions,
                  )
  
              hidden_states = layer_outputs[0]
  
              if output_attentions:
                  all_attentions = all_attentions + (layer_outputs[1],)
  
          if output_hidden_states:
              encoder_states = encoder_states + (hidden_states,)
  
          if not return_dict:
              return tuple(v for v in [hidden_states, encoder_states, all_attentions] if v is not None)
          return BaseModelOutput(
              last_hidden_state=hidden_states, hidden_states=encoder_states, attentions=all_attentions
          )
  
  
  class CLIPTextTransformer(nn.Module):
      def __init__(self, config: CLIPTextConfig):
          super().__init__()
          self.config = config
          embed_dim = config.hidden_size
          self.embeddings = CLIPTextEmbeddings(config)
          self.encoder = CLIPEncoder(config)
          self.final_layer_norm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps)
  
          # For `pooled_output` computation
          self.eos_token_id = config.eos_token_id
  
      @add_start_docstrings_to_model_forward(CLIP_TEXT_INPUTS_DOCSTRING)
      @replace_return_docstrings(output_type=BaseModelOutputWithPooling, config_class=CLIPTextConfig)
      def forward(
          self,
          input_ids: Optional[torch.Tensor] = None,
          attention_mask: Optional[torch.Tensor] = None,
          position_ids: Optional[torch.Tensor] = None,
          output_attentions: Optional[bool] = None,
          output_hidden_states: Optional[bool] = None,
          return_dict: Optional[bool] = None,
      ) -> Union[Tuple, BaseModelOutputWithPooling]:
          r"""
          Returns:
  
          """
          output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
          output_hidden_states = (
              output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
          )
          return_dict = return_dict if return_dict is not None else self.config.use_return_dict
  
          if input_ids is None:
              raise ValueError("You have to specify input_ids")
  
          input_shape = input_ids.size()
          input_ids = input_ids.view(-1, input_shape[-1])
  
          hidden_states = self.embeddings(input_ids=input_ids, position_ids=position_ids)
  
          # CLIP's text model uses causal mask, prepare it here.
  