





from __future__ import annotations

import collections
import tempfile
import unittest

import numpy as np

from transformers.testing_utils import require_tf, require_vision, slow
from transformers.utils import is_tf_available, is_vision_available

from ...test_modeling_tf_common import floats_tensor, ids_tensor, random_attention_mask
from ..bert.test_modeling_tf_bert import TFBertModelTester
from ..clip.test_modeling_tf_clip import TFCLIPVisionModelTester
from ..deit.test_modeling_tf_deit import TFDeiTModelTester
from ..roberta.test_modeling_tf_roberta import TFRobertaModelTester
from ..vit.test_modeling_tf_vit import TFViTModelTester


if is_tf_available():
		from transformers import (
		    TFBertModel,
		    TFCLIPVisionModel,
		    TFDeiTModel,
		    TFRobertaModel,
		    TFVisionTextDualEncoderModel,
		    TFViTModel,
		    VisionTextDualEncoderConfig,
		)

if is_vision_available():
		from PIL import Image

		from transformers import VisionTextDualEncoderProcessor

def    A__(		__lowerCAmelCase			):
				if isinstance(__lowerCAmelCase						,      collections.abc.Iterable			):
								return x
				return (x, x)


@require_tf
class       lowercase					:



						"""simple docstring"""

						def 		__UpperCAmelCase							(     self      :		Optional[int]			,					lowerCamelCase_      :		Optional[int]			,					lowerCamelCase_      :		Tuple       ):





										'''simple docstring'''







										pass







						def 		__UpperCAmelCase							(     self      :		Dict       ):





										'''simple docstring'''







										pass







						def 		__UpperCAmelCase							(     self      :		Optional[Any]       ):





										'''simple docstring'''







										pass







						def 		__UpperCAmelCase							(     self      :		Tuple			,					lowerCamelCase_      :		int			,					lowerCamelCase_      :		Tuple			,					lowerCamelCase_      :		Dict			,					lowerCamelCase_      :		Optional[Any]			,					lowerCamelCase_      :		Dict=None			,					**lowerCamelCase_      :		List[Any]       ):





										'''simple docstring'''







										_snake_case			:      Optional[int]	    = VisionTextDualEncoderConfig.from_vision_text_configs(snake_case_			,					snake_case_       )

										_snake_case			:      Dict	    = TFVisionTextDualEncoderModel(snake_case_       )

										_snake_case			:      Optional[int]	    = 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 		__UpperCAmelCase							(     self      :		Union[str, Any]			,					lowerCamelCase_      :		List[Any]			,					lowerCamelCase_      :		Any			,					lowerCamelCase_      :		Optional[int]			,					lowerCamelCase_      :		Dict			,					lowerCamelCase_      :		List[str]=None			,					**lowerCamelCase_      :		Dict       ):





										'''simple docstring'''







										_snake_case			:      Optional[int]	    = self.get_vision_text_model(snake_case_			,					snake_case_       )
										_snake_case			:      int	    = TFVisionTextDualEncoderModel(vision_model=snake_case_			,					text_model=snake_case_       )

										_snake_case			:      Tuple	    = 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 		__UpperCAmelCase							(     self      :		Optional[int]			,					lowerCamelCase_      :		Union[str, Any]			,					lowerCamelCase_      :		List[str]			,					lowerCamelCase_      :		Optional[Any]			,					lowerCamelCase_      :		str			,					lowerCamelCase_      :		str=None			,					**lowerCamelCase_      :		Any       ):





										'''simple docstring'''







										_snake_case			:      List[str]	    = self.get_vision_text_model(snake_case_			,					snake_case_       )
										_snake_case			:      Dict	    = {"""vision_model""": vision_model, """text_model""": text_model}
										_snake_case			:      List[str]	    = TFVisionTextDualEncoderModel.from_vision_text_pretrained(**snake_case_       )

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





										'''simple docstring'''







										_snake_case			:      List[Any]	    = self.get_vision_text_model(snake_case_			,					snake_case_       )
										_snake_case			:      Tuple	    = TFVisionTextDualEncoderModel(vision_model=snake_case_			,					text_model=snake_case_       )

										_snake_case			:      List[Any]	    = model(input_ids=snake_case_			,					pixel_values=snake_case_			,					attention_mask=snake_case_       )
										_snake_case			:      Optional[Any]	    = output[0].numpy()

										with tempfile.TemporaryDirectory() as tmpdirname:
														model.save_pretrained(snake_case_       )
														_snake_case			:      Dict	    = TFVisionTextDualEncoderModel.from_pretrained(snake_case_       )

														_snake_case			:      Dict	    = model(input_ids=snake_case_			,					pixel_values=snake_case_			,					attention_mask=snake_case_       )
														_snake_case			:      List[Any]	    = after_output[0].numpy()
														_snake_case			:      Any	    = np.amax(np.abs(out_a - out_a       )       )
														self.assertLessEqual(snake_case_			,					1e-5       )







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





										'''simple docstring'''







										_snake_case			:      Any	    = self.get_vision_text_model(snake_case_			,					snake_case_       )
										_snake_case			:      Any	    = TFVisionTextDualEncoderModel(vision_model=snake_case_			,					text_model=snake_case_       )

										_snake_case			:      Optional[int]	    = model(
										    input_ids=snake_case_			,					pixel_values=snake_case_			,					attention_mask=snake_case_			,					output_attentions=snake_case_       )

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

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





										'''simple docstring'''







										_snake_case			:      List[Any]	    = np.abs((a - b)       ).max()
										self.assertLessEqual(snake_case_			,					snake_case_			,					f'''Difference between torch and flax is {diff} (>= {tol}).'''       )







						def 		__UpperCAmelCase							(     self      :		int       ):





										'''simple docstring'''







										_snake_case			:      List[Any]	    = self.prepare_config_and_inputs()
										self.check_vision_text_dual_encoder_model(**snake_case_       )







						def 		__UpperCAmelCase							(     self      :		List[Any]       ):





										'''simple docstring'''







										_snake_case			:      Optional[int]	    = self.prepare_config_and_inputs()
										self.check_model_from_pretrained_configs(**snake_case_       )







						def 		__UpperCAmelCase							(     self      :		Dict       ):





										'''simple docstring'''







										_snake_case			:      Optional[Any]	    = self.prepare_config_and_inputs()
										self.check_vision_text_dual_encoder_from_pretrained(**snake_case_       )







						def 		__UpperCAmelCase							(     self      :		List[str]       ):





										'''simple docstring'''







										_snake_case			:      int	    = self.prepare_config_and_inputs()
										self.check_save_load(**snake_case_       )







						def 		__UpperCAmelCase							(     self      :		str       ):





										'''simple docstring'''







										_snake_case			:      str	    = self.prepare_config_and_inputs()
										self.check_vision_text_output_attention(**snake_case_       )







						@slow
						def 		__UpperCAmelCase							(     self      :		int       ):





										'''simple docstring'''







										_snake_case			:      int	    = self.get_pretrained_model_and_inputs()

										_snake_case			:      Union[str, Any]	    = model_a(**snake_case_       )
										_snake_case			:      Union[str, Any]	    = outputs[0].numpy()

										with tempfile.TemporaryDirectory() as tmp_dirname:
														model_a.save_pretrained(snake_case_       )
														_snake_case			:      Optional[int]	    = TFVisionTextDualEncoderModel.from_pretrained(snake_case_       )

														_snake_case			:      int	    = model_a(**snake_case_       )
														_snake_case			:      List[str]	    = after_outputs[0].numpy()
														_snake_case			:      Any	    = np.amax(np.abs(out_a - out_a       )       )
														self.assertLessEqual(snake_case_			,					1e-5       )



@require_tf
class       lowercase					(				__A		,      unittest.TestCase  ):



						"""simple docstring"""

						def 		__UpperCAmelCase							(     self      :		int       ):





										'''simple docstring'''







										_snake_case			:      str	    = TFVisionTextDualEncoderModel.from_vision_text_pretrained(
										    'hf-internal-testing/tiny-random-vit'			,					'hf-internal-testing/tiny-random-bert'       )
										_snake_case			:      str	    = 13
										_snake_case			:      str	    = floats_tensor(
										    [
										        batch_size,
										        model.vision_model.config.num_channels,
										        model.vision_model.config.image_size,
										        model.vision_model.config.image_size,
										    ]       )
										_snake_case			:      Any	    = ids_tensor([batch_size, 4]			,					model.text_model.config.vocab_size       )
										_snake_case			:      int	    = random_attention_mask([batch_size, 4]       )
										_snake_case			:      Union[str, Any]	    = {"""pixel_values""": pixel_values, """input_ids""": input_ids, """attention_mask""": attention_mask}

										return model, inputs







						def 		__UpperCAmelCase							(     self      :		Dict			,					lowerCamelCase_      :		Optional[Any]			,					lowerCamelCase_      :		Any       ):





										'''simple docstring'''







										_snake_case			:      Tuple	    = TFViTModel(snake_case_			,					name='vision_model'       )
										_snake_case			:      Dict	    = TFBertModel(snake_case_			,					name='text_model'       )
										return vision_model, text_model







						def 		__UpperCAmelCase							(     self      :		List[str]       ):





										'''simple docstring'''







										_snake_case			:      List[str]	    = TFViTModelTester(self       )
										_snake_case			:      Any	    = TFBertModelTester(self       )
										_snake_case			:      Dict	    = vit_model_tester.prepare_config_and_inputs()
										_snake_case			:      Dict	    = bert_model_tester.prepare_config_and_inputs()

										_snake_case			:      Tuple	    = vision_config_and_inputs

										(
										    _snake_case							
										)			:      List[str]	    = text_config_and_inputs

										return {
										    "text_config": text_config,
										    "vision_config": vision_config,
										    "pixel_values": pixel_values,
										    "attention_mask": input_mask,
										    "input_ids": input_ids,
										    "text_token_type_ids": token_type_ids,
										    "text_sequence_labels": sequence_labels,
										    "text_token_labels": token_labels,
										    "text_choice_labels": choice_labels,
										}



@require_tf
class       lowercase					(				__A		,      unittest.TestCase  ):



						"""simple docstring"""

						def 		__UpperCAmelCase							(     self      :		Any       ):





										'''simple docstring'''







										_snake_case			:      List[str]	    = TFVisionTextDualEncoderModel.from_vision_text_pretrained(
										    'Rocketknight1/tiny-random-deit-tf'			,					'hf-internal-testing/tiny-random-roberta'       )
										_snake_case			:      List[Any]	    = 13
										_snake_case			:      List[Any]	    = floats_tensor(
										    [
										        batch_size,
										        model.vision_model.config.num_channels,
										        model.vision_model.config.image_size,
										        model.vision_model.config.image_size,
										    ]       )
										_snake_case			:      Any	    = ids_tensor([batch_size, 4]			,					model.text_model.config.vocab_size       )
										_snake_case			:      Tuple	    = random_attention_mask([batch_size, 4]       )
										_snake_case			:      str	    = {"""pixel_values""": pixel_values, """input_ids""": input_ids, """attention_mask""": attention_mask}

										return model, inputs







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





										'''simple docstring'''







										_snake_case			:      List[str]	    = self.get_vision_text_model(snake_case_			,					snake_case_       )
										_snake_case			:      List[str]	    = TFVisionTextDualEncoderModel(vision_model=snake_case_			,					text_model=snake_case_       )

										_snake_case			:      Any	    = model(
										    input_ids=snake_case_			,					pixel_values=snake_case_			,					attention_mask=snake_case_			,					output_attentions=snake_case_       )

										_snake_case			:      Any	    = output.vision_model_output.attentions
										self.assertEqual(len(snake_case_       )			,					vision_config.num_hidden_layers       )

										# in DEiT, the seq_len equals the number of patches + 2 (we add 2 for the [CLS] and distillation tokens)
										_snake_case			:      List[str]	    = to_atuple(vision_model.config.image_size       )
										_snake_case			:      str	    = to_atuple(vision_model.config.patch_size       )
										_snake_case			:      Union[str, Any]	    = (image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0])
										_snake_case			:      int	    = num_patches + 2
										self.assertEqual(vision_attentions[0].shape[-3:]			,					(vision_config.num_attention_heads, seq_len, seq_len)       )

										_snake_case			:      List[str]	    = 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 		__UpperCAmelCase							(     self      :		Tuple			,					lowerCamelCase_      :		Optional[Any]			,					lowerCamelCase_      :		Any       ):





										'''simple docstring'''







										_snake_case			:      Dict	    = TFDeiTModel(snake_case_			,					name='vision_model'       )
										_snake_case			:      Any	    = TFRobertaModel(snake_case_			,					name='text_model'       )
										return vision_model, text_model







						def 		__UpperCAmelCase							(     self      :		Any       ):





										'''simple docstring'''







										_snake_case			:      Union[str, Any]	    = TFDeiTModelTester(self       )
										_snake_case			:      Optional[int]	    = TFRobertaModelTester(self       )
										_snake_case			:      Tuple	    = vit_model_tester.prepare_config_and_inputs()
										_snake_case			:      int	    = bert_model_tester.prepare_config_and_inputs()

										_snake_case			:      Union[str, Any]	    = vision_config_and_inputs

										(
										    _snake_case							
										)			:      List[str]	    = text_config_and_inputs

										return {
										    "text_config": text_config,
										    "vision_config": vision_config,
										    "pixel_values": pixel_values,
										    "attention_mask": input_mask,
										    "input_ids": input_ids,
										    "text_token_type_ids": token_type_ids,
										    "text_sequence_labels": sequence_labels,
										    "text_token_labels": token_labels,
										    "text_choice_labels": choice_labels,
										}



@require_tf
class       lowercase					(				__A		,      unittest.TestCase  ):



						"""simple docstring"""

						def 		__UpperCAmelCase							(     self      :		int       ):





										'''simple docstring'''







										_snake_case			:      Dict	    = TFVisionTextDualEncoderModel.from_vision_text_pretrained(
										    'Rocketknight1/tiny-random-clip-tf'			,					'hf-internal-testing/tiny-random-bert'       )
										_snake_case			:      Dict	    = 13
										_snake_case			:      List[str]	    = floats_tensor(
										    [
										        batch_size,
										        model.vision_model.config.num_channels,
										        model.vision_model.config.image_size,
										        model.vision_model.config.image_size,
										    ]       )
										_snake_case			:      int	    = ids_tensor([batch_size, 4]			,					model.text_model.config.vocab_size       )
										_snake_case			:      List[str]	    = random_attention_mask([batch_size, 4]       )
										_snake_case			:      List[str]	    = {"""pixel_values""": pixel_values, """input_ids""": input_ids, """attention_mask""": attention_mask}

										return model, inputs







						def 		__UpperCAmelCase							(     self      :		Optional[int]			,					lowerCamelCase_      :		Optional[int]			,					lowerCamelCase_      :		Optional[Any]       ):





										'''simple docstring'''







										_snake_case			:      str	    = TFCLIPVisionModel(snake_case_			,					name='vision_model'       )
										_snake_case			:      Any	    = TFBertModel(snake_case_			,					name='text_model'       )
										return vision_model, text_model







						def 		__UpperCAmelCase							(     self      :		Tuple       ):





										'''simple docstring'''







										_snake_case			:      List[Any]	    = TFCLIPVisionModelTester(self       )
										_snake_case			:      int	    = TFBertModelTester(self       )
										_snake_case			:      Optional[Any]	    = clip_model_tester.prepare_config_and_inputs()
										_snake_case			:      Optional[int]	    = bert_model_tester.prepare_config_and_inputs()

										_snake_case			:      Union[str, Any]	    = vision_config_and_inputs

										(
										    _snake_case							
										)			:      List[str]	    = text_config_and_inputs

										return {
										    "text_config": text_config,
										    "vision_config": vision_config,
										    "pixel_values": pixel_values,
										    "attention_mask": input_mask,
										    "input_ids": input_ids,
										    "text_token_type_ids": token_type_ids,
										    "text_sequence_labels": sequence_labels,
										    "text_token_labels": token_labels,
										    "text_choice_labels": choice_labels,
										}



@require_vision
@require_tf
class       lowercase					(				unittest.TestCase  ):



						"""simple docstring"""

						@slow
						def 		__UpperCAmelCase							(     self      :		Optional[Any]       ):





										'''simple docstring'''







										_snake_case			:      Tuple	    = TFVisionTextDualEncoderModel.from_pretrained(
										    'clip-italian/clip-italian'			,					logit_scale_init_value=1.0			,					from_pt=snake_case_       )
										_snake_case			:      List[str]	    = VisionTextDualEncoderProcessor.from_pretrained('clip-italian/clip-italian'       )

										_snake_case			:      Optional[Any]	    = Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png'       )
										_snake_case			:      List[str]	    = processor(
										    text=['una foto di un gatto', 'una foto di un cane']			,					images=snake_case_			,					padding=snake_case_			,					return_tensors='np'       )

										_snake_case			:      Any	    = 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])			,					)

										_snake_case			:      Optional[Any]	    = np.array([[1.228_4727, 0.310_4122]]       )

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







"""simple docstring"""



from typing import List

from .keymap import KEYMAP, get_character

def  _A(		lowerCAmelCase							):

							def decorator(lowerCAmelCase							):
														A__       :     Any								=			getattr(lowerCAmelCase			,  """handle_key"""			,  []							)
														handle += [key]
														setattr(lowerCAmelCase			,  """handle_key"""			,  lowerCAmelCase							)
														return func

							return decorator

def  _A(		*lowerCAmelCase							):

							def decorator(lowerCAmelCase							):
														A__       :     Optional[Any]								=			getattr(lowerCAmelCase			,  """handle_key"""			,  []							)
														handle += keys
														setattr(lowerCAmelCase			,  """handle_key"""			,  lowerCAmelCase							)
														return func

							return decorator
class 						__UpperCAmelCase    (__A     ):







			'''simple docstring'''





			def __new__(			cls   ,      snake_case_   ,      snake_case_   ,      snake_case_							):







										'''simple docstring'''


										A__       :     Union[str, Any]								=			super().__new__(cls   ,      snake_case_   ,      snake_case_   ,      snake_case_							)
										if not hasattr(snake_case_   ,      """key_handler"""							):
																	setattr(snake_case_   ,      """key_handler"""   ,      {}							)
										setattr(snake_case_   ,      """handle_input"""   ,      KeyHandler.handle_input							)

										for value in attrs.values():
																	A__       :     Dict								=			getattr(snake_case_   ,      """handle_key"""   ,      []							)
																	for key in handled_keys:
																								A__       :     Optional[Any]								=			value
										return new_cls





			@staticmethod
			def 		lowerCamelCase					(			cls							):







										'''simple docstring'''


										A__       :     int								=			get_character()
										if char != KEYMAP["undefined"]:
																	A__       :     Union[str, Any]								=			ord(snake_case_							)
										A__       :     int								=			cls.key_handler.get(snake_case_							)
										if handler:
																	A__       :     int								=			char
																	return handler(cls							)
										else:
																	return None





def  _A(		cls							):
							return KeyHandler(cls.__name__			,  cls.__bases__			,  cls.__dict__.copy()							)





