



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__	(					lowerCamelCase_	:	Any):





				'''simple docstring'''

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



@require_flax
class  lowerCamelCase__    :




		'''simple docstring'''



		def 	lowerCAmelCase__   (self	,__lowerCamelCase	,__lowerCamelCase				)  ->    str:

						"""simple docstring"""






						pass



		def 	lowerCAmelCase__   (self				)  ->    str:

						"""simple docstring"""






						pass



		def 	lowerCAmelCase__   (self				)  ->    List[str]:

						"""simple docstring"""






						pass



		def 	lowerCAmelCase__   (self	,__lowerCamelCase	,__lowerCamelCase	,__lowerCamelCase				)  ->    Tuple:

						"""simple docstring"""






						lowerCAmelCase__  :       Optional[int]				=  np.abs((a - b)				).max()
						self.assertLessEqual(__lowerCamelCase	,__lowerCamelCase	,f"""Difference between torch and flax is {diff} (>= {tol})."""				)



		def 	lowerCAmelCase__   (self	,__lowerCamelCase	,__lowerCamelCase	,__lowerCamelCase	,__lowerCamelCase	,__lowerCamelCase=None	,**__lowerCamelCase				)  ->    Dict:

						"""simple docstring"""






						lowerCAmelCase__  :       Optional[int]				=  VisionTextDualEncoderConfig.from_vision_text_configs(__lowerCamelCase	,__lowerCamelCase				)

						lowerCAmelCase__  :       List[str]				=  FlaxVisionTextDualEncoderModel(__lowerCamelCase				)

						lowerCAmelCase__  :       Dict				=  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 	lowerCAmelCase__   (self	,__lowerCamelCase	,__lowerCamelCase	,__lowerCamelCase	,__lowerCamelCase	,__lowerCamelCase=None	,**__lowerCamelCase				)  ->    List[Any]:

						"""simple docstring"""






						lowerCAmelCase__  :       Any				=  self.get_vision_text_model(__lowerCamelCase	,__lowerCamelCase				)
						lowerCAmelCase__  :       Optional[int]				=  {"vision_model": vision_model, "text_model": text_model}
						lowerCAmelCase__  :       Optional[int]				=  FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(**__lowerCamelCase				)

						lowerCAmelCase__  :       str				=  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 	lowerCAmelCase__   (self	,__lowerCamelCase	,__lowerCamelCase	,__lowerCamelCase	,__lowerCamelCase	,__lowerCamelCase=None	,**__lowerCamelCase				)  ->    Union[str, Any]:

						"""simple docstring"""






						lowerCAmelCase__  :       List[Any]				=  self.get_vision_text_model(__lowerCamelCase	,__lowerCamelCase				)
						lowerCAmelCase__  :       Optional[int]				=  {"vision_model": vision_model, "text_model": text_model}
						lowerCAmelCase__  :       Optional[int]				=  FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(**__lowerCamelCase				)

						lowerCAmelCase__  :       Any				=  model(input_ids=__lowerCamelCase	,pixel_values=__lowerCamelCase	,attention_mask=__lowerCamelCase				)
						lowerCAmelCase__  :       Any				=  output[0]

						with tempfile.TemporaryDirectory() as tmpdirname:
										model.save_pretrained(__lowerCamelCase				)
										lowerCAmelCase__  :       Any				=  FlaxVisionTextDualEncoderModel.from_pretrained(__lowerCamelCase				)

										lowerCAmelCase__  :       Any				=  model(input_ids=__lowerCamelCase	,pixel_values=__lowerCamelCase	,attention_mask=__lowerCamelCase				)
										lowerCAmelCase__  :       List[str]				=  after_output[0]
										lowerCAmelCase__  :       Any				=  np.amax(np.abs(out_a - out_a				)				)
										self.assertLessEqual(__lowerCamelCase	,1e-3				)



		def 	lowerCAmelCase__   (self	,__lowerCamelCase	,__lowerCamelCase	,__lowerCamelCase	,__lowerCamelCase	,__lowerCamelCase=None	,**__lowerCamelCase				)  ->    int:

						"""simple docstring"""






						lowerCAmelCase__  :       List[Any]				=  self.get_vision_text_model(__lowerCamelCase	,__lowerCamelCase				)
						lowerCAmelCase__  :       Union[str, Any]				=  {"vision_model": vision_model, "text_model": text_model}
						lowerCAmelCase__  :       Optional[Any]				=  FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(**__lowerCamelCase				)

						lowerCAmelCase__  :       int				=  model(
						    input_ids=__lowerCamelCase	,pixel_values=__lowerCamelCase	,attention_mask=__lowerCamelCase	,output_attentions=__lowerCamelCase				)

						lowerCAmelCase__  :       Tuple				=  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)
						lowerCAmelCase__  :       Optional[Any]				=  to_atuple(vision_model.config.image_size				)
						lowerCAmelCase__  :       int				=  to_atuple(vision_model.config.patch_size				)
						lowerCAmelCase__  :       List[Any]				=  (image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0])
						lowerCAmelCase__  :       List[str]				=  num_patches + 1
						self.assertEqual(vision_attentions[0].shape[-3:]	,(vision_config.num_attention_heads, seq_len, seq_len)				)

						lowerCAmelCase__  :       Optional[Any]				=  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 	lowerCAmelCase__   (self	,__lowerCamelCase	,__lowerCamelCase	,__lowerCamelCase				)  ->    Union[str, Any]:

						"""simple docstring"""






						pt_model.to(__lowerCamelCase				)
						pt_model.eval()

						# prepare inputs
						lowerCAmelCase__  :       Optional[int]				=  inputs_dict
						lowerCAmelCase__  :       List[Any]				=  {k: torch.tensor(v.tolist()				) for k, v in flax_inputs.items()}

						with torch.no_grad():
										lowerCAmelCase__  :       str				=  pt_model(**__lowerCamelCase				).to_tuple()

						lowerCAmelCase__  :       Optional[Any]				=  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				)
										lowerCAmelCase__  :       Optional[int]				=  FlaxVisionTextDualEncoderModel.from_pretrained(__lowerCamelCase	,from_pt=__lowerCamelCase				)

						lowerCAmelCase__  :       str				=  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				)
										lowerCAmelCase__  :       List[Any]				=  VisionTextDualEncoderModel.from_pretrained(__lowerCamelCase	,from_flax=__lowerCamelCase				)

						pt_model_loaded.to(__lowerCamelCase				)
						pt_model_loaded.eval()

						with torch.no_grad():
										lowerCAmelCase__  :       Tuple				=  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 	lowerCAmelCase__   (self	,__lowerCamelCase	,__lowerCamelCase	,__lowerCamelCase				)  ->    Optional[int]:

						"""simple docstring"""






						lowerCAmelCase__  :       List[Any]				=  VisionTextDualEncoderConfig.from_vision_text_configs(__lowerCamelCase	,__lowerCamelCase				)

						lowerCAmelCase__  :       str				=  VisionTextDualEncoderModel(__lowerCamelCase				)
						lowerCAmelCase__  :       int				=  FlaxVisionTextDualEncoderModel(__lowerCamelCase				)

						lowerCAmelCase__  :       Optional[Any]				=  convert_pytorch_state_dict_to_flax(pt_model.state_dict()	,__lowerCamelCase				)
						lowerCAmelCase__  :       int				=  fx_state

						self.check_pt_flax_equivalence(__lowerCamelCase	,__lowerCamelCase	,__lowerCamelCase				)



		def 	lowerCAmelCase__   (self	,__lowerCamelCase	,__lowerCamelCase	,__lowerCamelCase				)  ->    Optional[int]:

						"""simple docstring"""






						lowerCAmelCase__  :       Union[str, Any]				=  VisionTextDualEncoderConfig.from_vision_text_configs(__lowerCamelCase	,__lowerCamelCase				)

						lowerCAmelCase__  :       Tuple				=  VisionTextDualEncoderModel(__lowerCamelCase				)
						lowerCAmelCase__  :       Union[str, Any]				=  FlaxVisionTextDualEncoderModel(__lowerCamelCase				)

						lowerCAmelCase__  :       List[Any]				=  load_flax_weights_in_pytorch_model(__lowerCamelCase	,fx_model.params				)

						self.check_pt_flax_equivalence(__lowerCamelCase	,__lowerCamelCase	,__lowerCamelCase				)



		def 	lowerCAmelCase__   (self				)  ->    int:

						"""simple docstring"""






						lowerCAmelCase__  :       Optional[Any]				=  self.prepare_config_and_inputs()
						self.check_model_from_pretrained_configs(**__lowerCamelCase				)



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

						"""simple docstring"""






						lowerCAmelCase__  :       Union[str, Any]				=  self.prepare_config_and_inputs()
						self.check_vision_text_dual_encoder_from_pretrained(**__lowerCamelCase				)



		def 	lowerCAmelCase__   (self				)  ->    List[str]:

						"""simple docstring"""






						lowerCAmelCase__  :       Optional[int]				=  self.prepare_config_and_inputs()
						self.check_save_load(**__lowerCamelCase				)



		def 	lowerCAmelCase__   (self				)  ->    Optional[int]:

						"""simple docstring"""






						lowerCAmelCase__  :       List[str]				=  self.prepare_config_and_inputs()
						self.check_vision_text_output_attention(**__lowerCamelCase				)



		@is_pt_flax_cross_test
		def 	lowerCAmelCase__   (self				)  ->    int:

						"""simple docstring"""






						lowerCAmelCase__  :       List[Any]				=  self.prepare_config_and_inputs()
						lowerCAmelCase__  :       Optional[int]				=  config_inputs_dict.pop('''vision_config'''				)
						lowerCAmelCase__  :       Dict				=  config_inputs_dict.pop('''text_config'''				)

						lowerCAmelCase__  :       Optional[Any]				=  config_inputs_dict

						self.check_equivalence_pt_to_flax(__lowerCamelCase	,__lowerCamelCase	,__lowerCamelCase				)
						self.check_equivalence_flax_to_pt(__lowerCamelCase	,__lowerCamelCase	,__lowerCamelCase				)



		@slow
		def 	lowerCAmelCase__   (self				)  ->    Optional[int]:

						"""simple docstring"""






						lowerCAmelCase__  :       Dict				=  self.get_pretrained_model_and_inputs()

						lowerCAmelCase__  :       List[Any]				=  model_a(**__lowerCamelCase				)
						lowerCAmelCase__  :       Optional[int]				=  outputs[0]

						with tempfile.TemporaryDirectory() as tmp_dirname:
										model_a.save_pretrained(__lowerCamelCase				)
										lowerCAmelCase__  :       Optional[Any]				=  FlaxVisionTextDualEncoderModel.from_pretrained(__lowerCamelCase				)

										lowerCAmelCase__  :       List[str]				=  model_a(**__lowerCamelCase				)
										lowerCAmelCase__  :       Tuple				=  after_outputs[0]
										lowerCAmelCase__  :       Optional[int]				=  np.amax(np.abs(out_a - out_a				)				)
										self.assertLessEqual(__lowerCamelCase	,1e-5				)



@require_flax
class  lowerCamelCase__    (							_A   ,       unittest.TestCase):




		'''simple docstring'''



		def 	lowerCAmelCase__   (self				)  ->    List[str]:

						"""simple docstring"""






						lowerCAmelCase__  :       Any				=  FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(
						    '''hf-internal-testing/tiny-random-vit'''	,'''hf-internal-testing/tiny-bert'''	,vision_from_pt=__lowerCamelCase	,text_from_pt=__lowerCamelCase	,)
						lowerCAmelCase__  :       Union[str, Any]				=  13
						lowerCAmelCase__  :       List[Any]				=  floats_tensor(
						    [
						        batch_size,
						        model.config.vision_config.num_channels,
						        model.config.vision_config.image_size,
						        model.config.vision_config.image_size,
						    ]				)
						lowerCAmelCase__  :       Union[str, Any]				=  ids_tensor([batch_size, 4]	,model.config.text_config.vocab_size				)
						lowerCAmelCase__  :       Optional[int]				=  random_attention_mask([batch_size, 4]				)
						lowerCAmelCase__  :       Optional[Any]				=  {"pixel_values": pixel_values, "input_ids": input_ids, "attention_mask": attention_mask}

						return model, inputs



		def 	lowerCAmelCase__   (self	,__lowerCamelCase	,__lowerCamelCase				)  ->    List[str]:

						"""simple docstring"""






						lowerCAmelCase__  :       Dict				=  FlaxViTModel(__lowerCamelCase				)
						lowerCAmelCase__  :       Tuple				=  FlaxBertModel(__lowerCamelCase				)
						return vision_model, text_model



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

						"""simple docstring"""






						lowerCAmelCase__  :       Optional[int]				=  FlaxViTModelTester(self				)
						lowerCAmelCase__  :       Optional[Any]				=  FlaxBertModelTester(self				)
						lowerCAmelCase__  :       Dict				=  vit_model_tester.prepare_config_and_inputs()
						lowerCAmelCase__  :       int				=  bert_model_tester.prepare_config_and_inputs()

						lowerCAmelCase__  :       int				=  vision_config_and_inputs

						lowerCAmelCase__  :       int				=  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__    (							_A   ,       unittest.TestCase):




		'''simple docstring'''



		def 	lowerCAmelCase__   (self				)  ->    Any:

						"""simple docstring"""






						lowerCAmelCase__  :       int				=  FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(
						    '''hf-internal-testing/tiny-random-clip'''	,'''hf-internal-testing/tiny-bert'''	,vision_from_pt=__lowerCamelCase	,text_from_pt=__lowerCamelCase	,)
						lowerCAmelCase__  :       List[Any]				=  13
						lowerCAmelCase__  :       Dict				=  floats_tensor(
						    [
						        batch_size,
						        model.config.vision_config.num_channels,
						        model.config.vision_config.image_size,
						        model.config.vision_config.image_size,
						    ]				)
						lowerCAmelCase__  :       Union[str, Any]				=  ids_tensor([batch_size, 4]	,model.config.text_config.vocab_size				)
						lowerCAmelCase__  :       Optional[Any]				=  random_attention_mask([batch_size, 4]				)
						lowerCAmelCase__  :       Optional[int]				=  {"pixel_values": pixel_values, "input_ids": input_ids, "attention_mask": attention_mask}

						return model, inputs



		def 	lowerCAmelCase__   (self	,__lowerCamelCase	,__lowerCamelCase				)  ->    int:

						"""simple docstring"""






						lowerCAmelCase__  :       Any				=  FlaxCLIPVisionModel(__lowerCamelCase				)
						lowerCAmelCase__  :       Union[str, Any]				=  FlaxBertModel(__lowerCamelCase				)
						return vision_model, text_model



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

						"""simple docstring"""






						lowerCAmelCase__  :       Optional[Any]				=  FlaxCLIPVisionModelTester(self				)
						lowerCAmelCase__  :       Optional[int]				=  FlaxBertModelTester(self				)
						lowerCAmelCase__  :       Tuple				=  clip_model_tester.prepare_config_and_inputs()
						lowerCAmelCase__  :       int				=  bert_model_tester.prepare_config_and_inputs()

						lowerCAmelCase__  :       Optional[int]				=  vision_config_and_inputs

						lowerCAmelCase__  :       Optional[Any]				=  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				)  ->    Dict:

						"""simple docstring"""






						lowerCAmelCase__  :       Dict				=  FlaxVisionTextDualEncoderModel.from_pretrained('''clip-italian/clip-italian'''	,logit_scale_init_value=1.0				)
						lowerCAmelCase__  :       Any				=  VisionTextDualEncoderProcessor.from_pretrained('''clip-italian/clip-italian'''				)

						lowerCAmelCase__  :       Optional[int]				=  Image.open('''./tests/fixtures/tests_samples/COCO/000000039769.png'''				)
						lowerCAmelCase__  :       Tuple				=  processor(
						    text=['''una foto di un gatto''', '''una foto di un cane''']	,images=__lowerCamelCase	,padding=__lowerCamelCase	,return_tensors='''np'''				)

						lowerCAmelCase__  :       List[Any]				=  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])	,)

						lowerCAmelCase__  :       Dict				=  np.array([[1.228_4727, 0.310_4122]]				)

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






import io
import json
import unittest

from parameterized import parameterized

from transformers import FSMTForConditionalGeneration, FSMTTokenizer
from transformers.testing_utils import get_tests_dir, require_torch, slow, torch_device
from utils import calculate_bleu


__snake_case       :	Optional[int]            =get_tests_dir() + '/test_data/fsmt/fsmt_val_data.json'
with io.open(filename, 'r', encoding='utf-8') as f:
							__snake_case       :	Optional[Any]            =json.load(f)



@require_torch
class  lowerCamelCase__    (							unittest.TestCase):




		'''simple docstring'''



		def 	lowerCAmelCase__   (self	,__lowerCamelCase				)  ->    Dict:

						"""simple docstring"""






						return FSMTTokenizer.from_pretrained(__lowerCamelCase				)



		def 	lowerCAmelCase__   (self	,__lowerCamelCase				)  ->    Optional[int]:

						"""simple docstring"""






						lowerCAmelCase__  :       Optional[Any]				=  FSMTForConditionalGeneration.from_pretrained(__lowerCamelCase				).to(__lowerCamelCase				)
						if torch_device == "cuda":
										model.half()
						return model



		@parameterized.expand(
		    [
		        ['''en-ru''', 26.0],
		        ['''ru-en''', 22.0],
		        ['''en-de''', 22.0],
		        ['''de-en''', 29.0],
		    ]				)
		@slow
		def 	lowerCAmelCase__   (self	,__lowerCamelCase	,__lowerCamelCase				)  ->    Any:

						"""simple docstring"""






						lowerCAmelCase__  :       List[str]				=  f"""facebook/wmt19-{pair}"""
						lowerCAmelCase__  :       int				=  self.get_tokenizer(__lowerCamelCase				)
						lowerCAmelCase__  :       Optional[int]				=  self.get_model(__lowerCamelCase				)

						lowerCAmelCase__  :       List[Any]				=  bleu_data[pair]['''src''']
						lowerCAmelCase__  :       List[Any]				=  bleu_data[pair]['''tgt''']

						lowerCAmelCase__  :       List[str]				=  tokenizer(__lowerCamelCase	,return_tensors='''pt'''	,truncation=__lowerCamelCase	,padding='''longest'''				).to(__lowerCamelCase				)
						lowerCAmelCase__  :       List[str]				=  model.generate(
						    input_ids=batch.input_ids	,num_beams=8	,)
						lowerCAmelCase__  :       Tuple				=  tokenizer.batch_decode(
						    __lowerCamelCase	,skip_special_tokens=__lowerCamelCase	,clean_up_tokenization_spaces=__lowerCamelCase				)
						lowerCAmelCase__  :       Union[str, Any]				=  calculate_bleu(__lowerCamelCase	,__lowerCamelCase				)
						print(__lowerCamelCase				)
						self.assertGreaterEqual(scores['''bleu''']	,__lowerCamelCase				)

