

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 					_UpperCAmelCase				(						a__):

			'''simple docstring'''



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



@require_tf
class        A__:
						"""simple docstring"""






						def 			UpperCamelCase__					(    self  ,      _lowercase  ,      _lowercase    )     -> Dict:
									pass





						def 			UpperCamelCase__					(    self    )     -> Any:
									pass





						def 			UpperCamelCase__					(    self    )     -> List[str]:
									pass





						def 			UpperCamelCase__					(    self  ,      _lowercase  ,      _lowercase  ,      _lowercase  ,      _lowercase  ,      _lowercase=None  ,      **_lowercase    )     -> List[str]:
									a_						:						str 							= VisionTextDualEncoderConfig.from_vision_text_configs(_lowercase  ,      _lowercase    )

									a_						:						Union[str, Any] 							= TFVisionTextDualEncoderModel(_lowercase    )

									a_						:						Optional[Any] 							= model(input_ids=_lowercase  ,      pixel_values=_lowercase  ,      attention_mask=_lowercase    )

									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  ,      _lowercase  ,      _lowercase  ,      _lowercase  ,      _lowercase  ,      _lowercase=None  ,      **_lowercase    )     -> Optional[Any]:
									a_							,					a_						:						List[Any] 							= self.get_vision_text_model(_lowercase  ,      _lowercase    )
									a_						:						Any 							= TFVisionTextDualEncoderModel(vision_model=_lowercase  ,      text_model=_lowercase    )

									a_						:						List[Any] 							= model(input_ids=_lowercase  ,      pixel_values=_lowercase  ,      attention_mask=_lowercase    )

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

									a_						:						int 							= model(input_ids=_lowercase  ,      pixel_values=_lowercase  ,      attention_mask=_lowercase    )

									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  ,      _lowercase  ,      _lowercase  ,      _lowercase  ,      _lowercase  ,      _lowercase=None  ,      **_lowercase    )     -> str:
									a_							,					a_						:						Any 							= self.get_vision_text_model(_lowercase  ,      _lowercase    )
									a_						:						List[Any] 							= TFVisionTextDualEncoderModel(vision_model=_lowercase  ,      text_model=_lowercase    )

									a_						:						Optional[Any] 							= model(input_ids=_lowercase  ,      pixel_values=_lowercase  ,      attention_mask=_lowercase    )
									a_						:						Tuple 							= output[0].numpy()

									with tempfile.TemporaryDirectory() as tmpdirname:
												model.save_pretrained(_lowercase    )
												a_						:						Union[str, Any] 							= TFVisionTextDualEncoderModel.from_pretrained(_lowercase    )

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





						def 			UpperCamelCase__					(    self  ,      _lowercase  ,      _lowercase  ,      _lowercase  ,      _lowercase  ,      _lowercase=None  ,      **_lowercase    )     -> List[Any]:
									a_							,					a_						:						str 							= self.get_vision_text_model(_lowercase  ,      _lowercase    )
									a_						:						Any 							= TFVisionTextDualEncoderModel(vision_model=_lowercase  ,      text_model=_lowercase    )

									a_						:						Any 							= model(
									    input_ids=_lowercase  ,      pixel_values=_lowercase  ,      attention_mask=_lowercase  ,      output_attentions=_lowercase    )

									a_						:						Dict 							= output.vision_model_output.attentions
									self.assertEqual(len(_lowercase    )  ,      vision_config.num_hidden_layers    )

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

									a_						:						Any 							= output.text_model_output.attentions
									self.assertEqual(len(_lowercase    )  ,      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  ,      _lowercase  ,      _lowercase  ,      _lowercase    )     -> Optional[int]:
									a_						:						Tuple 							= np.abs((a - b)    ).max()
									self.assertLessEqual(_lowercase  ,      _lowercase  ,      F'''Difference between torch and flax is {diff} (>= {tol}).'''    )





						def 			UpperCamelCase__					(    self    )     -> Union[str, Any]:
									a_						:						Any 							= self.prepare_config_and_inputs()
									self.check_vision_text_dual_encoder_model(**_lowercase    )





						def 			UpperCamelCase__					(    self    )     -> Optional[int]:
									a_						:						Tuple 							= self.prepare_config_and_inputs()
									self.check_model_from_pretrained_configs(**_lowercase    )





						def 			UpperCamelCase__					(    self    )     -> Any:
									a_						:						Union[str, Any] 							= self.prepare_config_and_inputs()
									self.check_vision_text_dual_encoder_from_pretrained(**_lowercase    )





						def 			UpperCamelCase__					(    self    )     -> str:
									a_						:						Dict 							= self.prepare_config_and_inputs()
									self.check_save_load(**_lowercase    )





						def 			UpperCamelCase__					(    self    )     -> Union[str, Any]:
									a_						:						Tuple 							= self.prepare_config_and_inputs()
									self.check_vision_text_output_attention(**_lowercase    )





						@slow
						def 			UpperCamelCase__					(    self    )     -> Tuple:
									a_							,					a_						:						Optional[int] 							= self.get_pretrained_model_and_inputs()

									a_						:						Optional[Any] 							= model_a(**_lowercase    )
									a_						:						List[Any] 							= outputs[0].numpy()

									with tempfile.TemporaryDirectory() as tmp_dirname:
												model_a.save_pretrained(_lowercase    )
												a_						:						Tuple 							= TFVisionTextDualEncoderModel.from_pretrained(_lowercase    )

												a_						:						Optional[int] 							= model_a(**_lowercase    )
												a_						:						Optional[Any] 							= after_outputs[0].numpy()
												a_						:						int 							= np.amax(np.abs(out_a - out_a    )    )
												self.assertLessEqual(_lowercase  ,      1e-5    )







@require_tf
class        A__(a_,  unittest.TestCase							):
						"""simple docstring"""






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

									return model, inputs





						def 			UpperCamelCase__					(    self  ,      _lowercase  ,      _lowercase    )     -> Dict:
									a_						:						Union[str, Any] 							= TFViTModel(_lowercase  ,      name="""vision_model"""    )
									a_						:						Any 							= TFBertModel(_lowercase  ,      name="""text_model"""    )
									return vision_model, text_model





						def 			UpperCamelCase__					(    self    )     -> Tuple:
									a_						:						Tuple 							= TFViTModelTester(self    )
									a_						:						int 							= TFBertModelTester(self    )
									a_						:						str 							= vit_model_tester.prepare_config_and_inputs()
									a_						:						Optional[Any] 							= bert_model_tester.prepare_config_and_inputs()

									a_							,					a_							,					a_						:						Union[str, Any] 							= vision_config_and_inputs

									(
									    (
									    a_
									)							,					(
									    a_
									)							,					(
									    a_
									)							,					(
									    a_
									)							,					(
									    a_
									)							,					(
									    a_
									)							,					(
									    a_
									)							,					
									)						:						Union[str, Any] 							= 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        A__(a_,  unittest.TestCase							):
						"""simple docstring"""






						def 			UpperCamelCase__					(    self    )     -> List[str]:
									# DeiT repo doesn't have TF weights, but we don't actually use the weights at all so let's
									# just reinitialize it.
									a_						:						Union[str, Any] 							= TFVisionTextDualEncoderModel.from_vision_text_pretrained(
									    """Rocketknight1/tiny-random-deit-tf"""  ,      """hf-internal-testing/tiny-random-roberta"""    )
									a_						:						int 							= 13
									a_						:						Optional[Any] 							= floats_tensor(
									    [
									        batch_size,
									        model.vision_model.config.num_channels,
									        model.vision_model.config.image_size,
									        model.vision_model.config.image_size,
									    ]    )
									a_						:						Tuple 							= ids_tensor([batch_size, 4]  ,      model.text_model.config.vocab_size    )
									a_						:						Dict 							= random_attention_mask([batch_size, 4]    )
									a_						:						Optional[int] 							= {"""pixel_values""": pixel_values, """input_ids""": input_ids, """attention_mask""": attention_mask}

									return model, inputs





						def 			UpperCamelCase__					(    self  ,      _lowercase  ,      _lowercase  ,      _lowercase  ,      _lowercase  ,      _lowercase=None  ,      **_lowercase    )     -> List[Any]:
									a_							,					a_						:						Any 							= self.get_vision_text_model(_lowercase  ,      _lowercase    )
									a_						:						str 							= TFVisionTextDualEncoderModel(vision_model=_lowercase  ,      text_model=_lowercase    )

									a_						:						int 							= model(
									    input_ids=_lowercase  ,      pixel_values=_lowercase  ,      attention_mask=_lowercase  ,      output_attentions=_lowercase    )

									a_						:						Optional[int] 							= output.vision_model_output.attentions
									self.assertEqual(len(_lowercase    )  ,      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)
									a_						:						str 							= to_atuple(vision_model.config.image_size    )
									a_						:						str 							= to_atuple(vision_model.config.patch_size    )
									a_						:						int 							= (image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0])
									a_						:						Tuple 							= num_patches + 2
									self.assertEqual(vision_attentions[0].shape[-3:]  ,      (vision_config.num_attention_heads, seq_len, seq_len)    )

									a_						:						Union[str, Any] 							= output.text_model_output.attentions
									self.assertEqual(len(_lowercase    )  ,      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  ,      _lowercase  ,      _lowercase    )     -> Dict:
									a_						:						Optional[int] 							= TFDeiTModel(_lowercase  ,      name="""vision_model"""    )
									a_						:						int 							= TFRobertaModel(_lowercase  ,      name="""text_model"""    )
									return vision_model, text_model





						def 			UpperCamelCase__					(    self    )     -> Optional[Any]:
									a_						:						Optional[Any] 							= TFDeiTModelTester(self    )
									a_						:						Union[str, Any] 							= TFRobertaModelTester(self    )
									a_						:						Any 							= vit_model_tester.prepare_config_and_inputs()
									a_						:						Union[str, Any] 							= bert_model_tester.prepare_config_and_inputs()

									a_							,					a_							,					a_						:						List[Any] 							= vision_config_and_inputs

									(
									    (
									    a_
									)							,					(
									    a_
									)							,					(
									    a_
									)							,					(
									    a_
									)							,					(
									    a_
									)							,					(
									    a_
									)							,					(
									    a_
									)							,					
									)						:						Dict 							= 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        A__(a_,  unittest.TestCase							):
						"""simple docstring"""






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

									return model, inputs





						def 			UpperCamelCase__					(    self  ,      _lowercase  ,      _lowercase    )     -> Union[str, Any]:
									a_						:						int 							= TFCLIPVisionModel(_lowercase  ,      name="""vision_model"""    )
									a_						:						Optional[int] 							= TFBertModel(_lowercase  ,      name="""text_model"""    )
									return vision_model, text_model





						def 			UpperCamelCase__					(    self    )     -> Optional[Any]:
									a_						:						str 							= TFCLIPVisionModelTester(self    )
									a_						:						Optional[Any] 							= TFBertModelTester(self    )
									a_						:						Any 							= clip_model_tester.prepare_config_and_inputs()
									a_						:						Union[str, Any] 							= bert_model_tester.prepare_config_and_inputs()

									a_							,					a_						:						Dict 							= vision_config_and_inputs

									(
									    (
									    a_
									)							,					(
									    a_
									)							,					(
									    a_
									)							,					(
									    a_
									)							,					(
									    a_
									)							,					(
									    a_
									)							,					(
									    a_
									)							,					
									)						:						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        A__(unittest.TestCase							):
						"""simple docstring"""






						@slow
						def 			UpperCamelCase__					(    self    )     -> List[Any]:
									a_						:						Union[str, Any] 							= TFVisionTextDualEncoderModel.from_pretrained(
									    """clip-italian/clip-italian"""  ,      logit_scale_init_value=1.0  ,      from_pt=_lowercase    )
									a_						:						Union[str, Any] 							= VisionTextDualEncoderProcessor.from_pretrained("""clip-italian/clip-italian"""    )

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

									a_						:						Union[str, Any] 							= model(**_lowercase    )

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

									a_						:						str 							= np.array([[1.2_2_8_4_7_2_7, 0.3_1_0_4_1_2_2]]    )

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





from tempfile import TemporaryDirectory
from unittest import TestCase
from unittest.mock import MagicMock, patch

from transformers import AutoModel, TFAutoModel
from transformers.onnx import FeaturesManager
from transformers.testing_utils import SMALL_MODEL_IDENTIFIER, require_tf, require_torch



@require_torch
@require_tf
class        A__(a_							):
						"""simple docstring"""






						def 			UpperCamelCase__					(    self    )     -> List[str]:
									a_						:						List[Any] 							= SMALL_MODEL_IDENTIFIER
									a_						:						Optional[int] 							= """pt"""
									a_						:						Union[str, Any] 							= """tf"""





						def 			UpperCamelCase__					(    self  ,      _lowercase    )     -> str:
									a_						:						Union[str, Any] 							= AutoModel.from_pretrained(self.test_model    )
									model_pt.save_pretrained(_lowercase    )





						def 			UpperCamelCase__					(    self  ,      _lowercase    )     -> List[str]:
									a_						:						int 							= TFAutoModel.from_pretrained(self.test_model  ,      from_pt=_lowercase    )
									model_tf.save_pretrained(_lowercase    )





						def 			UpperCamelCase__					(    self    )     -> Optional[Any]:
									a_						:						str 							= """mock_framework"""

									# Framework provided - return whatever the user provides
									a_						:						Optional[int] 							= FeaturesManager.determine_framework(self.test_model  ,      _lowercase    )
									self.assertEqual(_lowercase  ,      _lowercase    )

									# Local checkpoint and framework provided - return provided framework
									# PyTorch checkpoint
									with TemporaryDirectory() as local_pt_ckpt:
												self._setup_pt_ckpt(_lowercase    )
												a_						:						Optional[int] 							= FeaturesManager.determine_framework(_lowercase  ,      _lowercase    )
												self.assertEqual(_lowercase  ,      _lowercase    )

									# TensorFlow checkpoint
									with TemporaryDirectory() as local_tf_ckpt:
												self._setup_tf_ckpt(_lowercase    )
												a_						:						List[str] 							= FeaturesManager.determine_framework(_lowercase  ,      _lowercase    )
												self.assertEqual(_lowercase  ,      _lowercase    )





						def 			UpperCamelCase__					(    self    )     -> Union[str, Any]:
									# PyTorch checkpoint
									with TemporaryDirectory() as local_pt_ckpt:
												self._setup_pt_ckpt(_lowercase    )
												a_						:						Union[str, Any] 							= FeaturesManager.determine_framework(_lowercase    )
												self.assertEqual(_lowercase  ,      self.framework_pt    )

									# TensorFlow checkpoint
									with TemporaryDirectory() as local_tf_ckpt:
												self._setup_tf_ckpt(_lowercase    )
												a_						:						Union[str, Any] 							= FeaturesManager.determine_framework(_lowercase    )
												self.assertEqual(_lowercase  ,      self.framework_tf    )

									# Invalid local checkpoint
									with TemporaryDirectory() as local_invalid_ckpt:
												with self.assertRaises(_lowercase    ):
															a_						:						Union[str, Any] 							= FeaturesManager.determine_framework(_lowercase    )





						def 			UpperCamelCase__					(    self    )     -> List[Any]:
									a_						:						int 							= MagicMock(return_value=_lowercase    )
									with patch("""transformers.onnx.features.is_tf_available"""  ,      _lowercase    ):
												a_						:						Union[str, Any] 							= FeaturesManager.determine_framework(self.test_model    )
												self.assertEqual(_lowercase  ,      self.framework_pt    )

									# PyTorch not in environment -> use TensorFlow
									a_						:						Dict 							= MagicMock(return_value=_lowercase    )
									with patch("""transformers.onnx.features.is_torch_available"""  ,      _lowercase    ):
												a_						:						Union[str, Any] 							= FeaturesManager.determine_framework(self.test_model    )
												self.assertEqual(_lowercase  ,      self.framework_tf    )

									# Both in environment -> use PyTorch
									a_						:						Optional[Any] 							= MagicMock(return_value=_lowercase    )
									a_						:						List[Any] 							= MagicMock(return_value=_lowercase    )
									with patch("""transformers.onnx.features.is_tf_available"""  ,      _lowercase    ), patch(
									    """transformers.onnx.features.is_torch_available"""  ,      _lowercase    ):
												a_						:						Any 							= FeaturesManager.determine_framework(self.test_model    )
												self.assertEqual(_lowercase  ,      self.framework_pt    )

									# Both not in environment -> raise error
									a_						:						List[str] 							= MagicMock(return_value=_lowercase    )
									a_						:						Optional[Any] 							= MagicMock(return_value=_lowercase    )
									with patch("""transformers.onnx.features.is_tf_available"""  ,      _lowercase    ), patch(
									    """transformers.onnx.features.is_torch_available"""  ,      _lowercase    ):
												with self.assertRaises(_lowercase    ):
															a_						:						Dict 							= FeaturesManager.determine_framework(self.test_model    )


