


'''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    lowerCamelCase__   ( A     :    Optional[Any]  ):




						'''simple docstring'''






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






@require_flax
class  UpperCamelCase__:




				def   a__(	self    :				Any   ,					lowerCAmelCase    :				Tuple   ,					lowerCAmelCase    :				Optional[int]   )->  Optional[int]:

										"""simple docstring"""




										pass




				def   a__(	self    :				Dict   )->  Any:

										"""simple docstring"""




										pass




				def   a__(	self    :				List[Any]   )->  Any:

										"""simple docstring"""




										pass




				def   a__(	self    :				Union[str, Any]   ,					lowerCAmelCase    :				np.ndarray   ,					lowerCAmelCase    :				np.ndarray   ,					lowerCAmelCase    :				float   )->  List[Any]:

										"""simple docstring"""




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




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

										"""simple docstring"""




										UpperCAmelCase         =			VisionTextDualEncoderConfig.from_vision_text_configs(lowerCAmelCase   ,					lowerCAmelCase   )

										UpperCAmelCase         =			FlaxVisionTextDualEncoderModel(lowerCAmelCase   )

										UpperCAmelCase         =			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   a__(	self    :				int   ,					lowerCAmelCase    :				int   ,					lowerCAmelCase    :				List[Any]   ,					lowerCAmelCase    :				int   ,					lowerCAmelCase    :				int   ,					lowerCAmelCase    :				Optional[Any]=None   ,					**lowerCAmelCase    :				Optional[int]   )->  Dict:

										"""simple docstring"""




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

										UpperCAmelCase         =			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   a__(	self    :				List[str]   ,					lowerCAmelCase    :				List[Any]   ,					lowerCAmelCase    :				Optional[int]   ,					lowerCAmelCase    :				Dict   ,					lowerCAmelCase    :				Optional[int]   ,					lowerCAmelCase    :				Any=None   ,					**lowerCAmelCase    :				List[Any]   )->  Optional[int]:

										"""simple docstring"""




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

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

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

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




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

										"""simple docstring"""




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

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

										UpperCAmelCase         =			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)
										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(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   a__(	self    :				Dict   ,					lowerCAmelCase    :				List[str]   ,					lowerCAmelCase    :				int   ,					lowerCAmelCase    :				List[Any]   )->  Any:

										"""simple docstring"""




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

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

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

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

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

										pt_model_loaded.to(lowerCAmelCase   )
										pt_model_loaded.eval()

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

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




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

										"""simple docstring"""




										UpperCAmelCase         =			VisionTextDualEncoderConfig.from_vision_text_configs(lowerCAmelCase   ,					lowerCAmelCase   )

										UpperCAmelCase         =			VisionTextDualEncoderModel(lowerCAmelCase   )
										UpperCAmelCase         =			FlaxVisionTextDualEncoderModel(lowerCAmelCase   )

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

										self.check_pt_flax_equivalence(lowerCAmelCase   ,					lowerCAmelCase   ,					lowerCAmelCase   )




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

										"""simple docstring"""




										UpperCAmelCase         =			VisionTextDualEncoderConfig.from_vision_text_configs(lowerCAmelCase   ,					lowerCAmelCase   )

										UpperCAmelCase         =			VisionTextDualEncoderModel(lowerCAmelCase   )
										UpperCAmelCase         =			FlaxVisionTextDualEncoderModel(lowerCAmelCase   )

										UpperCAmelCase         =			load_flax_weights_in_pytorch_model(lowerCAmelCase   ,					fx_model.params   )

										self.check_pt_flax_equivalence(lowerCAmelCase   ,					lowerCAmelCase   ,					lowerCAmelCase   )




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

										"""simple docstring"""




										UpperCAmelCase         =			self.prepare_config_and_inputs()
										self.check_model_from_pretrained_configs(**lowerCAmelCase   )




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

										"""simple docstring"""




										UpperCAmelCase         =			self.prepare_config_and_inputs()
										self.check_vision_text_dual_encoder_from_pretrained(**lowerCAmelCase   )




				def   a__(	self    :				Tuple   )->  Tuple:

										"""simple docstring"""




										UpperCAmelCase         =			self.prepare_config_and_inputs()
										self.check_save_load(**lowerCAmelCase   )




				def   a__(	self    :				List[Any]   )->  int:

										"""simple docstring"""




										UpperCAmelCase         =			self.prepare_config_and_inputs()
										self.check_vision_text_output_attention(**lowerCAmelCase   )




				@is_pt_flax_cross_test
				def   a__(	self    :				int   )->  List[str]:

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





				@slow
				def   a__(	self    :				Union[str, Any]   )->  Union[str, Any]:

										"""simple docstring"""




										UpperCAmelCase       ,       UpperCAmelCase         =			self.get_pretrained_model_and_inputs()

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

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

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







@require_flax
class  UpperCamelCase__(     lowerCAmelCase		,       unittest.TestCase ):




				def   a__(	self    :				str   )->  Dict:

										"""simple docstring"""




										UpperCAmelCase         =			FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(
										    '''hf-internal-testing/tiny-random-vit'''   ,					'''hf-internal-testing/tiny-bert'''   ,					vision_from_pt=lowerCAmelCase   ,					text_from_pt=lowerCAmelCase   ,					)
										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   a__(	self    :				Optional[int]   ,					lowerCAmelCase    :				Optional[Any]   ,					lowerCAmelCase    :				List[str]   )->  Union[str, Any]:

										"""simple docstring"""




										UpperCAmelCase         =			FlaxViTModel(lowerCAmelCase   )
										UpperCAmelCase         =			FlaxBertModel(lowerCAmelCase   )
										return vision_model, text_model





				def   a__(	self    :				Union[str, Any]   )->  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  UpperCamelCase__(     lowerCAmelCase		,       unittest.TestCase ):




				def   a__(	self    :				str   )->  Tuple:

										"""simple docstring"""




										UpperCAmelCase         =			FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(
										    '''hf-internal-testing/tiny-random-clip'''   ,					'''hf-internal-testing/tiny-bert'''   ,					vision_from_pt=lowerCAmelCase   ,					text_from_pt=lowerCAmelCase   ,					)
										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   a__(	self    :				Tuple   ,					lowerCAmelCase    :				Tuple   ,					lowerCAmelCase    :				int   )->  Tuple:

										"""simple docstring"""




										UpperCAmelCase         =			FlaxCLIPVisionModel(lowerCAmelCase   )
										UpperCAmelCase         =			FlaxBertModel(lowerCAmelCase   )
										return vision_model, text_model





				def   a__(	self    :				Dict   )->  Dict:

										"""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   a__(	self    :				str   )->  Dict:

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

										UpperCAmelCase         =			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])   ,					)

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

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









'''simple docstring'''



from argparse import ArgumentParser

from . import BaseTransformersCLICommand


def    lowerCamelCase__   ( A     :    List[Any]  ):




						'''simple docstring'''






						return DownloadCommand(args.model   ,     args.cache_dir   ,     args.force   ,     args.trust_remote_code  )






class  UpperCamelCase__(     lowerCAmelCase ):




				@staticmethod
				def   a__(	lowerCAmelCase    :				ArgumentParser   )->  Optional[Any]:

										"""simple docstring"""




										UpperCAmelCase         =			parser.add_parser('''download'''   )
										download_parser.add_argument(
										    '''--cache-dir'''   ,					type=lowerCAmelCase   ,					default=lowerCAmelCase   ,					help='''Path to location to store the models'''   )
										download_parser.add_argument(
										    '''--force'''   ,					action='''store_true'''   ,					help='''Force the model to be download even if already in cache-dir'''   )
										download_parser.add_argument(
										    '''--trust-remote-code'''   ,					action='''store_true'''   ,					help='''Whether or not to allow for custom models defined on the Hub in their own modeling files. Use only if you\'ve reviewed the code as it will execute on your local machine'''   ,					)
										download_parser.add_argument('''model'''   ,					type=lowerCAmelCase   ,					help='''Name of the model to download'''   )
										download_parser.set_defaults(func=lowerCAmelCase   )




				def __init__(	self    :				Dict   ,					lowerCAmelCase    :				str   ,					lowerCAmelCase    :				str   ,					lowerCAmelCase    :				bool   ,					lowerCAmelCase    :				bool   )->  Any:

										"""simple docstring"""




										UpperCAmelCase         =			model
										UpperCAmelCase         =			cache
										UpperCAmelCase         =			force
										UpperCAmelCase         =			trust_remote_code





				def   a__(	self    :				int   )->  Optional[Any]:

										"""simple docstring"""




										from ..models.auto import AutoModel, AutoTokenizer

										AutoModel.from_pretrained(
										    self._model   ,					cache_dir=self._cache   ,					force_download=self._force   ,					trust_remote_code=self._trust_remote_code   )
										AutoTokenizer.from_pretrained(
										    self._model   ,					cache_dir=self._cache   ,					force_download=self._force   ,					trust_remote_code=self._trust_remote_code   )





