



import collections.abc
from typing import Optional, Tuple, Union

import torch
import torch.utils.checkpoint
from torch import nn
from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss

from ...activations import ACTaFN
from ...modeling_outputs import BaseModelOutputWithNoAttention, ImageClassifierOutputWithNoAttention
from ...modeling_utils import PreTrainedModel
from ...utils import add_code_sample_docstrings, add_start_docstrings, add_start_docstrings_to_model_forward, logging
from .configuration_poolformer import PoolFormerConfig


a__    =	logging.get_logger(__name__)

# General docstring
a__    =	"""PoolFormerConfig"""

# Base docstring
a__    =	"""sail/poolformer_s12"""
a__    =	[1, 5_12, 7, 7]

# Image classification docstring
a__    =	"""sail/poolformer_s12"""
a__    =	"""tabby, tabby cat"""

a__    =	[
    """sail/poolformer_s12""",
    # See all PoolFormer models at https://huggingface.co/models?filter=poolformer
]



def       lowercase       (   SCREAMING_SNAKE_CASE__       :	str				,       SCREAMING_SNAKE_CASE__       :	Optional[int] = 0.0				,       SCREAMING_SNAKE_CASE__       :	Optional[Any] = False				)  ->     str:
 if drop_prob == 0.0 or not training:
  return input
 _snake_case    :  Optional[int]					=					1 - drop_prob
 _snake_case    :  Optional[Any]					=					(input.shape[0],) + (1,) * (input.ndim - 1)  # work with diff dim tensors, not just 2D ConvNets
 _snake_case    :  List[Any]					=					keep_prob + torch.rand(__snake_case				,       dtype=input.dtype				,       device=input.device				)
 random_tensor.floor_()  # binarize
 _snake_case    :  str					=					input.div(__snake_case				) * random_tensor
 return output







class 		snake_case      (	nn.Module  ):



   '''simple docstring'''






   def __init__(					self							:			int				,		lowerCAmelCase							:			Optional[float] = None)				->				Tuple:

    """simple docstring"""







    super().__init__()
    _snake_case    :  Union[str, Any]					=					drop_prob

   def      UpperCamelCase_  (					self							:			Union[str, Any]				,		lowerCAmelCase							:			torch.Tensor)				->				str:

    """simple docstring"""







    return drop_path(_A				,		self.drop_prob				,		self.training)


   def      UpperCamelCase_  (					self							:			List[str])				->				List[str]:

    """simple docstring"""







    return "p={}".format(self.drop_prob)







class 		snake_case      (	nn.Module  ):



   '''simple docstring'''






   def __init__(					self							:			List[Any]				,		lowerCAmelCase							:			int				,		lowerCAmelCase							:			Dict				,		lowerCAmelCase							:			List[str]				,		lowerCAmelCase							:			Any				,		lowerCAmelCase							:			List[str]				,		lowerCAmelCase							:			Union[str, Any]=None)				->				str:

    """simple docstring"""







    super().__init__()
    _snake_case    :  Any					=					patch_size if isinstance(_A				,		collections.abc.Iterable) else (patch_size, patch_size)
    _snake_case    :  Optional[int]					=					stride if isinstance(_A				,		collections.abc.Iterable) else (stride, stride)
    _snake_case    :  Dict					=					padding if isinstance(_A				,		collections.abc.Iterable) else (padding, padding)

    _snake_case    :  Dict					=					nn.Convad(_A				,		_A				,		kernel_size=_A				,		stride=_A				,		padding=_A)
    _snake_case    :  int					=					norm_layer(_A) if norm_layer else nn.Identity()


   def      UpperCamelCase_  (					self							:			str				,		lowerCAmelCase							:			Union[str, Any])				->				List[Any]:

    """simple docstring"""







    _snake_case    :  List[str]					=					self.projection(_A)
    _snake_case    :  Dict					=					self.norm(_A)
    return embeddings







class 		snake_case      (	nn.GroupNorm  ):



   '''simple docstring'''






   def __init__(					self							:			Optional[Any]				,		lowerCAmelCase							:			Any				,		**lowerCAmelCase							:			Union[str, Any])				->				int:

    """simple docstring"""







    super().__init__(1				,		_A				,		**_A)







class 		snake_case      (	nn.Module  ):



   '''simple docstring'''






   def __init__(					self							:			Any				,		lowerCAmelCase							:			List[Any])				->				Union[str, Any]:

    """simple docstring"""







    super().__init__()
    _snake_case    :  List[Any]					=					nn.AvgPoolad(_A				,		stride=1				,		padding=pool_size // 2				,		count_include_pad=_A)


   def      UpperCamelCase_  (					self							:			str				,		lowerCAmelCase							:			Optional[int])				->				List[str]:

    """simple docstring"""







    return self.pool(_A) - hidden_states







class 		snake_case      (	nn.Module  ):



   '''simple docstring'''






   def __init__(					self							:			Any				,		lowerCAmelCase							:			str				,		lowerCAmelCase							:			Tuple				,		lowerCAmelCase							:			List[str]				,		lowerCAmelCase							:			int)				->				int:

    """simple docstring"""







    super().__init__()
    _snake_case    :  Optional[int]					=					nn.Convad(_A				,		_A				,		1)
    _snake_case    :  str					=					nn.Convad(_A				,		_A				,		1)
    _snake_case    :  List[str]					=					PoolFormerDropPath(_A)
    if isinstance(config.hidden_act				,		_A):
     _snake_case    :  Tuple					=					ACTaFN[config.hidden_act]
    else:
     _snake_case    :  List[str]					=					config.hidden_act


   def      UpperCamelCase_  (					self							:			int				,		lowerCAmelCase							:			List[str])				->				str:

    """simple docstring"""







    _snake_case    :  int					=					self.conva(_A)
    _snake_case    :  List[Any]					=					self.act_fn(_A)
    _snake_case    :  Optional[Any]					=					self.drop(_A)
    _snake_case    :  Any					=					self.conva(_A)
    _snake_case    :  List[str]					=					self.drop(_A)

    return hidden_states







class 		snake_case      (	nn.Module  ):



   '''simple docstring'''






   def __init__(					self							:			List[str]				,		lowerCAmelCase							:			Union[str, Any]				,		lowerCAmelCase							:			Union[str, Any]				,		lowerCAmelCase							:			Optional[Any]				,		lowerCAmelCase							:			Optional[int]				,		lowerCAmelCase							:			List[str]				,		lowerCAmelCase							:			Tuple)				->				Dict:

    """simple docstring"""







    super().__init__()
    _snake_case    :  str					=					PoolFormerPooling(_A)
    _snake_case    :  Optional[int]					=					PoolFormerOutput(_A				,		_A				,		_A				,		_A)
    _snake_case    :  Optional[Any]					=					PoolFormerGroupNorm(_A)
    _snake_case    :  Dict					=					PoolFormerGroupNorm(_A)

    # Useful for training neural nets
    _snake_case    :  Union[str, Any]					=					PoolFormerDropPath(_A) if drop_path > 0.0 else nn.Identity()
    _snake_case    :  Tuple					=					config.use_layer_scale
    if config.use_layer_scale:
     _snake_case    :  Tuple					=					nn.Parameter(
         config.layer_scale_init_value * torch.ones((_A))				,		requires_grad=_A)
     _snake_case    :  Dict					=					nn.Parameter(
         config.layer_scale_init_value * torch.ones((_A))				,		requires_grad=_A)


   def      UpperCamelCase_  (					self							:			int				,		lowerCAmelCase							:			List[Any])				->				Dict:

    """simple docstring"""







    if self.use_layer_scale:
     _snake_case    :  str					=					self.pooling(self.before_norm(_A))
     _snake_case    :  List[str]					=					self.layer_scale_a.unsqueeze(-1).unsqueeze(-1) * pooling_output
     # First residual connection
     _snake_case    :  str					=					hidden_states + self.drop_path(_A)
     _snake_case    :  Dict					=					()

     _snake_case    :  str					=					self.output(self.after_norm(_A))
     _snake_case    :  Dict					=					self.layer_scale_a.unsqueeze(-1).unsqueeze(-1) * layer_output
     # Second residual connection
     _snake_case    :  str					=					hidden_states + self.drop_path(_A)

     _snake_case    :  str					=					(output,) + outputs
     return outputs

    else:
     _snake_case    :  int					=					self.drop_path(self.pooling(self.before_norm(_A)))
     # First residual connection
     _snake_case    :  Optional[int]					=					pooling_output + hidden_states
     _snake_case    :  Optional[Any]					=					()

     # Second residual connection inside the PoolFormerOutput block
     _snake_case    :  str					=					self.drop_path(self.output(self.after_norm(_A)))
     _snake_case    :  Any					=					hidden_states + layer_output

     _snake_case    :  Dict					=					(output,) + outputs
     return outputs







class 		snake_case      (	nn.Module  ):



   '''simple docstring'''






   def __init__(					self							:			int				,		lowerCAmelCase							:			Any)				->				Tuple:

    """simple docstring"""







    super().__init__()
    _snake_case    :  Dict					=					config
    # stochastic depth decay rule
    _snake_case    :  int					=					[x.item() for x in torch.linspace(0				,		config.drop_path_rate				,		sum(config.depths))]

    # patch embeddings
    _snake_case    :  Dict					=					[]
    for i in range(config.num_encoder_blocks):
     embeddings.append(
         PoolFormerEmbeddings(
             patch_size=config.patch_sizes[i]				,		stride=config.strides[i]				,		padding=config.padding[i]				,		num_channels=config.num_channels if i == 0 else config.hidden_sizes[i - 1]				,		hidden_size=config.hidden_sizes[i]				,		))
    _snake_case    :  List[Any]					=					nn.ModuleList(_A)

    # Transformer blocks
    _snake_case    :  List[str]					=					[]
    _snake_case    :  Optional[int]					=					0
    for i in range(config.num_encoder_blocks):
     # each block consists of layers
     _snake_case    :  Any					=					[]
     if i != 0:
      cur += config.depths[i - 1]
     for j in range(config.depths[i]):
      layers.append(
          PoolFormerLayer(
              _A				,		num_channels=config.hidden_sizes[i]				,		pool_size=config.pool_size				,		hidden_size=config.hidden_sizes[i]				,		intermediate_size=int(config.hidden_sizes[i] * config.mlp_ratio)				,		drop_path=dpr[cur + j]				,		))
     blocks.append(nn.ModuleList(_A))

    _snake_case    :  str					=					nn.ModuleList(_A)


   def      UpperCamelCase_  (					self							:			List[str]				,		lowerCAmelCase							:			Dict				,		lowerCAmelCase							:			str=False				,		lowerCAmelCase							:			List[Any]=True)				->				List[Any]:

    """simple docstring"""







    _snake_case    :  int					=					() if output_hidden_states else None

    _snake_case    :  Union[str, Any]					=					pixel_values
    for idx, layers in enumerate(zip(self.patch_embeddings				,		self.block)):
     _snake_case	,    _snake_case    :  Optional[Any]					=					layers
     # Get patch embeddings from hidden_states
     _snake_case    :  Optional[Any]					=					embedding_layer(_A)
     # Send the embeddings through the blocks
     for _, blk in enumerate(_A):
      _snake_case    :  Optional[Any]					=					blk(_A)
      _snake_case    :  int					=					layer_outputs[0]

     if output_hidden_states:
      _snake_case    :  Tuple					=					all_hidden_states + (hidden_states,)

    if not return_dict:
     return tuple(v for v in [hidden_states, all_hidden_states] if v is not None)

    return BaseModelOutputWithNoAttention(last_hidden_state=_A				,		hidden_states=_A)







class 		snake_case      (	__lowercase  ):



   '''simple docstring'''






   snake_case_					: List[str]          =		PoolFormerConfig
   snake_case_					: Optional[Any]          =		"""poolformer"""
   snake_case_					: Dict          =		"""pixel_values"""
   snake_case_					: Any          =		True

   def      UpperCamelCase_  (					self							:			Tuple				,		lowerCAmelCase							:			Any)				->				List[Any]:

    """simple docstring"""







    if isinstance(_A				,		(nn.Linear, nn.Convad)):
     module.weight.data.normal_(mean=0.0				,		std=self.config.initializer_range)
     if module.bias is not None:
      module.bias.data.zero_()
    elif isinstance(_A				,		nn.LayerNorm):
     module.bias.data.zero_()
     module.weight.data.fill_(1.0)


   def      UpperCamelCase_  (					self							:			Any				,		lowerCAmelCase							:			List[Any]				,		lowerCAmelCase							:			Optional[Any]=False)				->				Dict:

    """simple docstring"""







    if isinstance(_A				,		_A):
     _snake_case    :  Dict					=					value
a__    =	R"""\n    This model is a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) sub-class. Use\n    it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage and\n    behavior.\n\n    Parameters:\n        config ([`PoolFormerConfig`]): Model configuration class with all the parameters of the model.\n            Initializing with a config file does not load the weights associated with the model, only the\n            configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.\n"""

a__    =	R"""\n    Args:\n        pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):\n            Pixel values. Pixel values can be obtained using [`AutoImageProcessor`]. See\n            [`PoolFormerImageProcessor.__call__`] for details.\n"""







@add_start_docstrings(
    """The bare PoolFormer Model transformer outputting raw hidden-states without any specific head on top."""					,__lowercase					,)
class 		snake_case      (	__lowercase  ):



   '''simple docstring'''






   def __init__(					self							:			Dict				,		lowerCAmelCase							:			Optional[Any])				->				Union[str, Any]:

    """simple docstring"""







    super().__init__(_A)
    _snake_case    :  Dict					=					config

    _snake_case    :  Union[str, Any]					=					PoolFormerEncoder(_A)

    # Initialize weights and apply final processing
    self.post_init()

   def      UpperCamelCase_  (					self							:			int)				->				Union[str, Any]:

    """simple docstring"""







    return self.embeddings.patch_embeddings


   @add_start_docstrings_to_model_forward(_A)
   @add_code_sample_docstrings(
       checkpoint=_CHECKPOINT_FOR_DOC				,		output_type=_A				,		config_class=_CONFIG_FOR_DOC				,		modality="""vision"""				,		expected_output=_EXPECTED_OUTPUT_SHAPE				,		)
   def      UpperCamelCase_  (					self							:			Any				,		lowerCAmelCase							:			Optional[torch.FloatTensor] = None				,		lowerCAmelCase							:			Optional[bool] = None				,		lowerCAmelCase							:			Optional[bool] = None				,		)				->				Any:

    """simple docstring"""







    _snake_case    :  Optional[Any]					=					(
        output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
    )
    _snake_case    :  Dict					=					return_dict if return_dict is not None else self.config.use_return_dict

    if pixel_values is None:
     raise ValueError("""You have to specify pixel_values""")

    _snake_case    :  int					=					self.encoder(
        _A				,		output_hidden_states=_A				,		return_dict=_A				,		)
    _snake_case    :  Union[str, Any]					=					encoder_outputs[0]

    if not return_dict:
     return (sequence_output, None) + encoder_outputs[1:]

    return BaseModelOutputWithNoAttention(
        last_hidden_state=_A				,		hidden_states=encoder_outputs.hidden_states				,		)







class 		snake_case      (	nn.Module  ):



   '''simple docstring'''






   def __init__(					self							:			str				,		lowerCAmelCase							:			List[str])				->				Union[str, Any]:

    """simple docstring"""







    super().__init__()
    _snake_case    :  Any					=					nn.Linear(config.hidden_size				,		config.hidden_size)


   def      UpperCamelCase_  (					self							:			Dict				,		lowerCAmelCase							:			Tuple)				->				Optional[int]:

    """simple docstring"""







    _snake_case    :  str					=					self.dense(_A)
    return output







@add_start_docstrings(
    """\n    PoolFormer Model transformer with an image classification head on top\n    """					,__lowercase					,)
class 		snake_case      (	__lowercase  ):



   '''simple docstring'''






   def __init__(					self							:			Dict				,		lowerCAmelCase							:			int)				->				Tuple:

    """simple docstring"""







    super().__init__(_A)
    _snake_case    :  Any					=					config.num_labels
    _snake_case    :  Dict					=					PoolFormerModel(_A)

    # Final norm
    _snake_case    :  Union[str, Any]					=					PoolFormerGroupNorm(config.hidden_sizes[-1])
    # Classifier head
    _snake_case    :  int					=					(
        nn.Linear(config.hidden_sizes[-1]				,		config.num_labels) if config.num_labels > 0 else nn.Identity()
    )

    # Initialize weights and apply final processing
    self.post_init()


   @add_start_docstrings_to_model_forward(_A)
   @add_code_sample_docstrings(
       checkpoint=_IMAGE_CLASS_CHECKPOINT				,		output_type=_A				,		config_class=_CONFIG_FOR_DOC				,		expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT				,		)
   def      UpperCamelCase_  (					self							:			List[str]				,		lowerCAmelCase							:			Optional[torch.FloatTensor] = None				,		lowerCAmelCase							:			Optional[torch.LongTensor] = None				,		lowerCAmelCase							:			Optional[bool] = None				,		lowerCAmelCase							:			Optional[bool] = None				,		)				->				Any:

    """simple docstring"""







    _snake_case    :  Optional[int]					=					return_dict if return_dict is not None else self.config.use_return_dict

    _snake_case    :  List[str]					=					self.poolformer(
        _A				,		output_hidden_states=_A				,		return_dict=_A				,		)

    _snake_case    :  Tuple					=					outputs[0]

    _snake_case    :  str					=					self.classifier(self.norm(_A).mean([-2, -1]))

    _snake_case    :  Dict					=					None
    if labels is not None:
     if self.config.problem_type is None:
      if self.num_labels == 1:
       _snake_case    :  Optional[Any]					=					"""regression"""
      elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
       _snake_case    :  Tuple					=					"""single_label_classification"""
      else:
       _snake_case    :  List[Any]					=					"""multi_label_classification"""

     if self.config.problem_type == "regression":
      _snake_case    :  List[str]					=					MSELoss()
      if self.num_labels == 1:
       _snake_case    :  Optional[int]					=					loss_fct(logits.squeeze()				,		labels.squeeze())
      else:
       _snake_case    :  Dict					=					loss_fct(_A				,		_A)
     elif self.config.problem_type == "single_label_classification":
      _snake_case    :  Tuple					=					CrossEntropyLoss()
      _snake_case    :  List[Any]					=					loss_fct(logits.view(-1				,		self.num_labels)				,		labels.view(-1))
     elif self.config.problem_type == "multi_label_classification":
      _snake_case    :  List[str]					=					BCEWithLogitsLoss()
      _snake_case    :  List[Any]					=					loss_fct(_A				,		_A)

    if not return_dict:
     _snake_case    :  str					=					(logits,) + outputs[2:]
     return ((loss,) + output) if loss is not None else output

    return ImageClassifierOutputWithNoAttention(loss=_A				,		logits=_A				,		hidden_states=outputs.hidden_states)





from __future__ import annotations

import collections
import tempfile
import unittest

import numpy as np

from transformers.testing_utils import require_tf, require_vision, slow
from transformers.utils import is_tf_available, is_vision_available

from ...test_modeling_tf_common import floats_tensor, ids_tensor, random_attention_mask
from ..bert.test_modeling_tf_bert import TFBertModelTester
from ..clip.test_modeling_tf_clip import TFCLIPVisionModelTester
from ..deit.test_modeling_tf_deit import TFDeiTModelTester
from ..roberta.test_modeling_tf_roberta import TFRobertaModelTester
from ..vit.test_modeling_tf_vit import TFViTModelTester


if is_tf_available():
  from transformers import (
      TFBertModel,
      TFCLIPVisionModel,
      TFDeiTModel,
      TFRobertaModel,
      TFVisionTextDualEncoderModel,
      TFViTModel,
      VisionTextDualEncoderConfig,
  )

if is_vision_available():
  from PIL import Image

  from transformers import VisionTextDualEncoderProcessor



def       lowercase       (   SCREAMING_SNAKE_CASE__       :	Dict				)  ->     Optional[int]:
 if isinstance(SCREAMING_SNAKE_CASE__				,       collections.abc.Iterable				):
  return x
 return (x, x)







@require_tf
class 		snake_case						:



   '''simple docstring'''






   def      UpperCamelCase_  (					self							:			List[str]				,		lowerCAmelCase							:			Tuple				,		lowerCAmelCase							:			List[str])				->				Any:

    """simple docstring"""







    pass

   def      UpperCamelCase_  (					self							:			Dict)				->				Optional[Any]:

    """simple docstring"""







    pass

   def      UpperCamelCase_  (					self							:			str)				->				Union[str, Any]:

    """simple docstring"""







    pass

   def      UpperCamelCase_  (					self							:			Dict				,		lowerCAmelCase							:			Dict				,		lowerCAmelCase							:			Tuple				,		lowerCAmelCase							:			Optional[int]				,		lowerCAmelCase							:			Any				,		lowerCAmelCase							:			str=None				,		**lowerCAmelCase							:			List[Any])				->				List[str]:

    """simple docstring"""







    _snake_case    :  Dict					=					VisionTextDualEncoderConfig.from_vision_text_configs(lowerCAmelCase				,		lowerCAmelCase)

    _snake_case    :  Optional[int]					=					TFVisionTextDualEncoderModel(lowerCAmelCase)

    _snake_case    :  Dict					=					model(input_ids=lowerCAmelCase				,		pixel_values=lowerCAmelCase				,		attention_mask=lowerCAmelCase)

    self.assertEqual(output["""text_embeds"""].shape				,		(input_ids.shape[0], config.projection_dim))
    self.assertEqual(output["""image_embeds"""].shape				,		(pixel_values.shape[0], config.projection_dim))

   def      UpperCamelCase_  (					self							:			List[str]				,		lowerCAmelCase							:			Optional[Any]				,		lowerCAmelCase							:			Dict				,		lowerCAmelCase							:			Union[str, Any]				,		lowerCAmelCase							:			List[str]				,		lowerCAmelCase							:			Any=None				,		**lowerCAmelCase							:			Any)				->				Any:

    """simple docstring"""







    _snake_case	,    _snake_case    :  Dict					=					self.get_vision_text_model(lowerCAmelCase				,		lowerCAmelCase)
    _snake_case    :  Optional[int]					=					TFVisionTextDualEncoderModel(vision_model=lowerCAmelCase				,		text_model=lowerCAmelCase)

    _snake_case    :  List[str]					=					model(input_ids=lowerCAmelCase				,		pixel_values=lowerCAmelCase				,		attention_mask=lowerCAmelCase)

    self.assertEqual(output["""text_embeds"""].shape				,		(input_ids.shape[0], model.config.projection_dim))
    self.assertEqual(output["""image_embeds"""].shape				,		(pixel_values.shape[0], model.config.projection_dim))

   def      UpperCamelCase_  (					self							:			Optional[int]				,		lowerCAmelCase							:			str				,		lowerCAmelCase							:			List[Any]				,		lowerCAmelCase							:			Tuple				,		lowerCAmelCase							:			List[Any]				,		lowerCAmelCase							:			List[Any]=None				,		**lowerCAmelCase							:			List[Any])				->				Tuple:

    """simple docstring"""







    _snake_case	,    _snake_case    :  List[Any]					=					self.get_vision_text_model(lowerCAmelCase				,		lowerCAmelCase)
    _snake_case    :  Optional[Any]					=					{"""vision_model""": vision_model, """text_model""": text_model}
    _snake_case    :  str					=					TFVisionTextDualEncoderModel.from_vision_text_pretrained(**lowerCAmelCase)

    _snake_case    :  List[Any]					=					model(input_ids=lowerCAmelCase				,		pixel_values=lowerCAmelCase				,		attention_mask=lowerCAmelCase)

    self.assertEqual(output["""text_embeds"""].shape				,		(input_ids.shape[0], model.config.projection_dim))
    self.assertEqual(output["""image_embeds"""].shape				,		(pixel_values.shape[0], model.config.projection_dim))

   def      UpperCamelCase_  (					self							:			Optional[Any]				,		lowerCAmelCase							:			Any				,		lowerCAmelCase							:			List[Any]				,		lowerCAmelCase							:			Tuple				,		lowerCAmelCase							:			List[str]				,		lowerCAmelCase							:			Union[str, Any]=None				,		**lowerCAmelCase							:			int)				->				Tuple:

    """simple docstring"""







    _snake_case	,    _snake_case    :  int					=					self.get_vision_text_model(lowerCAmelCase				,		lowerCAmelCase)
    _snake_case    :  Union[str, Any]					=					TFVisionTextDualEncoderModel(vision_model=lowerCAmelCase				,		text_model=lowerCAmelCase)

    _snake_case    :  str					=					model(input_ids=lowerCAmelCase				,		pixel_values=lowerCAmelCase				,		attention_mask=lowerCAmelCase)
    _snake_case    :  Optional[Any]					=					output[0].numpy()

    with tempfile.TemporaryDirectory() as tmpdirname:
     model.save_pretrained(lowerCAmelCase)
     _snake_case    :  Optional[Any]					=					TFVisionTextDualEncoderModel.from_pretrained(lowerCAmelCase)

     _snake_case    :  Tuple					=					model(input_ids=lowerCAmelCase				,		pixel_values=lowerCAmelCase				,		attention_mask=lowerCAmelCase)
     _snake_case    :  Optional[int]					=					after_output[0].numpy()
     _snake_case    :  List[Any]					=					np.amax(np.abs(out_a - out_a))
     self.assertLessEqual(lowerCAmelCase				,		1E-5)

   def      UpperCamelCase_  (					self							:			Optional[int]				,		lowerCAmelCase							:			str				,		lowerCAmelCase							:			Optional[int]				,		lowerCAmelCase							:			Union[str, Any]				,		lowerCAmelCase							:			Union[str, Any]				,		lowerCAmelCase							:			Optional[int]=None				,		**lowerCAmelCase							:			Dict)				->				Any:

    """simple docstring"""







    _snake_case	,    _snake_case    :  List[str]					=					self.get_vision_text_model(lowerCAmelCase				,		lowerCAmelCase)
    _snake_case    :  Any					=					TFVisionTextDualEncoderModel(vision_model=lowerCAmelCase				,		text_model=lowerCAmelCase)

    _snake_case    :  Tuple					=					model(
        input_ids=lowerCAmelCase				,		pixel_values=lowerCAmelCase				,		attention_mask=lowerCAmelCase				,		output_attentions=lowerCAmelCase)

    _snake_case    :  List[Any]					=					output.vision_model_output.attentions
    self.assertEqual(len(lowerCAmelCase)				,		vision_config.num_hidden_layers)

    # in ViT, the seq_len equals the number of patches + 1 (we add 1 for the [CLS] token)
    _snake_case    :  Any					=					to_atuple(vision_model.config.image_size)
    _snake_case    :  int					=					to_atuple(vision_model.config.patch_size)
    _snake_case    :  List[str]					=					(image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0])
    _snake_case    :  Optional[int]					=					num_patches + 1
    self.assertEqual(vision_attentions[0].shape[-3:]				,		(vision_config.num_attention_heads, seq_len, seq_len))

    _snake_case    :  int					=					output.text_model_output.attentions
    self.assertEqual(len(lowerCAmelCase)				,		text_config.num_hidden_layers)

    self.assertEqual(
        text_attentions[0].shape[-3:]				,		(text_config.num_attention_heads, input_ids.shape[-1], input_ids.shape[-1])				,		)

   def      UpperCamelCase_  (					self							:			str				,		lowerCAmelCase							:			np.ndarray				,		lowerCAmelCase							:			np.ndarray				,		lowerCAmelCase							:			float)				->				Dict:

    """simple docstring"""







    _snake_case    :  Union[str, Any]					=					np.abs((a - b)).max()
    self.assertLessEqual(lowerCAmelCase				,		lowerCAmelCase				,		F'''Difference between torch and flax is {diff} (>= {tol}).''')

   def      UpperCamelCase_  (					self							:			Any)				->				int:

    """simple docstring"""







    _snake_case    :  Any					=					self.prepare_config_and_inputs()
    self.check_vision_text_dual_encoder_model(**lowerCAmelCase)

   def      UpperCamelCase_  (					self							:			List[str])				->				Optional[Any]:

    """simple docstring"""







    _snake_case    :  Union[str, Any]					=					self.prepare_config_and_inputs()
    self.check_model_from_pretrained_configs(**lowerCAmelCase)

   def      UpperCamelCase_  (					self							:			str)				->				Union[str, Any]:

    """simple docstring"""







    _snake_case    :  List[Any]					=					self.prepare_config_and_inputs()
    self.check_vision_text_dual_encoder_from_pretrained(**lowerCAmelCase)

   def      UpperCamelCase_  (					self							:			str)				->				Union[str, Any]:

    """simple docstring"""







    _snake_case    :  Optional[Any]					=					self.prepare_config_and_inputs()
    self.check_save_load(**lowerCAmelCase)

   def      UpperCamelCase_  (					self							:			Union[str, Any])				->				List[Any]:

    """simple docstring"""







    _snake_case    :  Tuple					=					self.prepare_config_and_inputs()
    self.check_vision_text_output_attention(**lowerCAmelCase)


   @slow
   def      UpperCamelCase_  (					self							:			str)				->				int:

    """simple docstring"""







    _snake_case	,    _snake_case    :  Dict					=					self.get_pretrained_model_and_inputs()

    _snake_case    :  Optional[Any]					=					model_a(**lowerCAmelCase)
    _snake_case    :  int					=					outputs[0].numpy()

    with tempfile.TemporaryDirectory() as tmp_dirname:
     model_a.save_pretrained(lowerCAmelCase)
     _snake_case    :  Union[str, Any]					=					TFVisionTextDualEncoderModel.from_pretrained(lowerCAmelCase)

     _snake_case    :  Dict					=					model_a(**lowerCAmelCase)
     _snake_case    :  int					=					after_outputs[0].numpy()
     _snake_case    :  str					=					np.amax(np.abs(out_a - out_a))
     self.assertLessEqual(lowerCAmelCase				,		1E-5)







@require_tf
class 		snake_case						(	SCREAMING_SNAKE_CASE_					,unittest.TestCase  ):



   '''simple docstring'''






   def      UpperCamelCase_  (					self							:			List[Any])				->				Tuple:

    """simple docstring"""







    _snake_case    :  List[Any]					=					TFVisionTextDualEncoderModel.from_vision_text_pretrained(
        """hf-internal-testing/tiny-random-vit"""				,		"""hf-internal-testing/tiny-random-bert""")
    _snake_case    :  Union[str, Any]					=					13
    _snake_case    :  int					=					floats_tensor(
        [
            batch_size,
            model.vision_model.config.num_channels,
            model.vision_model.config.image_size,
            model.vision_model.config.image_size,
        ])
    _snake_case    :  Any					=					ids_tensor([batch_size, 4]				,		model.text_model.config.vocab_size)
    _snake_case    :  Tuple					=					random_attention_mask([batch_size, 4])
    _snake_case    :  Tuple					=					{"""pixel_values""": pixel_values, """input_ids""": input_ids, """attention_mask""": attention_mask}

    return model, inputs

   def      UpperCamelCase_  (					self							:			List[str]				,		lowerCAmelCase							:			List[Any]				,		lowerCAmelCase							:			List[Any])				->				Dict:

    """simple docstring"""







    _snake_case    :  List[str]					=					TFViTModel(lowerCAmelCase				,		name="""vision_model""")
    _snake_case    :  List[Any]					=					TFBertModel(lowerCAmelCase				,		name="""text_model""")
    return vision_model, text_model


   def      UpperCamelCase_  (					self							:			List[str])				->				Union[str, Any]:

    """simple docstring"""







    _snake_case    :  Optional[Any]					=					TFViTModelTester(self)
    _snake_case    :  Optional[Any]					=					TFBertModelTester(self)
    _snake_case    :  Optional[int]					=					vit_model_tester.prepare_config_and_inputs()
    _snake_case    :  int					=					bert_model_tester.prepare_config_and_inputs()

    _snake_case	,    _snake_case	,    _snake_case    :  Dict					=					vision_config_and_inputs

    (
        (
        _snake_case
    )	,    (
        _snake_case
    )	,    (
        _snake_case
    )	,    (
        _snake_case
    )	,    (
        _snake_case
    )	,    (
        _snake_case
    )	,    (
        _snake_case
    )	,    
    )    :  List[str]					=					text_config_and_inputs

    return {
        "text_config": text_config,
        "vision_config": vision_config,
        "pixel_values": pixel_values,
        "attention_mask": input_mask,
        "input_ids": input_ids,
        "text_token_type_ids": token_type_ids,
        "text_sequence_labels": sequence_labels,
        "text_token_labels": token_labels,
        "text_choice_labels": choice_labels,
    }







@require_tf
class 		snake_case						(	SCREAMING_SNAKE_CASE_					,unittest.TestCase  ):



   '''simple docstring'''






   def      UpperCamelCase_  (					self							:			str)				->				Tuple:

    """simple docstring"""







    _snake_case    :  Union[str, Any]					=					TFVisionTextDualEncoderModel.from_vision_text_pretrained(
        """Rocketknight1/tiny-random-deit-tf"""				,		"""hf-internal-testing/tiny-random-roberta""")
    _snake_case    :  int					=					13
    _snake_case    :  List[str]					=					floats_tensor(
        [
            batch_size,
            model.vision_model.config.num_channels,
            model.vision_model.config.image_size,
            model.vision_model.config.image_size,
        ])
    _snake_case    :  Union[str, Any]					=					ids_tensor([batch_size, 4]				,		model.text_model.config.vocab_size)
    _snake_case    :  int					=					random_attention_mask([batch_size, 4])
    _snake_case    :  List[str]					=					{"""pixel_values""": pixel_values, """input_ids""": input_ids, """attention_mask""": attention_mask}

    return model, inputs

   def      UpperCamelCase_  (					self							:			str				,		lowerCAmelCase							:			str				,		lowerCAmelCase							:			Tuple				,		lowerCAmelCase							:			int				,		lowerCAmelCase							:			str				,		lowerCAmelCase							:			int=None				,		**lowerCAmelCase							:			Dict)				->				Optional[int]:

    """simple docstring"""







    _snake_case	,    _snake_case    :  Any					=					self.get_vision_text_model(lowerCAmelCase				,		lowerCAmelCase)
    _snake_case    :  int					=					TFVisionTextDualEncoderModel(vision_model=lowerCAmelCase				,		text_model=lowerCAmelCase)

    _snake_case    :  Optional[Any]					=					model(
        input_ids=lowerCAmelCase				,		pixel_values=lowerCAmelCase				,		attention_mask=lowerCAmelCase				,		output_attentions=lowerCAmelCase)

    _snake_case    :  Dict					=					output.vision_model_output.attentions
    self.assertEqual(len(lowerCAmelCase)				,		vision_config.num_hidden_layers)

    # in DEiT, the seq_len equals the number of patches + 2 (we add 2 for the [CLS] and distillation tokens)
    _snake_case    :  List[Any]					=					to_atuple(vision_model.config.image_size)
    _snake_case    :  str					=					to_atuple(vision_model.config.patch_size)
    _snake_case    :  str					=					(image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0])
    _snake_case    :  List[Any]					=					num_patches + 2
    self.assertEqual(vision_attentions[0].shape[-3:]				,		(vision_config.num_attention_heads, seq_len, seq_len))

    _snake_case    :  Dict					=					output.text_model_output.attentions
    self.assertEqual(len(lowerCAmelCase)				,		text_config.num_hidden_layers)

    self.assertEqual(
        text_attentions[0].shape[-3:]				,		(text_config.num_attention_heads, input_ids.shape[-1], input_ids.shape[-1])				,		)

   def      UpperCamelCase_  (					self							:			Union[str, Any]				,		lowerCAmelCase							:			Dict				,		lowerCAmelCase							:			Optional[Any])				->				Any:

    """simple docstring"""







    _snake_case    :  Optional[int]					=					TFDeiTModel(lowerCAmelCase				,		name="""vision_model""")
    _snake_case    :  Optional[int]					=					TFRobertaModel(lowerCAmelCase				,		name="""text_model""")
    return vision_model, text_model


   def      UpperCamelCase_  (					self							:			List[str])				->				int:

    """simple docstring"""







    _snake_case    :  Union[str, Any]					=					TFDeiTModelTester(self)
    _snake_case    :  Tuple					=					TFRobertaModelTester(self)
    _snake_case    :  Optional[int]					=					vit_model_tester.prepare_config_and_inputs()
    _snake_case    :  Optional[Any]					=					bert_model_tester.prepare_config_and_inputs()

    _snake_case	,    _snake_case	,    _snake_case    :  List[Any]					=					vision_config_and_inputs

    (
        (
        _snake_case
    )	,    (
        _snake_case
    )	,    (
        _snake_case
    )	,    (
        _snake_case
    )	,    (
        _snake_case
    )	,    (
        _snake_case
    )	,    (
        _snake_case
    )	,    
    )    :  Dict					=					text_config_and_inputs

    return {
        "text_config": text_config,
        "vision_config": vision_config,
        "pixel_values": pixel_values,
        "attention_mask": input_mask,
        "input_ids": input_ids,
        "text_token_type_ids": token_type_ids,
        "text_sequence_labels": sequence_labels,
        "text_token_labels": token_labels,
        "text_choice_labels": choice_labels,
    }







@require_tf
class 		snake_case						(	SCREAMING_SNAKE_CASE_					,unittest.TestCase  ):



   '''simple docstring'''






   def      UpperCamelCase_  (					self							:			Union[str, Any])				->				Union[str, Any]:

    """simple docstring"""







    _snake_case    :  str					=					TFVisionTextDualEncoderModel.from_vision_text_pretrained(
        """Rocketknight1/tiny-random-clip-tf"""				,		"""hf-internal-testing/tiny-random-bert""")
    _snake_case    :  str					=					13
    _snake_case    :  Dict					=					floats_tensor(
        [
            batch_size,
            model.vision_model.config.num_channels,
            model.vision_model.config.image_size,
            model.vision_model.config.image_size,
        ])
    _snake_case    :  Optional[int]					=					ids_tensor([batch_size, 4]				,		model.text_model.config.vocab_size)
    _snake_case    :  str					=					random_attention_mask([batch_size, 4])
    _snake_case    :  Optional[int]					=					{"""pixel_values""": pixel_values, """input_ids""": input_ids, """attention_mask""": attention_mask}

    return model, inputs

   def      UpperCamelCase_  (					self							:			Optional[Any]				,		lowerCAmelCase							:			Optional[Any]				,		lowerCAmelCase							:			Any)				->				Optional[Any]:

    """simple docstring"""







    _snake_case    :  Union[str, Any]					=					TFCLIPVisionModel(lowerCAmelCase				,		name="""vision_model""")
    _snake_case    :  List[Any]					=					TFBertModel(lowerCAmelCase				,		name="""text_model""")
    return vision_model, text_model


   def      UpperCamelCase_  (					self							:			str)				->				List[Any]:

    """simple docstring"""







    _snake_case    :  List[str]					=					TFCLIPVisionModelTester(self)
    _snake_case    :  int					=					TFBertModelTester(self)
    _snake_case    :  Any					=					clip_model_tester.prepare_config_and_inputs()
    _snake_case    :  str					=					bert_model_tester.prepare_config_and_inputs()

    _snake_case	,    _snake_case    :  Optional[Any]					=					vision_config_and_inputs

    (
        (
        _snake_case
    )	,    (
        _snake_case
    )	,    (
        _snake_case
    )	,    (
        _snake_case
    )	,    (
        _snake_case
    )	,    (
        _snake_case
    )	,    (
        _snake_case
    )	,    
    )    :  Union[str, Any]					=					text_config_and_inputs

    return {
        "text_config": text_config,
        "vision_config": vision_config,
        "pixel_values": pixel_values,
        "attention_mask": input_mask,
        "input_ids": input_ids,
        "text_token_type_ids": token_type_ids,
        "text_sequence_labels": sequence_labels,
        "text_token_labels": token_labels,
        "text_choice_labels": choice_labels,
    }







@require_vision
@require_tf
class 		snake_case						(	unittest.TestCase  ):



   '''simple docstring'''






   @slow
   def      UpperCamelCase_  (					self							:			Dict)				->				str:

    """simple docstring"""







    _snake_case    :  Optional[int]					=					TFVisionTextDualEncoderModel.from_pretrained(
        """clip-italian/clip-italian"""				,		logit_scale_init_value=1.0				,		from_pt=lowerCAmelCase)
    _snake_case    :  str					=					VisionTextDualEncoderProcessor.from_pretrained("""clip-italian/clip-italian""")

    _snake_case    :  List[str]					=					Image.open("""./tests/fixtures/tests_samples/COCO/000000039769.png""")
    _snake_case    :  List[Any]					=					processor(
        text=["""una foto di un gatto""", """una foto di un cane"""]				,		images=lowerCAmelCase				,		padding=lowerCAmelCase				,		return_tensors="""np""")

    _snake_case    :  Optional[int]					=					model(**lowerCAmelCase)

    # verify the logits
    self.assertEqual(outputs.logits_per_image.shape				,		(inputs.pixel_values.shape[0], inputs.input_ids.shape[0]))
    self.assertEqual(
        outputs.logits_per_text.shape				,		(inputs.input_ids.shape[0], inputs.pixel_values.shape[0])				,		)

    _snake_case    :  int					=					np.array([[1.2_284_727, 0.3_104_122]])

    self.assertTrue(np.allclose(outputs.logits_per_image.numpy()				,		lowerCAmelCase				,		atol=1E-3))
