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 						SCREAMING_SNAKE_CASE					(  __UpperCAmelCase       )							->   Optional[Any]:
				if isinstance(_A  ,					collections.abc.Iterable       ):
								return x
				return (x, x)

@require_flax
class lowerCamelCase		:






					'''simple docstring'''


					def 			lowerCAmelCase_					(						self			:				List[Any]   , _snake_case			:				Optional[int]   , _snake_case			:				Union[str, Any]		)							->       Optional[int]:
									pass



					def 			lowerCAmelCase_					(						self			:				Union[str, Any]		)							->       Optional[int]:
									pass



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



					def 			lowerCAmelCase_					(						self			:				Optional[int]   , _snake_case			:				Optional[int]   , _snake_case			:				Optional[Any]   , _snake_case			:				Dict		)							->       Any:
									SCREAMING_SNAKE_CASE__		  =				np.abs((a - b)		).max()
									self.assertLessEqual(_SCREAMING_SNAKE_CASE   , _SCREAMING_SNAKE_CASE   , F"""Difference between torch and flax is {diff} (>= {tol})."""		)



					def 			lowerCAmelCase_					(						self			:				Optional[Any]   , _snake_case			:				List[str]   , _snake_case			:				str   , _snake_case			:				Any   , _snake_case			:				Optional[int]   , _snake_case			:				str=None   , **_snake_case			:				List[str]		)							->       int:
									SCREAMING_SNAKE_CASE__		  =				VisionTextDualEncoderConfig.from_vision_text_configs(_SCREAMING_SNAKE_CASE   , _SCREAMING_SNAKE_CASE		)

									SCREAMING_SNAKE_CASE__		  =				FlaxVisionTextDualEncoderModel(_SCREAMING_SNAKE_CASE		)

									SCREAMING_SNAKE_CASE__		  =				model(input_ids=_SCREAMING_SNAKE_CASE   , pixel_values=_SCREAMING_SNAKE_CASE   , attention_mask=_SCREAMING_SNAKE_CASE		)

									self.assertEqual(output["text_embeds"].shape   , (input_ids.shape[0], config.projection_dim)		)
									self.assertEqual(output["image_embeds"].shape   , (pixel_values.shape[0], config.projection_dim)		)



					def 			lowerCAmelCase_					(						self			:				Tuple   , _snake_case			:				str   , _snake_case			:				Optional[int]   , _snake_case			:				Optional[int]   , _snake_case			:				List[Any]   , _snake_case			:				List[Any]=None   , **_snake_case			:				Dict		)							->       str:
									SCREAMING_SNAKE_CASE__  ,       SCREAMING_SNAKE_CASE__		  =				self.get_vision_text_model(_SCREAMING_SNAKE_CASE   , _SCREAMING_SNAKE_CASE		)
									SCREAMING_SNAKE_CASE__		  =				{"vision_model": vision_model, "text_model": text_model}
									SCREAMING_SNAKE_CASE__		  =				FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(**_SCREAMING_SNAKE_CASE		)

									SCREAMING_SNAKE_CASE__		  =				model(input_ids=_SCREAMING_SNAKE_CASE   , pixel_values=_SCREAMING_SNAKE_CASE   , attention_mask=_SCREAMING_SNAKE_CASE		)

									self.assertEqual(output["text_embeds"].shape   , (input_ids.shape[0], model.config.projection_dim)		)
									self.assertEqual(output["image_embeds"].shape   , (pixel_values.shape[0], model.config.projection_dim)		)



					def 			lowerCAmelCase_					(						self			:				Optional[int]   , _snake_case			:				int   , _snake_case			:				List[str]   , _snake_case			:				str   , _snake_case			:				str   , _snake_case			:				List[Any]=None   , **_snake_case			:				Dict		)							->       List[str]:
									SCREAMING_SNAKE_CASE__  ,       SCREAMING_SNAKE_CASE__		  =				self.get_vision_text_model(_SCREAMING_SNAKE_CASE   , _SCREAMING_SNAKE_CASE		)
									SCREAMING_SNAKE_CASE__		  =				{"vision_model": vision_model, "text_model": text_model}
									SCREAMING_SNAKE_CASE__		  =				FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(**_SCREAMING_SNAKE_CASE		)

									SCREAMING_SNAKE_CASE__		  =				model(input_ids=_SCREAMING_SNAKE_CASE   , pixel_values=_SCREAMING_SNAKE_CASE   , attention_mask=_SCREAMING_SNAKE_CASE		)
									SCREAMING_SNAKE_CASE__		  =				output[0]

									with tempfile.TemporaryDirectory() as tmpdirname:
													model.save_pretrained(_SCREAMING_SNAKE_CASE		)
													SCREAMING_SNAKE_CASE__		  =				FlaxVisionTextDualEncoderModel.from_pretrained(_SCREAMING_SNAKE_CASE		)

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



					def 			lowerCAmelCase_					(						self			:				Optional[Any]   , _snake_case			:				Any   , _snake_case			:				Tuple   , _snake_case			:				Union[str, Any]   , _snake_case			:				Optional[int]   , _snake_case			:				Any=None   , **_snake_case			:				Union[str, Any]		)							->       Tuple:
									SCREAMING_SNAKE_CASE__  ,       SCREAMING_SNAKE_CASE__		  =				self.get_vision_text_model(_SCREAMING_SNAKE_CASE   , _SCREAMING_SNAKE_CASE		)
									SCREAMING_SNAKE_CASE__		  =				{"vision_model": vision_model, "text_model": text_model}
									SCREAMING_SNAKE_CASE__		  =				FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(**_SCREAMING_SNAKE_CASE		)

									SCREAMING_SNAKE_CASE__		  =				model(
									    input_ids=_SCREAMING_SNAKE_CASE   , pixel_values=_SCREAMING_SNAKE_CASE   , attention_mask=_SCREAMING_SNAKE_CASE   , output_attentions=_SCREAMING_SNAKE_CASE		)

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

									SCREAMING_SNAKE_CASE__		  =				output.text_model_output.attentions
									self.assertEqual(len(_SCREAMING_SNAKE_CASE		)   , text_config.num_hidden_layers		)

									self.assertEqual(
									    text_attentions[0].shape[-3:]   , (text_config.num_attention_heads, input_ids.shape[-1], input_ids.shape[-1])   , )



					def 			lowerCAmelCase_					(						self			:				str   , _snake_case			:				int   , _snake_case			:				Any   , _snake_case			:				Dict		)							->       Any:
									pt_model.to(_SCREAMING_SNAKE_CASE		)
									pt_model.eval()

									# prepare inputs
									SCREAMING_SNAKE_CASE__		  =				inputs_dict
									SCREAMING_SNAKE_CASE__		  =				{k: torch.tensor(v.tolist()		) for k, v in flax_inputs.items()}

									with torch.no_grad():
													SCREAMING_SNAKE_CASE__		  =				pt_model(**_SCREAMING_SNAKE_CASE		).to_tuple()

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

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

									SCREAMING_SNAKE_CASE__		  =				fx_model_loaded(**_SCREAMING_SNAKE_CASE		).to_tuple()
									self.assertEqual(len(_SCREAMING_SNAKE_CASE		)   , len(_SCREAMING_SNAKE_CASE		)   , "Output lengths differ between Flax and PyTorch"		)
									for fx_output_loaded, pt_output in zip(fx_outputs_loaded[:4]   , pt_outputs[:4]		):
													self.assert_almost_equals(_SCREAMING_SNAKE_CASE   , pt_output.numpy()   , 4e-2		)

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

									pt_model_loaded.to(_SCREAMING_SNAKE_CASE		)
									pt_model_loaded.eval()

									with torch.no_grad():
													SCREAMING_SNAKE_CASE__		  =				pt_model_loaded(**_SCREAMING_SNAKE_CASE		).to_tuple()

									self.assertEqual(len(_SCREAMING_SNAKE_CASE		)   , len(_SCREAMING_SNAKE_CASE		)   , "Output lengths differ between Flax and PyTorch"		)
									for fx_output, pt_output_loaded in zip(fx_outputs[:4]   , pt_outputs_loaded[:4]		):
													self.assert_almost_equals(_SCREAMING_SNAKE_CASE   , pt_output_loaded.numpy()   , 4e-2		)



					def 			lowerCAmelCase_					(						self			:				Dict   , _snake_case			:				Any   , _snake_case			:				Tuple   , _snake_case			:				Union[str, Any]		)							->       Tuple:
									SCREAMING_SNAKE_CASE__		  =				VisionTextDualEncoderConfig.from_vision_text_configs(_SCREAMING_SNAKE_CASE   , _SCREAMING_SNAKE_CASE		)

									SCREAMING_SNAKE_CASE__		  =				VisionTextDualEncoderModel(_SCREAMING_SNAKE_CASE		)
									SCREAMING_SNAKE_CASE__		  =				FlaxVisionTextDualEncoderModel(_SCREAMING_SNAKE_CASE		)

									SCREAMING_SNAKE_CASE__		  =				convert_pytorch_state_dict_to_flax(pt_model.state_dict()   , _SCREAMING_SNAKE_CASE		)
									SCREAMING_SNAKE_CASE__		  =				fx_state

									self.check_pt_flax_equivalence(_SCREAMING_SNAKE_CASE   , _SCREAMING_SNAKE_CASE   , _SCREAMING_SNAKE_CASE		)



					def 			lowerCAmelCase_					(						self			:				List[Any]   , _snake_case			:				str   , _snake_case			:				Optional[int]   , _snake_case			:				int		)							->       str:
									SCREAMING_SNAKE_CASE__		  =				VisionTextDualEncoderConfig.from_vision_text_configs(_SCREAMING_SNAKE_CASE   , _SCREAMING_SNAKE_CASE		)

									SCREAMING_SNAKE_CASE__		  =				VisionTextDualEncoderModel(_SCREAMING_SNAKE_CASE		)
									SCREAMING_SNAKE_CASE__		  =				FlaxVisionTextDualEncoderModel(_SCREAMING_SNAKE_CASE		)

									SCREAMING_SNAKE_CASE__		  =				load_flax_weights_in_pytorch_model(_SCREAMING_SNAKE_CASE   , fx_model.params		)

									self.check_pt_flax_equivalence(_SCREAMING_SNAKE_CASE   , _SCREAMING_SNAKE_CASE   , _SCREAMING_SNAKE_CASE		)



					def 			lowerCAmelCase_					(						self			:				int		)							->       Tuple:
									SCREAMING_SNAKE_CASE__		  =				self.prepare_config_and_inputs()
									self.check_model_from_pretrained_configs(**_SCREAMING_SNAKE_CASE		)



					def 			lowerCAmelCase_					(						self			:				str		)							->       List[str]:
									SCREAMING_SNAKE_CASE__		  =				self.prepare_config_and_inputs()
									self.check_vision_text_dual_encoder_from_pretrained(**_SCREAMING_SNAKE_CASE		)



					def 			lowerCAmelCase_					(						self			:				List[str]		)							->       Optional[Any]:
									SCREAMING_SNAKE_CASE__		  =				self.prepare_config_and_inputs()
									self.check_save_load(**_SCREAMING_SNAKE_CASE		)



					def 			lowerCAmelCase_					(						self			:				str		)							->       Tuple:
									SCREAMING_SNAKE_CASE__		  =				self.prepare_config_and_inputs()
									self.check_vision_text_output_attention(**_SCREAMING_SNAKE_CASE		)



					@is_pt_flax_cross_test
					def 			lowerCAmelCase_					(						self			:				str		)							->       Dict:
									SCREAMING_SNAKE_CASE__		  =				self.prepare_config_and_inputs()
									SCREAMING_SNAKE_CASE__		  =				config_inputs_dict.pop("vision_config"		)
									SCREAMING_SNAKE_CASE__		  =				config_inputs_dict.pop("text_config"		)

									SCREAMING_SNAKE_CASE__		  =				config_inputs_dict

									self.check_equivalence_pt_to_flax(_SCREAMING_SNAKE_CASE   , _SCREAMING_SNAKE_CASE   , _SCREAMING_SNAKE_CASE		)
									self.check_equivalence_flax_to_pt(_SCREAMING_SNAKE_CASE   , _SCREAMING_SNAKE_CASE   , _SCREAMING_SNAKE_CASE		)



					@slow
					def 			lowerCAmelCase_					(						self			:				List[Any]		)							->       Union[str, Any]:
									SCREAMING_SNAKE_CASE__  ,       SCREAMING_SNAKE_CASE__		  =				self.get_pretrained_model_and_inputs()

									SCREAMING_SNAKE_CASE__		  =				model_a(**_SCREAMING_SNAKE_CASE		)
									SCREAMING_SNAKE_CASE__		  =				outputs[0]

									with tempfile.TemporaryDirectory() as tmp_dirname:
													model_a.save_pretrained(_SCREAMING_SNAKE_CASE		)
													SCREAMING_SNAKE_CASE__		  =				FlaxVisionTextDualEncoderModel.from_pretrained(_SCREAMING_SNAKE_CASE		)

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



@require_flax
class lowerCamelCase		(lowerCAmelCase__    ,    unittest.TestCase      ):






					'''simple docstring'''


					def 			lowerCAmelCase_					(						self			:				Optional[int]		)							->       Tuple:
									SCREAMING_SNAKE_CASE__		  =				FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(
									    "hf-internal-testing/tiny-random-vit"   , "hf-internal-testing/tiny-bert"   , vision_from_pt=_SCREAMING_SNAKE_CASE   , text_from_pt=_SCREAMING_SNAKE_CASE   , )
									SCREAMING_SNAKE_CASE__		  =				13
									SCREAMING_SNAKE_CASE__		  =				floats_tensor(
									    [
									        batch_size,
									        model.config.vision_config.num_channels,
									        model.config.vision_config.image_size,
									        model.config.vision_config.image_size,
									    ]		)
									SCREAMING_SNAKE_CASE__		  =				ids_tensor([batch_size, 4]   , model.config.text_config.vocab_size		)
									SCREAMING_SNAKE_CASE__		  =				random_attention_mask([batch_size, 4]		)
									SCREAMING_SNAKE_CASE__		  =				{"pixel_values": pixel_values, "input_ids": input_ids, "attention_mask": attention_mask}

									return model, inputs



					def 			lowerCAmelCase_					(						self			:				Dict   , _snake_case			:				Union[str, Any]   , _snake_case			:				Union[str, Any]		)							->       Tuple:
									SCREAMING_SNAKE_CASE__		  =				FlaxViTModel(_SCREAMING_SNAKE_CASE		)
									SCREAMING_SNAKE_CASE__		  =				FlaxBertModel(_SCREAMING_SNAKE_CASE		)
									return vision_model, text_model



					def 			lowerCAmelCase_					(						self			:				Dict		)							->       int:
									SCREAMING_SNAKE_CASE__		  =				FlaxViTModelTester(self		)
									SCREAMING_SNAKE_CASE__		  =				FlaxBertModelTester(self		)
									SCREAMING_SNAKE_CASE__		  =				vit_model_tester.prepare_config_and_inputs()
									SCREAMING_SNAKE_CASE__		  =				bert_model_tester.prepare_config_and_inputs()

									SCREAMING_SNAKE_CASE__  ,       SCREAMING_SNAKE_CASE__		  =				vision_config_and_inputs

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






					'''simple docstring'''


					def 			lowerCAmelCase_					(						self			:				Tuple		)							->       Tuple:
									SCREAMING_SNAKE_CASE__		  =				FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(
									    "hf-internal-testing/tiny-random-clip"   , "hf-internal-testing/tiny-bert"   , vision_from_pt=_SCREAMING_SNAKE_CASE   , text_from_pt=_SCREAMING_SNAKE_CASE   , )
									SCREAMING_SNAKE_CASE__		  =				13
									SCREAMING_SNAKE_CASE__		  =				floats_tensor(
									    [
									        batch_size,
									        model.config.vision_config.num_channels,
									        model.config.vision_config.image_size,
									        model.config.vision_config.image_size,
									    ]		)
									SCREAMING_SNAKE_CASE__		  =				ids_tensor([batch_size, 4]   , model.config.text_config.vocab_size		)
									SCREAMING_SNAKE_CASE__		  =				random_attention_mask([batch_size, 4]		)
									SCREAMING_SNAKE_CASE__		  =				{"pixel_values": pixel_values, "input_ids": input_ids, "attention_mask": attention_mask}

									return model, inputs



					def 			lowerCAmelCase_					(						self			:				Tuple   , _snake_case			:				Tuple   , _snake_case			:				Optional[Any]		)							->       List[Any]:
									SCREAMING_SNAKE_CASE__		  =				FlaxCLIPVisionModel(_SCREAMING_SNAKE_CASE		)
									SCREAMING_SNAKE_CASE__		  =				FlaxBertModel(_SCREAMING_SNAKE_CASE		)
									return vision_model, text_model



					def 			lowerCAmelCase_					(						self			:				List[str]		)							->       str:
									SCREAMING_SNAKE_CASE__		  =				FlaxCLIPVisionModelTester(self		)
									SCREAMING_SNAKE_CASE__		  =				FlaxBertModelTester(self		)
									SCREAMING_SNAKE_CASE__		  =				clip_model_tester.prepare_config_and_inputs()
									SCREAMING_SNAKE_CASE__		  =				bert_model_tester.prepare_config_and_inputs()

									SCREAMING_SNAKE_CASE__  ,       SCREAMING_SNAKE_CASE__		  =				vision_config_and_inputs

									SCREAMING_SNAKE_CASE__  ,       SCREAMING_SNAKE_CASE__  ,       SCREAMING_SNAKE_CASE__  ,       SCREAMING_SNAKE_CASE__		  =				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			:				List[Any]		)							->       List[str]:
									SCREAMING_SNAKE_CASE__		  =				FlaxVisionTextDualEncoderModel.from_pretrained("clip-italian/clip-italian"   , logit_scale_init_value=1.0		)
									SCREAMING_SNAKE_CASE__		  =				VisionTextDualEncoderProcessor.from_pretrained("clip-italian/clip-italian"		)

									SCREAMING_SNAKE_CASE__		  =				Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png"		)
									SCREAMING_SNAKE_CASE__		  =				processor(
									    text=["una foto di un gatto", "una foto di un cane"]   , images=_SCREAMING_SNAKE_CASE   , padding=_SCREAMING_SNAKE_CASE   , return_tensors="np"		)

									SCREAMING_SNAKE_CASE__		  =				model(**_SCREAMING_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])   , )

									SCREAMING_SNAKE_CASE__		  =				np.array([[1.228_4727, 0.310_4122]]		)

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





"""simple docstring"""






import unittest

from transformers import EsmConfig, is_torch_available
from transformers.testing_utils import TestCasePlus, require_torch, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
						import torch

						from transformers import EsmForMaskedLM, EsmForSequenceClassification, EsmForTokenClassification, EsmModel
						from transformers.models.esm.modeling_esm import (
						    ESM_PRETRAINED_MODEL_ARCHIVE_LIST,
						    EsmEmbeddings,
						    create_position_ids_from_input_ids,
						)

class lowerCamelCase		:






					'''simple docstring'''


					def __init__(						self			:				Union[str, Any]   , _snake_case			:				Optional[int]   , _snake_case			:				Dict=13   , _snake_case			:				Optional[Any]=7   , _snake_case			:				Union[str, Any]=False   , _snake_case			:				Any=True   , _snake_case			:				int=False   , _snake_case			:				int=True   , _snake_case			:				Tuple=33   , _snake_case			:				Optional[int]=32   , _snake_case			:				List[Any]=5   , _snake_case			:				Union[str, Any]=4   , _snake_case			:				int=37   , _snake_case			:				Tuple="gelu"   , _snake_case			:				Dict=0.1   , _snake_case			:				Dict=0.1   , _snake_case			:				Tuple=512   , _snake_case			:				Any=16   , _snake_case			:				Union[str, Any]=2   , _snake_case			:				List[str]=0.02   , _snake_case			:				Optional[Any]=3   , _snake_case			:				int=4   , _snake_case			:				List[str]=None   , )							->       Tuple:
									SCREAMING_SNAKE_CASE__		  =				parent
									SCREAMING_SNAKE_CASE__		  =				batch_size
									SCREAMING_SNAKE_CASE__		  =				seq_length
									SCREAMING_SNAKE_CASE__		  =				is_training
									SCREAMING_SNAKE_CASE__		  =				use_input_mask
									SCREAMING_SNAKE_CASE__		  =				use_token_type_ids
									SCREAMING_SNAKE_CASE__		  =				use_labels
									SCREAMING_SNAKE_CASE__		  =				vocab_size
									SCREAMING_SNAKE_CASE__		  =				hidden_size
									SCREAMING_SNAKE_CASE__		  =				num_hidden_layers
									SCREAMING_SNAKE_CASE__		  =				num_attention_heads
									SCREAMING_SNAKE_CASE__		  =				intermediate_size
									SCREAMING_SNAKE_CASE__		  =				hidden_act
									SCREAMING_SNAKE_CASE__		  =				hidden_dropout_prob
									SCREAMING_SNAKE_CASE__		  =				attention_probs_dropout_prob
									SCREAMING_SNAKE_CASE__		  =				max_position_embeddings
									SCREAMING_SNAKE_CASE__		  =				type_vocab_size
									SCREAMING_SNAKE_CASE__		  =				type_sequence_label_size
									SCREAMING_SNAKE_CASE__		  =				initializer_range
									SCREAMING_SNAKE_CASE__		  =				num_labels
									SCREAMING_SNAKE_CASE__		  =				num_choices
									SCREAMING_SNAKE_CASE__		  =				scope



					def 			lowerCAmelCase_					(						self			:				Union[str, Any]		)							->       str:
									SCREAMING_SNAKE_CASE__		  =				ids_tensor([self.batch_size, self.seq_length]   , self.vocab_size		)

									SCREAMING_SNAKE_CASE__		  =				None
									if self.use_input_mask:
													SCREAMING_SNAKE_CASE__		  =				random_attention_mask([self.batch_size, self.seq_length]		)

									SCREAMING_SNAKE_CASE__		  =				None
									SCREAMING_SNAKE_CASE__		  =				None
									SCREAMING_SNAKE_CASE__		  =				None
									if self.use_labels:
													SCREAMING_SNAKE_CASE__		  =				ids_tensor([self.batch_size]   , self.type_sequence_label_size		)
													SCREAMING_SNAKE_CASE__		  =				ids_tensor([self.batch_size, self.seq_length]   , self.num_labels		)
													SCREAMING_SNAKE_CASE__		  =				ids_tensor([self.batch_size]   , self.num_choices		)

									SCREAMING_SNAKE_CASE__		  =				self.get_config()

									return config, input_ids, input_mask, sequence_labels, token_labels, choice_labels



					def 			lowerCAmelCase_					(						self			:				Tuple		)							->       Any:
									return EsmConfig(
									    vocab_size=self.vocab_size   , hidden_size=self.hidden_size   , pad_token_id=1   , num_hidden_layers=self.num_hidden_layers   , num_attention_heads=self.num_attention_heads   , intermediate_size=self.intermediate_size   , hidden_act=self.hidden_act   , hidden_dropout_prob=self.hidden_dropout_prob   , attention_probs_dropout_prob=self.attention_probs_dropout_prob   , max_position_embeddings=self.max_position_embeddings   , type_vocab_size=self.type_vocab_size   , initializer_range=self.initializer_range   , )



					def 			lowerCAmelCase_					(						self			:				str   , _snake_case			:				int   , _snake_case			:				Dict   , _snake_case			:				Union[str, Any]   , _snake_case			:				List[str]   , _snake_case			:				Any   , _snake_case			:				List[str]		)							->       Optional[int]:
									SCREAMING_SNAKE_CASE__		  =				EsmModel(config=_snake_case		)
									model.to(_snake_case		)
									model.eval()
									SCREAMING_SNAKE_CASE__		  =				model(_snake_case   , attention_mask=_snake_case		)
									SCREAMING_SNAKE_CASE__		  =				model(_snake_case		)
									SCREAMING_SNAKE_CASE__		  =				model(_snake_case		)

									self.parent.assertEqual(result.last_hidden_state.shape   , (self.batch_size, self.seq_length, self.hidden_size)		)
									self.parent.assertEqual(result.pooler_output.shape   , (self.batch_size, self.hidden_size)		)



					def 			lowerCAmelCase_					(						self			:				Tuple   , _snake_case			:				Union[str, Any]   , _snake_case			:				Optional[Any]   , _snake_case			:				Dict   , _snake_case			:				List[Any]   , _snake_case			:				Optional[int]   , _snake_case			:				int		)							->       int:
									SCREAMING_SNAKE_CASE__		  =				EsmForMaskedLM(config=_snake_case		)
									model.to(_snake_case		)
									model.eval()
									SCREAMING_SNAKE_CASE__		  =				model(_snake_case   , attention_mask=_snake_case   , labels=_snake_case		)
									self.parent.assertEqual(result.logits.shape   , (self.batch_size, self.seq_length, self.vocab_size)		)



					def 			lowerCAmelCase_					(						self			:				int   , _snake_case			:				List[Any]   , _snake_case			:				Union[str, Any]   , _snake_case			:				Tuple   , _snake_case			:				List[str]   , _snake_case			:				Dict   , _snake_case			:				Tuple		)							->       Dict:
									SCREAMING_SNAKE_CASE__		  =				self.num_labels
									SCREAMING_SNAKE_CASE__		  =				EsmForTokenClassification(config=_snake_case		)
									model.to(_snake_case		)
									model.eval()
									SCREAMING_SNAKE_CASE__		  =				model(_snake_case   , attention_mask=_snake_case   , labels=_snake_case		)
									self.parent.assertEqual(result.logits.shape   , (self.batch_size, self.seq_length, self.num_labels)		)



					def 			lowerCAmelCase_					(						self			:				Dict		)							->       List[Any]:
									SCREAMING_SNAKE_CASE__		  =				self.prepare_config_and_inputs()
									(
									    (
									    SCREAMING_SNAKE_CASE__
									)  ,       (
									    SCREAMING_SNAKE_CASE__
									)  ,       (
									    SCREAMING_SNAKE_CASE__
									)  ,       (
									    SCREAMING_SNAKE_CASE__
									)  ,       (
									    SCREAMING_SNAKE_CASE__
									)  ,       (
									    SCREAMING_SNAKE_CASE__
									)  ,       
									)		  =				config_and_inputs
									SCREAMING_SNAKE_CASE__		  =				{"input_ids": input_ids, "attention_mask": input_mask}
									return config, inputs_dict



@require_torch
class lowerCamelCase		(_SCREAMING_SNAKE_CASE    ,    _SCREAMING_SNAKE_CASE    ,    unittest.TestCase      ):






					'''simple docstring'''


					a   	=		False

					a   	=		(
					    (
					        EsmForMaskedLM,
					        EsmModel,
					        EsmForSequenceClassification,
					        EsmForTokenClassification,
					    )
					    if is_torch_available()
					    else ()
					)
					a   	=		()
					a   	=		(
					    {
					        "feature-extraction": EsmModel,
					        "fill-mask": EsmForMaskedLM,
					        "text-classification": EsmForSequenceClassification,
					        "token-classification": EsmForTokenClassification,
					        "zero-shot": EsmForSequenceClassification,
					    }
					    if is_torch_available()
					    else {}
					)
					a   	=		True



					def 			lowerCAmelCase_					(						self			:				List[str]		)							->       Tuple:
									SCREAMING_SNAKE_CASE__		  =				EsmModelTester(self		)
									SCREAMING_SNAKE_CASE__		  =				ConfigTester(self   , config_class=_snake_case   , hidden_size=37		)



					def 			lowerCAmelCase_					(						self			:				str		)							->       List[str]:
									self.config_tester.run_common_tests()



					def 			lowerCAmelCase_					(						self			:				str		)							->       Dict:
									SCREAMING_SNAKE_CASE__		  =				self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_model(*_snake_case		)



					def 			lowerCAmelCase_					(						self			:				Any		)							->       List[str]:
									SCREAMING_SNAKE_CASE__		  =				self.model_tester.prepare_config_and_inputs()
									for type in ["absolute", "relative_key", "relative_key_query"]:
													SCREAMING_SNAKE_CASE__		  =				type
													self.model_tester.create_and_check_model(*_snake_case		)



					def 			lowerCAmelCase_					(						self			:				str		)							->       Optional[Any]:
									SCREAMING_SNAKE_CASE__		  =				self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_for_masked_lm(*_snake_case		)



					def 			lowerCAmelCase_					(						self			:				Tuple		)							->       str:
									SCREAMING_SNAKE_CASE__		  =				self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_for_token_classification(*_snake_case		)



					@slow
					def 			lowerCAmelCase_					(						self			:				Tuple		)							->       List[Any]:
									for model_name in ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
													SCREAMING_SNAKE_CASE__		  =				EsmModel.from_pretrained(_snake_case		)
													self.assertIsNotNone(_snake_case		)



					def 			lowerCAmelCase_					(						self			:				Union[str, Any]		)							->       str:
									SCREAMING_SNAKE_CASE__		  =				self.model_tester.prepare_config_and_inputs()[0]
									SCREAMING_SNAKE_CASE__		  =				EsmEmbeddings(config=_snake_case		)

									SCREAMING_SNAKE_CASE__		  =				torch.as_tensor([[12, 31, 13, model.padding_idx]]		)
									SCREAMING_SNAKE_CASE__		  =				torch.as_tensor(
									    [
									        [
									            0 + model.padding_idx + 1,
									            1 + model.padding_idx + 1,
									            2 + model.padding_idx + 1,
									            model.padding_idx,
									        ]
									    ]		)
									SCREAMING_SNAKE_CASE__		  =				create_position_ids_from_input_ids(_snake_case   , model.padding_idx		)
									self.assertEqual(position_ids.shape   , expected_positions.shape		)
									self.assertTrue(torch.all(torch.eq(_snake_case   , _snake_case		)		)		)



					def 			lowerCAmelCase_					(						self			:				Dict		)							->       List[str]:
									SCREAMING_SNAKE_CASE__		  =				self.model_tester.prepare_config_and_inputs()[0]
									SCREAMING_SNAKE_CASE__		  =				EsmEmbeddings(config=_snake_case		)

									SCREAMING_SNAKE_CASE__		  =				torch.empty(2   , 4   , 30		)
									SCREAMING_SNAKE_CASE__		  =				[
									    0 + embeddings.padding_idx + 1,
									    1 + embeddings.padding_idx + 1,
									    2 + embeddings.padding_idx + 1,
									    3 + embeddings.padding_idx + 1,
									]
									SCREAMING_SNAKE_CASE__		  =				torch.as_tensor([expected_single_positions, expected_single_positions]		)
									SCREAMING_SNAKE_CASE__		  =				embeddings.create_position_ids_from_inputs_embeds(_snake_case		)
									self.assertEqual(position_ids.shape   , expected_positions.shape		)
									self.assertTrue(torch.all(torch.eq(_snake_case   , _snake_case		)		)		)



					@unittest.skip("Esm does not support embedding resizing"		)
					def 			lowerCAmelCase_					(						self			:				List[str]		)							->       int:
									pass



					@unittest.skip("Esm does not support embedding resizing"		)
					def 			lowerCAmelCase_					(						self			:				Any		)							->       Tuple:
									pass



					@unittest.skip("Will be fixed soon by reducing the size of the model used for common tests."		)
					def 			lowerCAmelCase_					(						self			:				Union[str, Any]		)							->       Optional[int]:
									pass



@require_torch
class lowerCamelCase		(_SCREAMING_SNAKE_CASE      ):






					'''simple docstring'''


					@slow
					def 			lowerCAmelCase_					(						self			:				Dict		)							->       str:
									with torch.no_grad():
													SCREAMING_SNAKE_CASE__		  =				EsmForMaskedLM.from_pretrained("facebook/esm2_t6_8M_UR50D"		)
													model.eval()
													SCREAMING_SNAKE_CASE__		  =				torch.tensor([[0, 1, 2, 3, 4, 5]]		)
													SCREAMING_SNAKE_CASE__		  =				model(_snake_case		)[0]

													SCREAMING_SNAKE_CASE__		  =				33

													SCREAMING_SNAKE_CASE__		  =				torch.Size((1, 6, vocab_size)		)
													self.assertEqual(output.shape   , _snake_case		)

													SCREAMING_SNAKE_CASE__		  =				torch.tensor(
													    [[[8.9215, -10.5898, -6.4671], [-6.3967, -13.9114, -1.1212], [-7.7812, -13.9516, -3.7406]]]		)
													self.assertTrue(torch.allclose(output[:, :3, :3]   , _snake_case   , atol=1e-4		)		)



					@slow
					def 			lowerCAmelCase_					(						self			:				Union[str, Any]		)							->       Any:
									with torch.no_grad():
													SCREAMING_SNAKE_CASE__		  =				EsmModel.from_pretrained("facebook/esm2_t6_8M_UR50D"		)
													model.eval()

													SCREAMING_SNAKE_CASE__		  =				torch.tensor([[0, 6, 4, 13, 5, 4, 16, 12, 11, 7, 2]]		)
													SCREAMING_SNAKE_CASE__		  =				model(_snake_case		)[0]
													# compare the actual values for a slice.
													SCREAMING_SNAKE_CASE__		  =				torch.tensor(
													    [[[0.1444, 0.5413, 0.3248], [0.3034, 0.0053, 0.3108], [0.3228, -0.2499, 0.3415]]]		)
													self.assertTrue(torch.allclose(output[:, :3, :3]   , _snake_case   , atol=1e-4		)		)




