


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 				_lowercase     (		SCREAMING_SNAKE_CASE_ :		int						):





			"""simple docstring"""




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




@require_flax
class  UpperCAmelCase       :




			def 	lowerCamelCase_       (    self				:		Tuple							,							__magic_name__				:		str							,							__magic_name__				:		Dict    ):







						"""simple docstring"""







						pass




			def 	lowerCamelCase_       (    self				:		List[str]    ):







						"""simple docstring"""







						pass




			def 	lowerCamelCase_       (    self				:		Optional[int]    ):







						"""simple docstring"""







						pass




			def 	lowerCamelCase_       (    self				:		str							,							__magic_name__				:		Optional[int]							,							__magic_name__				:		List[Any]							,							__magic_name__				:		Optional[int]    ):







						"""simple docstring"""







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




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







						"""simple docstring"""







						UpperCamelCase								=       VisionTextDualEncoderConfig.from_vision_text_configs(__snake_case							,							__snake_case    )

						UpperCamelCase								=       FlaxVisionTextDualEncoderModel(__snake_case    )

						UpperCamelCase								=       model(input_ids=__snake_case							,							pixel_values=__snake_case							,							attention_mask=__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[Any]							,							__magic_name__				:		int							,							__magic_name__				:		Optional[int]							,							__magic_name__				:		List[Any]							,							__magic_name__				:		Dict							,							__magic_name__				:		Any=None							,							**__magic_name__				:		List[str]    ):







						"""simple docstring"""







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

						UpperCamelCase								=       model(input_ids=__snake_case							,							pixel_values=__snake_case							,							attention_mask=__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				:		str							,							__magic_name__				:		str							,							__magic_name__				:		str							,							__magic_name__				:		Tuple							,							__magic_name__				:		Union[str, Any]							,							__magic_name__				:		List[Any]=None							,							**__magic_name__				:		Optional[Any]    ):







						"""simple docstring"""







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

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

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

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




			def 	lowerCamelCase_       (    self				:		List[Any]							,							__magic_name__				:		Union[str, Any]							,							__magic_name__				:		int							,							__magic_name__				:		int							,							__magic_name__				:		Dict							,							__magic_name__				:		Any=None							,							**__magic_name__				:		str    ):







						"""simple docstring"""







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

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

						UpperCamelCase								=       output.vision_model_output.attentions
						self.assertEqual(len(__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(__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				:		Optional[Any]							,							__magic_name__				:		Optional[Any]							,							__magic_name__				:		int							,							__magic_name__				:		Optional[Any]    ):







						"""simple docstring"""







						pt_model.to(__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(**__snake_case    ).to_tuple()

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

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

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

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

						pt_model_loaded.to(__snake_case    )
						pt_model_loaded.eval()

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

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




			def 	lowerCamelCase_       (    self				:		Tuple							,							__magic_name__				:		Optional[Any]							,							__magic_name__				:		Tuple							,							__magic_name__				:		Tuple    ):







						"""simple docstring"""







						UpperCamelCase								=       VisionTextDualEncoderConfig.from_vision_text_configs(__snake_case							,							__snake_case    )

						UpperCamelCase								=       VisionTextDualEncoderModel(__snake_case    )
						UpperCamelCase								=       FlaxVisionTextDualEncoderModel(__snake_case    )

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

						self.check_pt_flax_equivalence(__snake_case							,							__snake_case							,							__snake_case    )




			def 	lowerCamelCase_       (    self				:		Optional[int]							,							__magic_name__				:		List[str]							,							__magic_name__				:		str							,							__magic_name__				:		Optional[int]    ):







						"""simple docstring"""







						UpperCamelCase								=       VisionTextDualEncoderConfig.from_vision_text_configs(__snake_case							,							__snake_case    )

						UpperCamelCase								=       VisionTextDualEncoderModel(__snake_case    )
						UpperCamelCase								=       FlaxVisionTextDualEncoderModel(__snake_case    )

						UpperCamelCase								=       load_flax_weights_in_pytorch_model(__snake_case							,							fx_model.params    )

						self.check_pt_flax_equivalence(__snake_case							,							__snake_case							,							__snake_case    )




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







						"""simple docstring"""







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




			def 	lowerCamelCase_       (    self				:		List[str]    ):







						"""simple docstring"""







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




			def 	lowerCamelCase_       (    self				:		Dict    ):







						"""simple docstring"""







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




			def 	lowerCamelCase_       (    self				:		str    ):







						"""simple docstring"""







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




			@is_pt_flax_cross_test
			def 	lowerCamelCase_       (    self				:		int    ):







						"""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(__snake_case							,							__snake_case							,							__snake_case    )
						self.check_equivalence_flax_to_pt(__snake_case							,							__snake_case							,							__snake_case    )





			@slow
			def 	lowerCamelCase_       (    self				:		Dict    ):







						"""simple docstring"""







						UpperCamelCase	,	UpperCamelCase								=       self.get_pretrained_model_and_inputs()

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

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

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




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




			def 	lowerCamelCase_       (    self				:		Dict    ):







						"""simple docstring"""







						UpperCamelCase								=       FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(
						    """hf-internal-testing/tiny-random-vit"""							,							"""hf-internal-testing/tiny-bert"""							,							vision_from_pt=__snake_case							,							text_from_pt=__snake_case							,							)
						UpperCamelCase								=       1_3
						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				:		str							,							__magic_name__				:		Tuple							,							__magic_name__				:		int    ):







						"""simple docstring"""







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





			def 	lowerCamelCase_       (    self				:		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  UpperCAmelCase       (							_SCREAMING_SNAKE_CASE   ,			unittest.TestCase     ):




			def 	lowerCamelCase_       (    self				:		int    ):







						"""simple docstring"""







						UpperCamelCase								=       FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(
						    """hf-internal-testing/tiny-random-clip"""							,							"""hf-internal-testing/tiny-bert"""							,							vision_from_pt=__snake_case							,							text_from_pt=__snake_case							,							)
						UpperCamelCase								=       1_3
						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[int]							,							__magic_name__				:		int							,							__magic_name__				:		List[str]    ):







						"""simple docstring"""







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





			def 	lowerCamelCase_       (    self				:		List[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  UpperCAmelCase       (							unittest.TestCase     ):




			@slow
			def 	lowerCamelCase_       (    self				:		Tuple    ):







						"""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=__snake_case							,							padding=__snake_case							,							return_tensors="""np"""    )

						UpperCamelCase								=       model(**__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.2_284_727, 0.3_104_122]]    )

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











import shutil
import tempfile
import unittest

from transformers import ClapFeatureExtractor, ClapProcessor, RobertaTokenizer, RobertaTokenizerFast
from transformers.testing_utils import require_sentencepiece, require_torchaudio

from .test_feature_extraction_clap import floats_list





@require_torchaudio
@require_sentencepiece
class UpperCAmelCase_				( unittest.TestCase ):


  def 						__UpperCAmelCase     (							self ):



       """simple docstring"""





       A_          =       '''laion/clap-htsat-unfused'''
       A_          =       tempfile.mkdtemp()


  def 						__UpperCAmelCase     (							self ,**__snake_case ):



       """simple docstring"""





       return RobertaTokenizer.from_pretrained(self.checkpoint ,**__snake_case )


  def 						__UpperCAmelCase     (							self ,**__snake_case ):



       """simple docstring"""





       return ClapFeatureExtractor.from_pretrained(self.checkpoint ,**__snake_case )


  def 						__UpperCAmelCase     (							self ):



       """simple docstring"""





       shutil.rmtree(self.tmpdirname )


  def 						__UpperCAmelCase     (							self ):



       """simple docstring"""





       A_          =       self.get_tokenizer()
       A_          =       self.get_feature_extractor()

       A_          =       ClapProcessor(tokenizer=__snake_case ,feature_extractor=__snake_case )

       processor.save_pretrained(self.tmpdirname )
       A_          =       ClapProcessor.from_pretrained(self.tmpdirname )

       self.assertEqual(processor.tokenizer.get_vocab() ,tokenizer.get_vocab() )
       self.assertIsInstance(processor.tokenizer ,__snake_case )

       self.assertEqual(processor.feature_extractor.to_json_string() ,feature_extractor.to_json_string() )
       self.assertIsInstance(processor.feature_extractor ,__snake_case )


  def 						__UpperCAmelCase     (							self ):



       """simple docstring"""





       A_          =       ClapProcessor(tokenizer=self.get_tokenizer() ,feature_extractor=self.get_feature_extractor() )
       processor.save_pretrained(self.tmpdirname )

       A_          =       self.get_tokenizer(bos_token='''(BOS)''' ,eos_token='''(EOS)''' )
       A_          =       self.get_feature_extractor(do_normalize=__snake_case ,padding_value=1.0 )

       A_          =       ClapProcessor.from_pretrained(
           self.tmpdirname ,bos_token='''(BOS)''' ,eos_token='''(EOS)''' ,do_normalize=__snake_case ,padding_value=1.0 )

       self.assertEqual(processor.tokenizer.get_vocab() ,tokenizer_add_kwargs.get_vocab() )
       self.assertIsInstance(processor.tokenizer ,__snake_case )

       self.assertEqual(processor.feature_extractor.to_json_string() ,feature_extractor_add_kwargs.to_json_string() )
       self.assertIsInstance(processor.feature_extractor ,__snake_case )


  def 						__UpperCAmelCase     (							self ):



       """simple docstring"""





       A_          =       self.get_feature_extractor()
       A_          =       self.get_tokenizer()

       A_          =       ClapProcessor(tokenizer=__snake_case ,feature_extractor=__snake_case )

       A_          =       floats_list((3, 1_0_0_0) )

       A_          =       feature_extractor(__snake_case ,return_tensors='''np''' )
       A_          =       processor(audios=__snake_case ,return_tensors='''np''' )

       for key in input_feat_extract.keys():
            self.assertAlmostEqual(input_feat_extract[key].sum() ,input_processor[key].sum() ,delta=1E-2 )


  def 						__UpperCAmelCase     (							self ):



       """simple docstring"""





       A_          =       self.get_feature_extractor()
       A_          =       self.get_tokenizer()

       A_          =       ClapProcessor(tokenizer=__snake_case ,feature_extractor=__snake_case )

       A_          =       '''This is a test string'''

       A_          =       processor(text=__snake_case )

       A_          =       tokenizer(__snake_case )

       for key in encoded_tok.keys():
            self.assertListEqual(encoded_tok[key] ,encoded_processor[key] )


  def 						__UpperCAmelCase     (							self ):



       """simple docstring"""





       A_          =       self.get_feature_extractor()
       A_          =       self.get_tokenizer()

       A_          =       ClapProcessor(tokenizer=__snake_case ,feature_extractor=__snake_case )

       A_          =       [[1, 4, 5, 8, 1, 0, 8], [3, 4, 3, 1, 1, 8, 9]]

       A_          =       processor.batch_decode(__snake_case )
       A_          =       tokenizer.batch_decode(__snake_case )

       self.assertListEqual(__snake_case ,__snake_case )




  def 						__UpperCAmelCase     (							self ):



       """simple docstring"""





       A_          =       self.get_feature_extractor()
       A_          =       self.get_tokenizer()

       A_          =       ClapProcessor(tokenizer=__snake_case ,feature_extractor=__snake_case )

       self.assertListEqual(
           processor.model_input_names[2:] ,feature_extractor.model_input_names ,msg='''`processor` and `feature_extractor` model input names do not match''' ,)



