





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 				SCREAMING_SNAKE_CASE__							(	snake_case__	:Optional[Any]  )   ->    int:
  if isinstance(snake_case__      ,				collections.abc.Iterable  ):
    return x
  return (x, x)





@require_flax
class 					A_					:







       """simple docstring"""





       def  __UpperCAmelCase			(     self						:						str    ,__A						:						List[str]    ,__A						:						Dict	)    ->       Union[str, Any]:
         pass







       def  __UpperCAmelCase			(     self						:						Any	)    ->       int:
         pass







       def  __UpperCAmelCase			(     self						:						Any	)    ->       Tuple:
         pass







       def  __UpperCAmelCase			(     self						:						Union[str, Any]    ,__A						:						np.ndarray    ,__A						:						np.ndarray    ,__A						:						float	)    ->       Optional[Any]:
         _lowercase						    =			np.abs((a - b)	).max()
         self.assertLessEqual(__A    ,__A    ,F"""Difference between torch and flax is {diff} (>= {tol})."""	)







       def  __UpperCAmelCase			(     self						:						Optional[Any]    ,__A						:						Optional[int]    ,__A						:						int    ,__A						:						str    ,__A						:						Optional[Any]    ,__A						:						int=None    ,**__A						:						int	)    ->       Any:
         _lowercase						    =			VisionTextDualEncoderConfig.from_vision_text_configs(__A    ,__A	)

         _lowercase						    =			FlaxVisionTextDualEncoderModel(__A	)

         _lowercase						    =			model(input_ids=__A    ,pixel_values=__A    ,attention_mask=__A	)

         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						:						str    ,__A						:						List[str]    ,__A						:						Optional[Any]    ,__A						:						Union[str, Any]    ,__A						:						int    ,__A						:						List[str]=None    ,**__A						:						Tuple	)    ->       Dict:
         _lowercase    ,		_lowercase						    =			self.get_vision_text_model(__A    ,__A	)
         _lowercase						    =			{'vision_model': vision_model, 'text_model': text_model}
         _lowercase						    =			FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(**__A	)

         _lowercase						    =			model(input_ids=__A    ,pixel_values=__A    ,attention_mask=__A	)

         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						:						List[str]    ,__A						:						Optional[int]    ,__A						:						List[str]    ,__A						:						Optional[int]    ,__A						:						Union[str, Any]    ,__A						:						Union[str, Any]=None    ,**__A						:						Optional[Any]	)    ->       List[Any]:
         _lowercase    ,		_lowercase						    =			self.get_vision_text_model(__A    ,__A	)
         _lowercase						    =			{'vision_model': vision_model, 'text_model': text_model}
         _lowercase						    =			FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(**__A	)

         _lowercase						    =			model(input_ids=__A    ,pixel_values=__A    ,attention_mask=__A	)
         _lowercase						    =			output[0]

         with tempfile.TemporaryDirectory() as tmpdirname:
           model.save_pretrained(__A	)
           _lowercase						    =			FlaxVisionTextDualEncoderModel.from_pretrained(__A	)

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







       def  __UpperCAmelCase			(     self						:						List[str]    ,__A						:						Any    ,__A						:						Union[str, Any]    ,__A						:						Tuple    ,__A						:						Tuple    ,__A						:						Any=None    ,**__A						:						Union[str, Any]	)    ->       Optional[Any]:
         _lowercase    ,		_lowercase						    =			self.get_vision_text_model(__A    ,__A	)
         _lowercase						    =			{'vision_model': vision_model, 'text_model': text_model}
         _lowercase						    =			FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(**__A	)

         _lowercase						    =			model(
             input_ids=__A    ,pixel_values=__A    ,attention_mask=__A    ,output_attentions=__A	)

         _lowercase						    =			output.vision_model_output.attentions
         self.assertEqual(len(__A	)    ,vision_config.num_hidden_layers	)

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

         _lowercase						    =			output.text_model_output.attentions
         self.assertEqual(len(__A	)    ,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						:						Any    ,__A						:						Optional[int]    ,__A						:						Union[str, Any]    ,__A						:						Union[str, Any]	)    ->       List[str]:
         pt_model.to(__A	)
         pt_model.eval()

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

         with torch.no_grad():
           _lowercase						    =			pt_model(**__A	).to_tuple()

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

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

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

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

         pt_model_loaded.to(__A	)
         pt_model_loaded.eval()

         with torch.no_grad():
           _lowercase						    =			pt_model_loaded(**__A	).to_tuple()

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







       def  __UpperCAmelCase			(     self						:						Optional[int]    ,__A						:						Optional[int]    ,__A						:						Tuple    ,__A						:						int	)    ->       Optional[int]:
         _lowercase						    =			VisionTextDualEncoderConfig.from_vision_text_configs(__A    ,__A	)

         _lowercase						    =			VisionTextDualEncoderModel(__A	)
         _lowercase						    =			FlaxVisionTextDualEncoderModel(__A	)

         _lowercase						    =			convert_pytorch_state_dict_to_flax(pt_model.state_dict()    ,__A	)
         _lowercase						    =			fx_state

         self.check_pt_flax_equivalence(__A    ,__A    ,__A	)







       def  __UpperCAmelCase			(     self						:						Any    ,__A						:						int    ,__A						:						Tuple    ,__A						:						Tuple	)    ->       Optional[Any]:
         _lowercase						    =			VisionTextDualEncoderConfig.from_vision_text_configs(__A    ,__A	)

         _lowercase						    =			VisionTextDualEncoderModel(__A	)
         _lowercase						    =			FlaxVisionTextDualEncoderModel(__A	)

         _lowercase						    =			load_flax_weights_in_pytorch_model(__A    ,fx_model.params	)

         self.check_pt_flax_equivalence(__A    ,__A    ,__A	)







       def  __UpperCAmelCase			(     self						:						str	)    ->       Union[str, Any]:
         _lowercase						    =			self.prepare_config_and_inputs()
         self.check_model_from_pretrained_configs(**__A	)







       def  __UpperCAmelCase			(     self						:						int	)    ->       Optional[Any]:
         _lowercase						    =			self.prepare_config_and_inputs()
         self.check_vision_text_dual_encoder_from_pretrained(**__A	)







       def  __UpperCAmelCase			(     self						:						Union[str, Any]	)    ->       Any:
         _lowercase						    =			self.prepare_config_and_inputs()
         self.check_save_load(**__A	)







       def  __UpperCAmelCase			(     self						:						List[str]	)    ->       Any:
         _lowercase						    =			self.prepare_config_and_inputs()
         self.check_vision_text_output_attention(**__A	)







       @is_pt_flax_cross_test
       def  __UpperCAmelCase			(     self						:						Tuple	)    ->       Any:
         _lowercase						    =			self.prepare_config_and_inputs()
         _lowercase						    =			config_inputs_dict.pop('vision_config'	)
         _lowercase						    =			config_inputs_dict.pop('text_config'	)

         _lowercase						    =			config_inputs_dict

         self.check_equivalence_pt_to_flax(__A    ,__A    ,__A	)
         self.check_equivalence_flax_to_pt(__A    ,__A    ,__A	)







       @slow
       def  __UpperCAmelCase			(     self						:						List[Any]	)    ->       Optional[Any]:
         _lowercase    ,		_lowercase						    =			self.get_pretrained_model_and_inputs()

         _lowercase						    =			model_a(**__A	)
         _lowercase						    =			outputs[0]

         with tempfile.TemporaryDirectory() as tmp_dirname:
           model_a.save_pretrained(__A	)
           _lowercase						    =			FlaxVisionTextDualEncoderModel.from_pretrained(__A	)

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





@require_flax
class 					A_					(   UpperCAmelCase       ,						unittest.TestCase						):







       """simple docstring"""





       def  __UpperCAmelCase			(     self						:						Tuple	)    ->       Optional[Any]:
         _lowercase						    =			FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(
             'hf-internal-testing/tiny-random-vit'    ,'hf-internal-testing/tiny-bert'    ,vision_from_pt=__A    ,text_from_pt=__A    ,)
         _lowercase						    =			13
         _lowercase						    =			floats_tensor(
             [
                 batch_size,
                 model.config.vision_config.num_channels,
                 model.config.vision_config.image_size,
                 model.config.vision_config.image_size,
             ]	)
         _lowercase						    =			ids_tensor([batch_size, 4]    ,model.config.text_config.vocab_size	)
         _lowercase						    =			random_attention_mask([batch_size, 4]	)
         _lowercase						    =			{'pixel_values': pixel_values, 'input_ids': input_ids, 'attention_mask': attention_mask}

         return model, inputs







       def  __UpperCAmelCase			(     self						:						Optional[Any]    ,__A						:						Tuple    ,__A						:						Tuple	)    ->       Optional[int]:
         _lowercase						    =			FlaxViTModel(__A	)
         _lowercase						    =			FlaxBertModel(__A	)
         return vision_model, text_model







       def  __UpperCAmelCase			(     self						:						Union[str, Any]	)    ->       Dict:
         _lowercase						    =			FlaxViTModelTester(self	)
         _lowercase						    =			FlaxBertModelTester(self	)
         _lowercase						    =			vit_model_tester.prepare_config_and_inputs()
         _lowercase						    =			bert_model_tester.prepare_config_and_inputs()

         _lowercase    ,		_lowercase						    =			vision_config_and_inputs

         _lowercase    ,		_lowercase    ,		_lowercase    ,		_lowercase						    =			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 					A_					(   UpperCAmelCase       ,						unittest.TestCase						):







       """simple docstring"""





       def  __UpperCAmelCase			(     self						:						List[Any]	)    ->       Tuple:
         _lowercase						    =			FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(
             'hf-internal-testing/tiny-random-clip'    ,'hf-internal-testing/tiny-bert'    ,vision_from_pt=__A    ,text_from_pt=__A    ,)
         _lowercase						    =			13
         _lowercase						    =			floats_tensor(
             [
                 batch_size,
                 model.config.vision_config.num_channels,
                 model.config.vision_config.image_size,
                 model.config.vision_config.image_size,
             ]	)
         _lowercase						    =			ids_tensor([batch_size, 4]    ,model.config.text_config.vocab_size	)
         _lowercase						    =			random_attention_mask([batch_size, 4]	)
         _lowercase						    =			{'pixel_values': pixel_values, 'input_ids': input_ids, 'attention_mask': attention_mask}

         return model, inputs







       def  __UpperCAmelCase			(     self						:						Optional[Any]    ,__A						:						Dict    ,__A						:						Optional[int]	)    ->       str:
         _lowercase						    =			FlaxCLIPVisionModel(__A	)
         _lowercase						    =			FlaxBertModel(__A	)
         return vision_model, text_model







       def  __UpperCAmelCase			(     self						:						Optional[int]	)    ->       Optional[Any]:
         _lowercase						    =			FlaxCLIPVisionModelTester(self	)
         _lowercase						    =			FlaxBertModelTester(self	)
         _lowercase						    =			clip_model_tester.prepare_config_and_inputs()
         _lowercase						    =			bert_model_tester.prepare_config_and_inputs()

         _lowercase    ,		_lowercase						    =			vision_config_and_inputs

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







       """simple docstring"""





       @slow
       def  __UpperCAmelCase			(     self						:						Dict	)    ->       Any:
         _lowercase						    =			FlaxVisionTextDualEncoderModel.from_pretrained('clip-italian/clip-italian'    ,logit_scale_init_value=1.0	)
         _lowercase						    =			VisionTextDualEncoderProcessor.from_pretrained('clip-italian/clip-italian'	)

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

         _lowercase						    =			model(**__A	)

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

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

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




'''simple docstring'''





import warnings

from ...processing_utils import ProcessorMixin
from ...tokenization_utils_base import BatchEncoding



class  snake_case__						(    SCREAMING_SNAKE_CASE_				):
 A__				= ['''image_processor''', '''tokenizer''']
 A__				= '''CLIPImageProcessor'''
 A__				= ('''CLIPTokenizer''', '''CLIPTokenizerFast''')
 def __init__(   self						:							Any      ,     __a						:							str=None      ,     __a						:							List[Any]=None      ,     **__a						:							List[str]  )      ->   Optional[Any]:







       '''simple docstring'''
       __snake_case :   Optional[Any]													=    None
       if "feature_extractor" in kwargs:
             warnings.warn(
                 'The `feature_extractor` argument is deprecated and will be removed in v5, use `image_processor`'
                 ' instead.'      ,     __a      ,     )
             __snake_case :   List[str]													=    kwargs.pop('feature_extractor'  )

       __snake_case :   List[Any]													=    image_processor if image_processor is not None else feature_extractor
       if image_processor is None:
             raise ValueError('You need to specify an `image_processor`.'  )
       if tokenizer is None:
             raise ValueError('You need to specify a `tokenizer`.'  )

       super().__init__(__a      ,     __a  )
 def __call__(   self						:							List[Any]      ,     __a						:							Optional[int]=None      ,     __a						:							Optional[int]=None      ,     __a						:							Union[str, Any]=None      ,     **__a						:							Union[str, Any]  )      ->   List[Any]:







       '''simple docstring'''

       if text is None and images is None:
             raise ValueError('You have to specify either text or images. Both cannot be none.'  )

       if text is not None:
             __snake_case :   Any													=    self.tokenizer(__a      ,     return_tensors=__a      ,     **__a  )

       if images is not None:
             __snake_case :   str													=    self.image_processor(__a      ,     return_tensors=__a      ,     **__a  )

       if text is not None and images is not None:
             __snake_case :   Optional[Any]													=    image_features.pixel_values
             return encoding
       elif text is not None:
             return encoding
       else:
             return BatchEncoding(data=dict(**__a  )      ,     tensor_type=__a  )
 def A_   (   self						:							List[Any]      ,     *__a						:							Dict      ,     **__a						:							Dict  )      ->   Union[str, Any]:







       '''simple docstring'''
       return self.tokenizer.batch_decode(*__a      ,     **__a  )
 def A_   (   self						:							str      ,     *__a						:							Tuple      ,     **__a						:							List[str]  )      ->   int:







       '''simple docstring'''
       return self.tokenizer.decode(*__a      ,     **__a  )
 @property
 def A_   (   self						:							Union[str, Any]  )      ->   Optional[Any]:







       '''simple docstring'''
       __snake_case :   Tuple													=    self.tokenizer.model_input_names
       __snake_case :   Tuple													=    self.image_processor.model_input_names
       return list(dict.fromkeys(tokenizer_input_names + image_processor_input_names  )  )
 @property
 def A_   (   self						:							Any  )      ->   Any:







       '''simple docstring'''
       warnings.warn(
           '`feature_extractor_class` is deprecated and will be removed in v5. Use `image_processor_class` instead.'      ,     __a      ,     )
       return self.image_processor_class





 @property
 def A_   (   self						:							Optional[int]  )      ->   Optional[Any]:







       '''simple docstring'''
       warnings.warn(
           '`feature_extractor` is deprecated and will be removed in v5. Use `image_processor` instead.'      ,     __a      ,     )
       return self.image_processor



