

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 					__lowercase  :





							"""simple docstring"""

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




							def       snake_case			(					self						)		->       int:
														A      :					Optional[Any]      		=					ids_tensor([self.batch_size, self.seq_length]    ,						self.vocab_size						)

														A      :					Dict      		=					None
														if self.use_input_mask:
																					A      :					Optional[int]      		=					random_attention_mask([self.batch_size, self.seq_length]						)

														A      :					Optional[int]      		=					None
														A      :					int      		=					None
														A      :					Dict      		=					None
														if self.use_labels:
																					A      :					Union[str, Any]      		=					ids_tensor([self.batch_size]    ,						self.type_sequence_label_size						)
																					A      :					Optional[int]      		=					ids_tensor([self.batch_size, self.seq_length]    ,						self.num_labels						)
																					A      :					str      		=					ids_tensor([self.batch_size]    ,						self.num_choices						)

														A      :					str      		=					self.get_config()

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




							def       snake_case			(					self						)		->       Tuple:
														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       snake_case			(					self    ,						__UpperCAmelCase    ,						__UpperCAmelCase    ,						__UpperCAmelCase    ,						__UpperCAmelCase    ,						__UpperCAmelCase    ,						__UpperCAmelCase						)		->       List[Any]:
														A      :					int      		=					EsmModel(config=__UpperCAmelCase						)
														model.to(__UpperCAmelCase						)
														model.eval()
														A      :					int      		=					model(__UpperCAmelCase    ,						attention_mask=__UpperCAmelCase						)
														A      :					str      		=					model(__UpperCAmelCase						)
														A      :					Dict      		=					model(__UpperCAmelCase						)

														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       snake_case			(					self    ,						__UpperCAmelCase    ,						__UpperCAmelCase    ,						__UpperCAmelCase    ,						__UpperCAmelCase    ,						__UpperCAmelCase    ,						__UpperCAmelCase						)		->       Union[str, Any]:
														A      :					Dict      		=					EsmForMaskedLM(config=__UpperCAmelCase						)
														model.to(__UpperCAmelCase						)
														model.eval()
														A      :					List[Any]      		=					model(__UpperCAmelCase    ,						attention_mask=__UpperCAmelCase    ,						labels=__UpperCAmelCase						)
														self.parent.assertEqual(result.logits.shape    ,						(self.batch_size, self.seq_length, self.vocab_size)						)




							def       snake_case			(					self    ,						__UpperCAmelCase    ,						__UpperCAmelCase    ,						__UpperCAmelCase    ,						__UpperCAmelCase    ,						__UpperCAmelCase    ,						__UpperCAmelCase						)		->       List[Any]:
														A      :					List[Any]      		=					self.num_labels
														A      :					Union[str, Any]      		=					EsmForTokenClassification(config=__UpperCAmelCase						)
														model.to(__UpperCAmelCase						)
														model.eval()
														A      :					Optional[Any]      		=					model(__UpperCAmelCase    ,						attention_mask=__UpperCAmelCase    ,						labels=__UpperCAmelCase						)
														self.parent.assertEqual(result.logits.shape    ,						(self.batch_size, self.seq_length, self.num_labels)						)






							def       snake_case			(					self						)		->       List[str]:
														A      :					List[Any]      		=					self.prepare_config_and_inputs()
														(
														    (
														    A
														)		,		(
														    A
														)		,		(
														    A
														)		,		(
														    A
														)		,		(
														    A
														)		,		(
														    A
														)		,		
														)      :					Tuple      		=					config_and_inputs
														A      :					List[str]      		=					{'''input_ids''': input_ids, '''attention_mask''': input_mask}
														return config, inputs_dict







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





							"""simple docstring"""

							UpperCAmelCase_						:		int								=	False

							UpperCAmelCase_						:		Tuple								=	(
							    (
							        EsmForMaskedLM,
							        EsmModel,
							        EsmForSequenceClassification,
							        EsmForTokenClassification,
							    )
							    if is_torch_available()
							    else ()
							)
							UpperCAmelCase_						:		Union[str, Any]								=	()
							UpperCAmelCase_						:		Optional[Any]								=	(
							    {
							        '''feature-extraction''': EsmModel,
							        '''fill-mask''': EsmForMaskedLM,
							        '''text-classification''': EsmForSequenceClassification,
							        '''token-classification''': EsmForTokenClassification,
							        '''zero-shot''': EsmForSequenceClassification,
							    }
							    if is_torch_available()
							    else {}
							)
							UpperCAmelCase_						:		Union[str, Any]								=	True




							def       snake_case			(					self						)		->       Optional[Any]:
														A      :					str      		=					EsmModelTester(self						)
														A      :					Dict      		=					ConfigTester(self    ,						config_class=__UpperCAmelCase    ,						hidden_size=37						)




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




							def       snake_case			(					self						)		->       Any:
														A      :					Tuple      		=					self.model_tester.prepare_config_and_inputs()
														self.model_tester.create_and_check_model(*__UpperCAmelCase						)




							def       snake_case			(					self						)		->       int:
														A      :					Optional[int]      		=					self.model_tester.prepare_config_and_inputs()
														for type in ["absolute", "relative_key", "relative_key_query"]:
																					A      :					List[Any]      		=					type
																					self.model_tester.create_and_check_model(*__UpperCAmelCase						)




							def       snake_case			(					self						)		->       List[str]:
														A      :					List[Any]      		=					self.model_tester.prepare_config_and_inputs()
														self.model_tester.create_and_check_for_masked_lm(*__UpperCAmelCase						)




							def       snake_case			(					self						)		->       Any:
														A      :					Dict      		=					self.model_tester.prepare_config_and_inputs()
														self.model_tester.create_and_check_for_token_classification(*__UpperCAmelCase						)




							@slow
							def       snake_case			(					self						)		->       Dict:
														for model_name in ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																					A      :					Union[str, Any]      		=					EsmModel.from_pretrained(__UpperCAmelCase						)
																					self.assertIsNotNone(__UpperCAmelCase						)




							def       snake_case			(					self						)		->       int:
														A      :					Optional[int]      		=					self.model_tester.prepare_config_and_inputs()[0]
														A      :					Any      		=					EsmEmbeddings(config=__UpperCAmelCase						)

														A      :					Tuple      		=					torch.as_tensor([[12, 31, 13, model.padding_idx]]						)
														A      :					str      		=					torch.as_tensor(
														    [
														        [
														            0 + model.padding_idx + 1,
														            1 + model.padding_idx + 1,
														            2 + model.padding_idx + 1,
														            model.padding_idx,
														        ]
														    ]						)
														A      :					str      		=					create_position_ids_from_input_ids(__UpperCAmelCase    ,						model.padding_idx						)
														self.assertEqual(position_ids.shape    ,						expected_positions.shape						)
														self.assertTrue(torch.all(torch.eq(__UpperCAmelCase    ,						__UpperCAmelCase						)						)						)




							def       snake_case			(					self						)		->       Dict:
														A      :					str      		=					self.model_tester.prepare_config_and_inputs()[0]
														A      :					Tuple      		=					EsmEmbeddings(config=__UpperCAmelCase						)

														A      :					List[str]      		=					torch.empty(2    ,						4    ,						30						)
														A      :					Optional[Any]      		=					[
														    0 + embeddings.padding_idx + 1,
														    1 + embeddings.padding_idx + 1,
														    2 + embeddings.padding_idx + 1,
														    3 + embeddings.padding_idx + 1,
														]
														A      :					int      		=					torch.as_tensor([expected_single_positions, expected_single_positions]						)
														A      :					Optional[Any]      		=					embeddings.create_position_ids_from_inputs_embeds(__UpperCAmelCase						)
														self.assertEqual(position_ids.shape    ,						expected_positions.shape						)
														self.assertTrue(torch.all(torch.eq(__UpperCAmelCase    ,						__UpperCAmelCase						)						)						)




							@unittest.skip('''Esm does not support embedding resizing'''						)
							def       snake_case			(					self						)		->       str:
														pass




							@unittest.skip('''Esm does not support embedding resizing'''						)
							def       snake_case			(					self						)		->       Optional[int]:
														pass






							@unittest.skip('''Will be fixed soon by reducing the size of the model used for common tests.'''						)
							def       snake_case			(					self						)		->       Optional[Any]:
														pass







@require_torch
class 					__lowercase  (						_SCREAMING_SNAKE_CASE				):





							"""simple docstring"""

							@slow
							def       snake_case			(					self						)		->       Tuple:
														with torch.no_grad():
																					A      :					Optional[int]      		=					EsmForMaskedLM.from_pretrained('''facebook/esm2_t6_8M_UR50D'''						)
																					model.eval()
																					A      :					Optional[int]      		=					torch.tensor([[0, 1, 2, 3, 4, 5]]						)
																					A      :					Union[str, Any]      		=					model(__UpperCAmelCase						)[0]

																					A      :					Optional[int]      		=					33

																					A      :					Dict      		=					torch.Size((1, 6, vocab_size)						)
																					self.assertEqual(output.shape    ,						__UpperCAmelCase						)

																					A      :					List[str]      		=					torch.tensor(
																					    [[[8.9_2_1_5, -1_0.5_8_9_8, -6.4_6_7_1], [-6.3_9_6_7, -1_3.9_1_1_4, -1.1_2_1_2], [-7.7_8_1_2, -1_3.9_5_1_6, -3.7_4_0_6]]]						)
																					self.assertTrue(torch.allclose(output[:, :3, :3]    ,						__UpperCAmelCase    ,						atol=1E-4						)						)






							@slow
							def       snake_case			(					self						)		->       int:
														with torch.no_grad():
																					A      :					str      		=					EsmModel.from_pretrained('''facebook/esm2_t6_8M_UR50D'''						)
																					model.eval()

																					A      :					Any      		=					torch.tensor([[0, 6, 4, 13, 5, 4, 16, 12, 11, 7, 2]]						)
																					A      :					Optional[int]      		=					model(__UpperCAmelCase						)[0]
																					# compare the actual values for a slice.
																					A      :					str      		=					torch.tensor(
																					    [[[0.1_4_4_4, 0.5_4_1_3, 0.3_2_4_8], [0.3_0_3_4, 0.0_0_5_3, 0.3_1_0_8], [0.3_2_2_8, -0.2_4_9_9, 0.3_4_1_5]]]						)
																					self.assertTrue(torch.allclose(output[:, :3, :3]    ,						__UpperCAmelCase    ,						atol=1E-4						)						)









from typing import Dict, List, Optional, Union

import numpy as np

from ...image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict
from ...image_transforms import center_crop, normalize, rescale, resize, to_channel_dimension_format
from ...image_utils import (
    IMAGENET_STANDARD_MEAN,
    IMAGENET_STANDARD_STD,
    ChannelDimension,
    ImageInput,
    PILImageResampling,
    make_list_of_images,
    to_numpy_array,
    valid_images,
)
from ...utils import TensorType, is_vision_available, logging


if is_vision_available():
		import PIL


lowercase			:				Tuple   =     logging.get_logger(__name__)







class 					__lowercase  (						_SCREAMING_SNAKE_CASE				):





							"""simple docstring"""

							UpperCAmelCase_						:		Optional[Any]								=	['''pixel_values''']




							def __init__(					self    ,						__UpperCAmelCase = True    ,						__UpperCAmelCase = None    ,						__UpperCAmelCase = PIL.Image.BICUBIC    ,						__UpperCAmelCase = True    ,						__UpperCAmelCase = None    ,						__UpperCAmelCase = 1 / 2_55    ,						__UpperCAmelCase = True    ,						__UpperCAmelCase = True    ,						__UpperCAmelCase = None    ,						__UpperCAmelCase = None    ,						**__UpperCAmelCase    ,						)		->       None:
														super().__init__(**__UpperCAmelCase						)
														A      :					Any      		=					size if size is not None else {'''height''': 2_56, '''width''': 2_56}
														A      :					Any      		=					get_size_dict(__UpperCAmelCase						)
														A      :					List[Any]      		=					crop_size if crop_size is not None else {'''height''': 2_24, '''width''': 2_24}
														A      :					List[Any]      		=					get_size_dict(__UpperCAmelCase    ,						param_name='''crop_size'''						)

														A      :					Dict      		=					do_resize
														A      :					Tuple      		=					size
														A      :					Union[str, Any]      		=					resample
														A      :					Dict      		=					do_center_crop
														A      :					int      		=					crop_size
														A      :					Union[str, Any]      		=					do_rescale
														A      :					str      		=					rescale_factor
														A      :					Optional[Any]      		=					do_normalize
														A      :					Dict      		=					image_mean if image_mean is not None else IMAGENET_STANDARD_MEAN
														A      :					Dict      		=					image_std if image_std is not None else IMAGENET_STANDARD_STD




							def       snake_case			(					self    ,						__UpperCAmelCase    ,						__UpperCAmelCase    ,						__UpperCAmelCase = PIL.Image.BICUBIC    ,						__UpperCAmelCase = None    ,						**__UpperCAmelCase    ,						)		->       np.ndarray:
														A      :					int      		=					get_size_dict(__UpperCAmelCase						)
														if "height" not in size or "width" not in size:
																					raise ValueError(f'The size dictionary must have keys \'height\' and \'width\'. Got {size.keys()}'						)
														return resize(
														    __UpperCAmelCase    ,						size=(size['''height'''], size['''width'''])    ,						resample=__UpperCAmelCase    ,						data_format=__UpperCAmelCase    ,						**__UpperCAmelCase						)




							def       snake_case			(					self    ,						__UpperCAmelCase    ,						__UpperCAmelCase    ,						__UpperCAmelCase = None    ,						**__UpperCAmelCase    ,						)		->       np.ndarray:
														A      :					Optional[Any]      		=					get_size_dict(__UpperCAmelCase						)
														if "height" not in size or "width" not in size:
																					raise ValueError(f'The size dictionary must have keys \'height\' and \'width\'. Got {size.keys()}'						)
														return center_crop(__UpperCAmelCase    ,						size=(size['''height'''], size['''width'''])    ,						data_format=__UpperCAmelCase    ,						**__UpperCAmelCase						)




							def       snake_case			(					self    ,						__UpperCAmelCase    ,						__UpperCAmelCase    ,						__UpperCAmelCase = None    ,						**__UpperCAmelCase    ,						)		->       List[str]:
														return rescale(__UpperCAmelCase    ,						scale=__UpperCAmelCase    ,						data_format=__UpperCAmelCase    ,						**__UpperCAmelCase						)




							def       snake_case			(					self    ,						__UpperCAmelCase    ,						__UpperCAmelCase    ,						__UpperCAmelCase    ,						__UpperCAmelCase = None    ,						**__UpperCAmelCase    ,						)		->       np.ndarray:
														return normalize(__UpperCAmelCase    ,						mean=__UpperCAmelCase    ,						std=__UpperCAmelCase    ,						data_format=__UpperCAmelCase    ,						**__UpperCAmelCase						)






							def       snake_case			(					self    ,						__UpperCAmelCase    ,						__UpperCAmelCase = None    ,						__UpperCAmelCase = None    ,						__UpperCAmelCase=None    ,						__UpperCAmelCase = None    ,						__UpperCAmelCase = None    ,						__UpperCAmelCase = None    ,						__UpperCAmelCase = None    ,						__UpperCAmelCase = None    ,						__UpperCAmelCase = None    ,						__UpperCAmelCase = None    ,						__UpperCAmelCase = None    ,						__UpperCAmelCase = ChannelDimension.FIRST    ,						**__UpperCAmelCase    ,						)		->       PIL.Image.Image:
														A      :					Optional[int]      		=					do_resize if do_resize is not None else self.do_resize
														A      :					Optional[Any]      		=					resample if resample is not None else self.resample
														A      :					Tuple      		=					do_center_crop if do_center_crop is not None else self.do_center_crop
														A      :					Optional[Any]      		=					do_rescale if do_rescale is not None else self.do_rescale
														A      :					str      		=					rescale_factor if rescale_factor is not None else self.rescale_factor
														A      :					Tuple      		=					do_normalize if do_normalize is not None else self.do_normalize
														A      :					Optional[int]      		=					image_mean if image_mean is not None else self.image_mean
														A      :					Dict      		=					image_std if image_std is not None else self.image_std

														A      :					List[Any]      		=					size if size is not None else self.size
														A      :					List[Any]      		=					get_size_dict(__UpperCAmelCase						)
														A      :					Dict      		=					crop_size if crop_size is not None else self.crop_size
														A      :					Tuple      		=					get_size_dict(__UpperCAmelCase    ,						param_name='''crop_size'''						)

														A      :					Any      		=					make_list_of_images(__UpperCAmelCase						)

														if not valid_images(__UpperCAmelCase						):
																					raise ValueError(
																					    '''Invalid image type. Must be of type PIL.Image.Image, numpy.ndarray, '''
																					    '''torch.Tensor, tf.Tensor or jax.ndarray.'''						)

														if do_resize and size is None or resample is None:
																					raise ValueError('''Size and resample must be specified if do_resize is True.'''						)

														if do_center_crop and crop_size is None:
																					raise ValueError('''Crop size must be specified if do_center_crop is True.'''						)

														if do_rescale and rescale_factor is None:
																					raise ValueError('''Rescale factor must be specified if do_rescale is True.'''						)

														if do_normalize and (image_mean is None or image_std is None):
																					raise ValueError('''Image mean and std must be specified if do_normalize is True.'''						)

														# All transformations expect numpy arrays.
														A      :					List[Any]      		=					[to_numpy_array(__UpperCAmelCase						) for image in images]

														if do_resize:
																					A      :					List[str]      		=					[self.resize(image=__UpperCAmelCase    ,						size=__UpperCAmelCase    ,						resample=__UpperCAmelCase						) for image in images]

														if do_center_crop:
																					A      :					Dict      		=					[self.center_crop(image=__UpperCAmelCase    ,						size=__UpperCAmelCase						) for image in images]

														if do_rescale:
																					A      :					Union[str, Any]      		=					[self.rescale(image=__UpperCAmelCase    ,						scale=__UpperCAmelCase						) for image in images]

														if do_normalize:
																					A      :					Tuple      		=					[self.normalize(image=__UpperCAmelCase    ,						mean=__UpperCAmelCase    ,						std=__UpperCAmelCase						) for image in images]

														A      :					Optional[Any]      		=					[to_channel_dimension_format(__UpperCAmelCase    ,						__UpperCAmelCase						) for image in images]

														A      :					List[Any]      		=					{'''pixel_values''': images}
														return BatchFeature(data=__UpperCAmelCase    ,						tensor_type=__UpperCAmelCase						)






