
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  __lowerCamelCase					( _lowercase			)       ->      List[Any]:
     if isinstance(_lowercase						,     collections.abc.Iterable			):
          return x
     return (x, x)
@require_flax
class       _lowerCAmelCase  :

 """simple docstring"""

 def      __lowerCAmelCase  (						self			:			Dict					, SCREAMING_SNAKE_CASE__			:			Dict					, SCREAMING_SNAKE_CASE__			:			List[Any]   ):






      """simple docstring"""



      pass

 def      __lowerCAmelCase  (						self			:			Union[str, Any]   ):






      """simple docstring"""



      pass

 def      __lowerCAmelCase  (						self			:			Tuple   ):






      """simple docstring"""



      pass

 def      __lowerCAmelCase  (						self			:			List[str]					, SCREAMING_SNAKE_CASE__			:			np.ndarray					, SCREAMING_SNAKE_CASE__			:			np.ndarray					, SCREAMING_SNAKE_CASE__			:			float   ):






      """simple docstring"""



      UpperCamelCase						    =					np.abs((a - b)   ).max()
      self.assertLessEqual(SCREAMING_SNAKE_CASE__					, SCREAMING_SNAKE_CASE__					, F'Difference between torch and flax is {diff} (>= {tol}).'   )

 def      __lowerCAmelCase  (						self			:			int					, SCREAMING_SNAKE_CASE__			:			List[str]					, SCREAMING_SNAKE_CASE__			:			Optional[Any]					, SCREAMING_SNAKE_CASE__			:			Union[str, Any]					, SCREAMING_SNAKE_CASE__			:			str					, SCREAMING_SNAKE_CASE__			:			str=None					, **SCREAMING_SNAKE_CASE__			:			Optional[int]   ):






      """simple docstring"""



      UpperCamelCase						    =					VisionTextDualEncoderConfig.from_vision_text_configs(SCREAMING_SNAKE_CASE__					, SCREAMING_SNAKE_CASE__   )

      UpperCamelCase						    =					FlaxVisionTextDualEncoderModel(SCREAMING_SNAKE_CASE__   )

      UpperCamelCase						    =					model(input_ids=SCREAMING_SNAKE_CASE__					, pixel_values=SCREAMING_SNAKE_CASE__					, attention_mask=SCREAMING_SNAKE_CASE__   )

      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      __lowerCAmelCase  (						self			:			List[str]					, SCREAMING_SNAKE_CASE__			:			Any					, SCREAMING_SNAKE_CASE__			:			Optional[Any]					, SCREAMING_SNAKE_CASE__			:			Dict					, SCREAMING_SNAKE_CASE__			:			Union[str, Any]					, SCREAMING_SNAKE_CASE__			:			Union[str, Any]=None					, **SCREAMING_SNAKE_CASE__			:			Tuple   ):






      """simple docstring"""



      UpperCamelCase    ,		UpperCamelCase						    =					self.get_vision_text_model(SCREAMING_SNAKE_CASE__					, SCREAMING_SNAKE_CASE__   )
      UpperCamelCase						    =					{'vision_model': vision_model, 'text_model': text_model}
      UpperCamelCase						    =					FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(**SCREAMING_SNAKE_CASE__   )

      UpperCamelCase						    =					model(input_ids=SCREAMING_SNAKE_CASE__					, pixel_values=SCREAMING_SNAKE_CASE__					, attention_mask=SCREAMING_SNAKE_CASE__   )

      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      __lowerCAmelCase  (						self			:			List[str]					, SCREAMING_SNAKE_CASE__			:			str					, SCREAMING_SNAKE_CASE__			:			Tuple					, SCREAMING_SNAKE_CASE__			:			Dict					, SCREAMING_SNAKE_CASE__			:			int					, SCREAMING_SNAKE_CASE__			:			Dict=None					, **SCREAMING_SNAKE_CASE__			:			Dict   ):






      """simple docstring"""



      UpperCamelCase    ,		UpperCamelCase						    =					self.get_vision_text_model(SCREAMING_SNAKE_CASE__					, SCREAMING_SNAKE_CASE__   )
      UpperCamelCase						    =					{'vision_model': vision_model, 'text_model': text_model}
      UpperCamelCase						    =					FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(**SCREAMING_SNAKE_CASE__   )

      UpperCamelCase						    =					model(input_ids=SCREAMING_SNAKE_CASE__					, pixel_values=SCREAMING_SNAKE_CASE__					, attention_mask=SCREAMING_SNAKE_CASE__   )
      UpperCamelCase						    =					output[0]

      with tempfile.TemporaryDirectory() as tmpdirname:
           model.save_pretrained(SCREAMING_SNAKE_CASE__   )
           UpperCamelCase						    =					FlaxVisionTextDualEncoderModel.from_pretrained(SCREAMING_SNAKE_CASE__   )

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

 def      __lowerCAmelCase  (						self			:			str					, SCREAMING_SNAKE_CASE__			:			Any					, SCREAMING_SNAKE_CASE__			:			Optional[Any]					, SCREAMING_SNAKE_CASE__			:			Optional[Any]					, SCREAMING_SNAKE_CASE__			:			Union[str, Any]					, SCREAMING_SNAKE_CASE__			:			Tuple=None					, **SCREAMING_SNAKE_CASE__			:			Optional[int]   ):






      """simple docstring"""



      UpperCamelCase    ,		UpperCamelCase						    =					self.get_vision_text_model(SCREAMING_SNAKE_CASE__					, SCREAMING_SNAKE_CASE__   )
      UpperCamelCase						    =					{'vision_model': vision_model, 'text_model': text_model}
      UpperCamelCase						    =					FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(**SCREAMING_SNAKE_CASE__   )

      UpperCamelCase						    =					model(
          input_ids=SCREAMING_SNAKE_CASE__					, pixel_values=SCREAMING_SNAKE_CASE__					, attention_mask=SCREAMING_SNAKE_CASE__					, output_attentions=SCREAMING_SNAKE_CASE__   )

      UpperCamelCase						    =					output.vision_model_output.attentions
      self.assertEqual(len(SCREAMING_SNAKE_CASE__   )					, vision_config.num_hidden_layers   )

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

      UpperCamelCase						    =					output.text_model_output.attentions
      self.assertEqual(len(SCREAMING_SNAKE_CASE__   )					, 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      __lowerCAmelCase  (						self			:			Tuple					, SCREAMING_SNAKE_CASE__			:			Tuple					, SCREAMING_SNAKE_CASE__			:			str					, SCREAMING_SNAKE_CASE__			:			Dict   ):






      """simple docstring"""



      pt_model.to(SCREAMING_SNAKE_CASE__   )
      pt_model.eval()

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

      with torch.no_grad():
           UpperCamelCase						    =					pt_model(**SCREAMING_SNAKE_CASE__   ).to_tuple()

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

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

      UpperCamelCase						    =					fx_model_loaded(**SCREAMING_SNAKE_CASE__   ).to_tuple()
      self.assertEqual(len(SCREAMING_SNAKE_CASE__   )					, len(SCREAMING_SNAKE_CASE__   )					, '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(SCREAMING_SNAKE_CASE__					, pt_output.numpy()					, 4e-2   )

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

      pt_model_loaded.to(SCREAMING_SNAKE_CASE__   )
      pt_model_loaded.eval()

      with torch.no_grad():
           UpperCamelCase						    =					pt_model_loaded(**SCREAMING_SNAKE_CASE__   ).to_tuple()

      self.assertEqual(len(SCREAMING_SNAKE_CASE__   )					, len(SCREAMING_SNAKE_CASE__   )					, '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(SCREAMING_SNAKE_CASE__					, pt_output_loaded.numpy()					, 4e-2   )

 def      __lowerCAmelCase  (						self			:			Union[str, Any]					, SCREAMING_SNAKE_CASE__			:			Dict					, SCREAMING_SNAKE_CASE__			:			List[str]					, SCREAMING_SNAKE_CASE__			:			int   ):






      """simple docstring"""



      UpperCamelCase						    =					VisionTextDualEncoderConfig.from_vision_text_configs(SCREAMING_SNAKE_CASE__					, SCREAMING_SNAKE_CASE__   )

      UpperCamelCase						    =					VisionTextDualEncoderModel(SCREAMING_SNAKE_CASE__   )
      UpperCamelCase						    =					FlaxVisionTextDualEncoderModel(SCREAMING_SNAKE_CASE__   )

      UpperCamelCase						    =					convert_pytorch_state_dict_to_flax(pt_model.state_dict()					, SCREAMING_SNAKE_CASE__   )
      UpperCamelCase						    =					fx_state

      self.check_pt_flax_equivalence(SCREAMING_SNAKE_CASE__					, SCREAMING_SNAKE_CASE__					, SCREAMING_SNAKE_CASE__   )

 def      __lowerCAmelCase  (						self			:			Tuple					, SCREAMING_SNAKE_CASE__			:			str					, SCREAMING_SNAKE_CASE__			:			Any					, SCREAMING_SNAKE_CASE__			:			Dict   ):






      """simple docstring"""



      UpperCamelCase						    =					VisionTextDualEncoderConfig.from_vision_text_configs(SCREAMING_SNAKE_CASE__					, SCREAMING_SNAKE_CASE__   )

      UpperCamelCase						    =					VisionTextDualEncoderModel(SCREAMING_SNAKE_CASE__   )
      UpperCamelCase						    =					FlaxVisionTextDualEncoderModel(SCREAMING_SNAKE_CASE__   )

      UpperCamelCase						    =					load_flax_weights_in_pytorch_model(SCREAMING_SNAKE_CASE__					, fx_model.params   )

      self.check_pt_flax_equivalence(SCREAMING_SNAKE_CASE__					, SCREAMING_SNAKE_CASE__					, SCREAMING_SNAKE_CASE__   )

 def      __lowerCAmelCase  (						self			:			List[str]   ):






      """simple docstring"""



      UpperCamelCase						    =					self.prepare_config_and_inputs()
      self.check_model_from_pretrained_configs(**SCREAMING_SNAKE_CASE__   )

 def      __lowerCAmelCase  (						self			:			Dict   ):






      """simple docstring"""



      UpperCamelCase						    =					self.prepare_config_and_inputs()
      self.check_vision_text_dual_encoder_from_pretrained(**SCREAMING_SNAKE_CASE__   )

 def      __lowerCAmelCase  (						self			:			Any   ):






      """simple docstring"""



      UpperCamelCase						    =					self.prepare_config_and_inputs()
      self.check_save_load(**SCREAMING_SNAKE_CASE__   )

 def      __lowerCAmelCase  (						self			:			Union[str, Any]   ):






      """simple docstring"""



      UpperCamelCase						    =					self.prepare_config_and_inputs()
      self.check_vision_text_output_attention(**SCREAMING_SNAKE_CASE__   )

 @is_pt_flax_cross_test
 def      __lowerCAmelCase  (						self			:			Union[str, Any]   ):






      """simple docstring"""



      UpperCamelCase						    =					self.prepare_config_and_inputs()
      UpperCamelCase						    =					config_inputs_dict.pop('vision_config'   )
      UpperCamelCase						    =					config_inputs_dict.pop('text_config'   )

      UpperCamelCase						    =					config_inputs_dict

      self.check_equivalence_pt_to_flax(SCREAMING_SNAKE_CASE__					, SCREAMING_SNAKE_CASE__					, SCREAMING_SNAKE_CASE__   )
      self.check_equivalence_flax_to_pt(SCREAMING_SNAKE_CASE__					, SCREAMING_SNAKE_CASE__					, SCREAMING_SNAKE_CASE__   )






 @slow
 def      __lowerCAmelCase  (						self			:			Any   ):






      """simple docstring"""



      UpperCamelCase    ,		UpperCamelCase						    =					self.get_pretrained_model_and_inputs()

      UpperCamelCase						    =					model_a(**SCREAMING_SNAKE_CASE__   )
      UpperCamelCase						    =					outputs[0]

      with tempfile.TemporaryDirectory() as tmp_dirname:
           model_a.save_pretrained(SCREAMING_SNAKE_CASE__   )
           UpperCamelCase						    =					FlaxVisionTextDualEncoderModel.from_pretrained(SCREAMING_SNAKE_CASE__   )

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




@require_flax
class       _lowerCAmelCase  (		__magic_name__			,	unittest.TestCase							):

 """simple docstring"""

 def      __lowerCAmelCase  (						self			:			Any   ):






      """simple docstring"""



      UpperCamelCase						    =					FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(
          'hf-internal-testing/tiny-random-vit'					, 'hf-internal-testing/tiny-bert'					, vision_from_pt=SCREAMING_SNAKE_CASE__					, text_from_pt=SCREAMING_SNAKE_CASE__					, )
      UpperCamelCase						    =					13
      UpperCamelCase						    =					floats_tensor(
          [
              batch_size,
              model.config.vision_config.num_channels,
              model.config.vision_config.image_size,
              model.config.vision_config.image_size,
          ]   )
      UpperCamelCase						    =					ids_tensor([batch_size, 4]					, model.config.text_config.vocab_size   )
      UpperCamelCase						    =					random_attention_mask([batch_size, 4]   )
      UpperCamelCase						    =					{'pixel_values': pixel_values, 'input_ids': input_ids, 'attention_mask': attention_mask}

      return model, inputs

 def      __lowerCAmelCase  (						self			:			Optional[Any]					, SCREAMING_SNAKE_CASE__			:			Dict					, SCREAMING_SNAKE_CASE__			:			List[Any]   ):






      """simple docstring"""



      UpperCamelCase						    =					FlaxViTModel(SCREAMING_SNAKE_CASE__   )
      UpperCamelCase						    =					FlaxBertModel(SCREAMING_SNAKE_CASE__   )
      return vision_model, text_model






 def      __lowerCAmelCase  (						self			:			Optional[int]   ):






      """simple docstring"""



      UpperCamelCase						    =					FlaxViTModelTester(self   )
      UpperCamelCase						    =					FlaxBertModelTester(self   )
      UpperCamelCase						    =					vit_model_tester.prepare_config_and_inputs()
      UpperCamelCase						    =					bert_model_tester.prepare_config_and_inputs()

      UpperCamelCase    ,		UpperCamelCase						    =					vision_config_and_inputs

      UpperCamelCase    ,		UpperCamelCase    ,		UpperCamelCase    ,		UpperCamelCase						    =					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       _lowerCAmelCase  (		__magic_name__			,	unittest.TestCase							):

 """simple docstring"""

 def      __lowerCAmelCase  (						self			:			str   ):






      """simple docstring"""



      UpperCamelCase						    =					FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(
          'hf-internal-testing/tiny-random-clip'					, 'hf-internal-testing/tiny-bert'					, vision_from_pt=SCREAMING_SNAKE_CASE__					, text_from_pt=SCREAMING_SNAKE_CASE__					, )
      UpperCamelCase						    =					13
      UpperCamelCase						    =					floats_tensor(
          [
              batch_size,
              model.config.vision_config.num_channels,
              model.config.vision_config.image_size,
              model.config.vision_config.image_size,
          ]   )
      UpperCamelCase						    =					ids_tensor([batch_size, 4]					, model.config.text_config.vocab_size   )
      UpperCamelCase						    =					random_attention_mask([batch_size, 4]   )
      UpperCamelCase						    =					{'pixel_values': pixel_values, 'input_ids': input_ids, 'attention_mask': attention_mask}

      return model, inputs

 def      __lowerCAmelCase  (						self			:			Dict					, SCREAMING_SNAKE_CASE__			:			Optional[Any]					, SCREAMING_SNAKE_CASE__			:			Union[str, Any]   ):






      """simple docstring"""



      UpperCamelCase						    =					FlaxCLIPVisionModel(SCREAMING_SNAKE_CASE__   )
      UpperCamelCase						    =					FlaxBertModel(SCREAMING_SNAKE_CASE__   )
      return vision_model, text_model






 def      __lowerCAmelCase  (						self			:			str   ):






      """simple docstring"""



      UpperCamelCase						    =					FlaxCLIPVisionModelTester(self   )
      UpperCamelCase						    =					FlaxBertModelTester(self   )
      UpperCamelCase						    =					clip_model_tester.prepare_config_and_inputs()
      UpperCamelCase						    =					bert_model_tester.prepare_config_and_inputs()

      UpperCamelCase    ,		UpperCamelCase						    =					vision_config_and_inputs

      UpperCamelCase    ,		UpperCamelCase    ,		UpperCamelCase    ,		UpperCamelCase						    =					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       _lowerCAmelCase  (		unittest.TestCase							):

 """simple docstring"""

 @slow
 def      __lowerCAmelCase  (						self			:			Union[str, Any]   ):






      """simple docstring"""



      UpperCamelCase						    =					FlaxVisionTextDualEncoderModel.from_pretrained('clip-italian/clip-italian'					, logit_scale_init_value=1.0   )
      UpperCamelCase						    =					VisionTextDualEncoderProcessor.from_pretrained('clip-italian/clip-italian'   )

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

      UpperCamelCase						    =					model(**SCREAMING_SNAKE_CASE__   )

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

      UpperCamelCase						    =					np.array([[1.2284727, 0.3104122]]   )

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




def  __lowerCamelCase					( _lowercase			)       ->      str:
     return "".join(chr(ord(_lowercase			) - 32			) if 'a' <= char <= 'z' else char for char in word			)


if __name__ == "__main__":
 from doctest import testmod

 testmod()


