





"""simple docstring"""





import collections
import tempfile
import unittest

import numpy as np

from transformers.testing_utils import (
    is_pt_flax_cross_test,
    require_flax,
    require_torch,
    require_vision,
    slow,
    torch_device,
)
from transformers.utils import is_flax_available, is_torch_available, is_vision_available

from ...test_modeling_flax_common import floats_tensor, ids_tensor, random_attention_mask
from ..bert.test_modeling_flax_bert import FlaxBertModelTester
from ..clip.test_modeling_flax_clip import FlaxCLIPVisionModelTester
from ..vit.test_modeling_flax_vit import FlaxViTModelTester


if is_flax_available():
       from transformers import (
           FlaxBertModel,
           FlaxCLIPVisionModel,
           FlaxVisionTextDualEncoderModel,
           FlaxViTModel,
           VisionTextDualEncoderConfig,
           VisionTextDualEncoderProcessor,
       )
       from transformers.modeling_flax_pytorch_utils import (
           convert_pytorch_state_dict_to_flax,
           load_flax_weights_in_pytorch_model,
       )


if is_torch_available():
       import torch

       from transformers import VisionTextDualEncoderModel

if is_vision_available():
       from PIL import Image


def 			UpperCAmelCase__							(  SCREAMING_SNAKE_CASE     :						List[str]			):







      '''simple docstring'''




      if isinstance(SCREAMING_SNAKE_CASE							,						collections.abc.Iterable			):
            return x
      return (x, x)

@require_flax
class 		lowercase					:





      def     _snake_case			(       self      ,						lowercase      ,						lowercase )		->     List[str]:
            pass





      def     _snake_case			(       self )		->     Any:
            pass





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





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





      def     _snake_case			(       self      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase=None      ,						**lowercase )		->     Optional[int]:
            lowerCAmelCase  						=       VisionTextDualEncoderConfig.from_vision_text_configs(lowercase      ,						lowercase )

            lowerCAmelCase  						=       FlaxVisionTextDualEncoderModel(lowercase )

            lowerCAmelCase  						=       model(input_ids=lowercase      ,						pixel_values=lowercase      ,						attention_mask=lowercase )

            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     _snake_case			(       self      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase=None      ,						**lowercase )		->     Optional[Any]:
            lowerCAmelCase						,		lowerCAmelCase  						=       self.get_vision_text_model(lowercase      ,						lowercase )
            lowerCAmelCase  						=       {"""vision_model""": vision_model, """text_model""": text_model}
            lowerCAmelCase  						=       FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(**lowercase )

            lowerCAmelCase  						=       model(input_ids=lowercase      ,						pixel_values=lowercase      ,						attention_mask=lowercase )

            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     _snake_case			(       self      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase=None      ,						**lowercase )		->     Union[str, Any]:
            lowerCAmelCase						,		lowerCAmelCase  						=       self.get_vision_text_model(lowercase      ,						lowercase )
            lowerCAmelCase  						=       {"""vision_model""": vision_model, """text_model""": text_model}
            lowerCAmelCase  						=       FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(**lowercase )

            lowerCAmelCase  						=       model(input_ids=lowercase      ,						pixel_values=lowercase      ,						attention_mask=lowercase )
            lowerCAmelCase  						=       output[0]

            with tempfile.TemporaryDirectory() as tmpdirname:
                  model.save_pretrained(lowercase )
                  lowerCAmelCase  						=       FlaxVisionTextDualEncoderModel.from_pretrained(lowercase )

                  lowerCAmelCase  						=       model(input_ids=lowercase      ,						pixel_values=lowercase      ,						attention_mask=lowercase )
                  lowerCAmelCase  						=       after_output[0]
                  lowerCAmelCase  						=       np.amax(np.abs(out_a - out_a ) )
                  self.assertLessEqual(lowercase      ,						1e-3 )





      def     _snake_case			(       self      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase=None      ,						**lowercase )		->     Optional[int]:
            lowerCAmelCase						,		lowerCAmelCase  						=       self.get_vision_text_model(lowercase      ,						lowercase )
            lowerCAmelCase  						=       {"""vision_model""": vision_model, """text_model""": text_model}
            lowerCAmelCase  						=       FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(**lowercase )

            lowerCAmelCase  						=       model(
                input_ids=lowercase      ,						pixel_values=lowercase      ,						attention_mask=lowercase      ,						output_attentions=lowercase )

            lowerCAmelCase  						=       output.vision_model_output.attentions
            self.assertEqual(len(lowercase )      ,						vision_config.num_hidden_layers )

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

            lowerCAmelCase  						=       output.text_model_output.attentions
            self.assertEqual(len(lowercase )      ,						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     _snake_case			(       self      ,						lowercase      ,						lowercase      ,						lowercase )		->     Dict:
            pt_model.to(lowercase )
            pt_model.eval()

            # prepare inputs
            lowerCAmelCase  						=       inputs_dict
            lowerCAmelCase  						=       {k: torch.tensor(v.tolist() ) for k, v in flax_inputs.items()}

            with torch.no_grad():
                  lowerCAmelCase  						=       pt_model(**lowercase ).to_tuple()

            lowerCAmelCase  						=       fx_model(**lowercase ).to_tuple()
            self.assertEqual(len(lowercase )      ,						len(lowercase )      ,						"""Output lengths differ between Flax and PyTorch""" )
            for fx_output, pt_output in zip(fx_outputs[:4]      ,						pt_outputs[:4] ):
                  self.assert_almost_equals(lowercase      ,						pt_output.numpy()      ,						4e-2 )

            # PT -> Flax
            with tempfile.TemporaryDirectory() as tmpdirname:
                  pt_model.save_pretrained(lowercase )
                  lowerCAmelCase  						=       FlaxVisionTextDualEncoderModel.from_pretrained(lowercase      ,						from_pt=lowercase )

            lowerCAmelCase  						=       fx_model_loaded(**lowercase ).to_tuple()
            self.assertEqual(len(lowercase )      ,						len(lowercase )      ,						"""Output lengths differ between Flax and PyTorch""" )
            for fx_output_loaded, pt_output in zip(fx_outputs_loaded[:4]      ,						pt_outputs[:4] ):
                  self.assert_almost_equals(lowercase      ,						pt_output.numpy()      ,						4e-2 )

            # Flax -> PT
            with tempfile.TemporaryDirectory() as tmpdirname:
                  fx_model.save_pretrained(lowercase )
                  lowerCAmelCase  						=       VisionTextDualEncoderModel.from_pretrained(lowercase      ,						from_flax=lowercase )

            pt_model_loaded.to(lowercase )
            pt_model_loaded.eval()

            with torch.no_grad():
                  lowerCAmelCase  						=       pt_model_loaded(**lowercase ).to_tuple()

            self.assertEqual(len(lowercase )      ,						len(lowercase )      ,						"""Output lengths differ between Flax and PyTorch""" )
            for fx_output, pt_output_loaded in zip(fx_outputs[:4]      ,						pt_outputs_loaded[:4] ):
                  self.assert_almost_equals(lowercase      ,						pt_output_loaded.numpy()      ,						4e-2 )





      def     _snake_case			(       self      ,						lowercase      ,						lowercase      ,						lowercase )		->     Dict:
            lowerCAmelCase  						=       VisionTextDualEncoderConfig.from_vision_text_configs(lowercase      ,						lowercase )

            lowerCAmelCase  						=       VisionTextDualEncoderModel(lowercase )
            lowerCAmelCase  						=       FlaxVisionTextDualEncoderModel(lowercase )

            lowerCAmelCase  						=       convert_pytorch_state_dict_to_flax(pt_model.state_dict()      ,						lowercase )
            lowerCAmelCase  						=       fx_state

            self.check_pt_flax_equivalence(lowercase      ,						lowercase      ,						lowercase )





      def     _snake_case			(       self      ,						lowercase      ,						lowercase      ,						lowercase )		->     Any:
            lowerCAmelCase  						=       VisionTextDualEncoderConfig.from_vision_text_configs(lowercase      ,						lowercase )

            lowerCAmelCase  						=       VisionTextDualEncoderModel(lowercase )
            lowerCAmelCase  						=       FlaxVisionTextDualEncoderModel(lowercase )

            lowerCAmelCase  						=       load_flax_weights_in_pytorch_model(lowercase      ,						fx_model.params )

            self.check_pt_flax_equivalence(lowercase      ,						lowercase      ,						lowercase )





      def     _snake_case			(       self )		->     Any:
            lowerCAmelCase  						=       self.prepare_config_and_inputs()
            self.check_model_from_pretrained_configs(**lowercase )





      def     _snake_case			(       self )		->     List[Any]:
            lowerCAmelCase  						=       self.prepare_config_and_inputs()
            self.check_vision_text_dual_encoder_from_pretrained(**lowercase )





      def     _snake_case			(       self )		->     Dict:
            lowerCAmelCase  						=       self.prepare_config_and_inputs()
            self.check_save_load(**lowercase )





      def     _snake_case			(       self )		->     Dict:
            lowerCAmelCase  						=       self.prepare_config_and_inputs()
            self.check_vision_text_output_attention(**lowercase )





      @is_pt_flax_cross_test
      def     _snake_case			(       self )		->     str:
            lowerCAmelCase  						=       self.prepare_config_and_inputs()
            lowerCAmelCase  						=       config_inputs_dict.pop("""vision_config""" )
            lowerCAmelCase  						=       config_inputs_dict.pop("""text_config""" )

            lowerCAmelCase  						=       config_inputs_dict

            self.check_equivalence_pt_to_flax(lowercase      ,						lowercase      ,						lowercase )
            self.check_equivalence_flax_to_pt(lowercase      ,						lowercase      ,						lowercase )





      @slow
      def     _snake_case			(       self )		->     Optional[Any]:
            lowerCAmelCase						,		lowerCAmelCase  						=       self.get_pretrained_model_and_inputs()

            lowerCAmelCase  						=       model_a(**lowercase )
            lowerCAmelCase  						=       outputs[0]

            with tempfile.TemporaryDirectory() as tmp_dirname:
                  model_a.save_pretrained(lowercase )
                  lowerCAmelCase  						=       FlaxVisionTextDualEncoderModel.from_pretrained(lowercase )

                  lowerCAmelCase  						=       model_a(**lowercase )
                  lowerCAmelCase  						=       after_outputs[0]
                  lowerCAmelCase  						=       np.amax(np.abs(out_a - out_a ) )
                  self.assertLessEqual(lowercase      ,						1e-5 )







@require_flax
class 		lowercase					(					_UpperCAmelCase     ,  unittest.TestCase					):





      def     _snake_case			(       self )		->     Optional[int]:
            lowerCAmelCase  						=       FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(
                """hf-internal-testing/tiny-random-vit"""      ,						"""hf-internal-testing/tiny-bert"""      ,						vision_from_pt=lowercase      ,						text_from_pt=lowercase      ,						)
            lowerCAmelCase  						=       13
            lowerCAmelCase  						=       floats_tensor(
                [
                    batch_size,
                    model.config.vision_config.num_channels,
                    model.config.vision_config.image_size,
                    model.config.vision_config.image_size,
                ] )
            lowerCAmelCase  						=       ids_tensor([batch_size, 4]      ,						model.config.text_config.vocab_size )
            lowerCAmelCase  						=       random_attention_mask([batch_size, 4] )
            lowerCAmelCase  						=       {"""pixel_values""": pixel_values, """input_ids""": input_ids, """attention_mask""": attention_mask}

            return model, inputs





      def     _snake_case			(       self      ,						lowercase      ,						lowercase )		->     Dict:
            lowerCAmelCase  						=       FlaxViTModel(lowercase )
            lowerCAmelCase  						=       FlaxBertModel(lowercase )
            return vision_model, text_model





      def     _snake_case			(       self )		->     Tuple:
            lowerCAmelCase  						=       FlaxViTModelTester(self )
            lowerCAmelCase  						=       FlaxBertModelTester(self )
            lowerCAmelCase  						=       vit_model_tester.prepare_config_and_inputs()
            lowerCAmelCase  						=       bert_model_tester.prepare_config_and_inputs()

            lowerCAmelCase						,		lowerCAmelCase  						=       vision_config_and_inputs

            lowerCAmelCase						,		lowerCAmelCase						,		lowerCAmelCase						,		lowerCAmelCase  						=       text_config_and_inputs

            # make sure that cross attention layers are added
            return {
                "text_config": text_config,
                "vision_config": vision_config,
                "pixel_values": pixel_values,
                "attention_mask": attention_mask,
                "input_ids": input_ids,
                "token_type_ids": token_type_ids,
            }







@require_torch
class 		lowercase					(					_UpperCAmelCase     ,  unittest.TestCase					):





      def     _snake_case			(       self )		->     str:
            lowerCAmelCase  						=       FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(
                """hf-internal-testing/tiny-random-clip"""      ,						"""hf-internal-testing/tiny-bert"""      ,						vision_from_pt=lowercase      ,						text_from_pt=lowercase      ,						)
            lowerCAmelCase  						=       13
            lowerCAmelCase  						=       floats_tensor(
                [
                    batch_size,
                    model.config.vision_config.num_channels,
                    model.config.vision_config.image_size,
                    model.config.vision_config.image_size,
                ] )
            lowerCAmelCase  						=       ids_tensor([batch_size, 4]      ,						model.config.text_config.vocab_size )
            lowerCAmelCase  						=       random_attention_mask([batch_size, 4] )
            lowerCAmelCase  						=       {"""pixel_values""": pixel_values, """input_ids""": input_ids, """attention_mask""": attention_mask}

            return model, inputs





      def     _snake_case			(       self      ,						lowercase      ,						lowercase )		->     Optional[Any]:
            lowerCAmelCase  						=       FlaxCLIPVisionModel(lowercase )
            lowerCAmelCase  						=       FlaxBertModel(lowercase )
            return vision_model, text_model





      def     _snake_case			(       self )		->     Optional[int]:
            lowerCAmelCase  						=       FlaxCLIPVisionModelTester(self )
            lowerCAmelCase  						=       FlaxBertModelTester(self )
            lowerCAmelCase  						=       clip_model_tester.prepare_config_and_inputs()
            lowerCAmelCase  						=       bert_model_tester.prepare_config_and_inputs()

            lowerCAmelCase						,		lowerCAmelCase  						=       vision_config_and_inputs

            lowerCAmelCase						,		lowerCAmelCase						,		lowerCAmelCase						,		lowerCAmelCase  						=       text_config_and_inputs

            # make sure that cross attention layers are added
            return {
                "text_config": text_config,
                "vision_config": vision_config,
                "pixel_values": pixel_values,
                "attention_mask": attention_mask,
                "input_ids": input_ids,
                "token_type_ids": token_type_ids,
            }







@require_flax
@require_vision
class 		lowercase					(					unittest.TestCase					):





      @slow
      def     _snake_case			(       self )		->     int:
            lowerCAmelCase  						=       FlaxVisionTextDualEncoderModel.from_pretrained("""clip-italian/clip-italian"""      ,						logit_scale_init_value=1.0 )
            lowerCAmelCase  						=       VisionTextDualEncoderProcessor.from_pretrained("""clip-italian/clip-italian""" )

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

            lowerCAmelCase  						=       model(**lowercase )

            # 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])      ,						)

            lowerCAmelCase  						=       np.array([[1.2_284_727, 0.3_104_122]] )

            self.assertTrue(np.allclose(outputs.logits_per_image      ,						lowercase      ,						atol=1e-3 ) )













"""simple docstring"""





from typing import TYPE_CHECKING

# rely on isort to merge the imports
from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available


SCREAMING_SNAKE_CASE__					=   {"configuration_focalnet": ["FOCALNET_PRETRAINED_CONFIG_ARCHIVE_MAP", "FocalNetConfig"]}


try:
       if not is_torch_available():
              raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
       pass
else:
       SCREAMING_SNAKE_CASE__					=   [
           "FOCALNET_PRETRAINED_MODEL_ARCHIVE_LIST",
           "FocalNetForImageClassification",
           "FocalNetForMaskedImageModeling",
           "FocalNetBackbone",
           "FocalNetModel",
           "FocalNetPreTrainedModel",
       ]

if TYPE_CHECKING:
       from .configuration_focalnet import FOCALNET_PRETRAINED_CONFIG_ARCHIVE_MAP, FocalNetConfig

       try:
              if not is_torch_available():
                     raise OptionalDependencyNotAvailable()
       except OptionalDependencyNotAvailable:
              pass
       else:
              from .modeling_focalnet import (
                  FOCALNET_PRETRAINED_MODEL_ARCHIVE_LIST,
                  FocalNetBackbone,
                  FocalNetForImageClassification,
                  FocalNetForMaskedImageModeling,
                  FocalNetModel,
                  FocalNetPreTrainedModel,
              )

else:
       import sys

       SCREAMING_SNAKE_CASE__					=   _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)






