
"""simple docstring"""






from __future__ import annotations

import unittest

from transformers import EsmConfig, is_tf_available
from transformers.testing_utils import require_tf, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
				import numpy
				import tensorflow as tf

				from transformers.models.esm.modeling_tf_esm import (
				    TF_ESM_PRETRAINED_MODEL_ARCHIVE_LIST,
				    TFEsmForMaskedLM,
				    TFEsmForSequenceClassification,
				    TFEsmForTokenClassification,
				    TFEsmModel,
				)




class 			_SCREAMING_SNAKE_CASE					:






			def __init__(					self       ,   __A       ,   )			->     Any:
						lowerCAmelCase_    :int			     =  parent
						lowerCAmelCase_    :Tuple			     =  13
						lowerCAmelCase_    :Optional[Any]			     =  7
						lowerCAmelCase_    :List[str]			     =  True
						lowerCAmelCase_    :Union[str, Any]			     =  True
						lowerCAmelCase_    :Tuple			     =  True
						lowerCAmelCase_    :int			     =  99
						lowerCAmelCase_    :Optional[Any]			     =  32
						lowerCAmelCase_    :Optional[int]			     =  2
						lowerCAmelCase_    :Optional[Any]			     =  4
						lowerCAmelCase_    :Any			     =  37
						lowerCAmelCase_    :List[Any]			     =  """gelu"""
						lowerCAmelCase_    :Optional[Any]			     =  0.1
						lowerCAmelCase_    :Dict			     =  0.1
						lowerCAmelCase_    :Union[str, Any]			     =  512
						lowerCAmelCase_    :Union[str, Any]			     =  16
						lowerCAmelCase_    :Optional[int]			     =  2
						lowerCAmelCase_    :str			     =  0.0_2
						lowerCAmelCase_    :List[Any]			     =  3
						lowerCAmelCase_    :Union[str, Any]			     =  4
						lowerCAmelCase_    :int			     =  None






			def  __lowerCAmelCase (					self       )			->     int:
						lowerCAmelCase_    :str			     =  ids_tensor([self.batch_size, self.seq_length]       ,   self.vocab_size       )

						lowerCAmelCase_    :Any			     =  None
						if self.use_input_mask:
									lowerCAmelCase_    :Union[str, Any]			     =  random_attention_mask([self.batch_size, self.seq_length]       )

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

						lowerCAmelCase_    :Optional[Any]			     =  EsmConfig(
						    vocab_size=self.vocab_size       ,   hidden_size=self.hidden_size       ,   num_hidden_layers=self.num_hidden_layers       ,   pad_token_id=1       ,   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       ,   )

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






			def  __lowerCAmelCase (					self       )			->     List[str]:
						(
						    (
						    lowerCAmelCase_
						)     ,       (
						    lowerCAmelCase_
						)     ,       (
						    lowerCAmelCase_
						)     ,       (
						    lowerCAmelCase_
						)     ,       (
						    lowerCAmelCase_
						)     ,       (
						    lowerCAmelCase_
						)     ,       
						)    :Tuple			     =  self.prepare_config_and_inputs()

						lowerCAmelCase_    :Union[str, Any]			     =  True
						lowerCAmelCase_    :Any			     =  floats_tensor([self.batch_size, self.seq_length, self.hidden_size]       )
						lowerCAmelCase_    :int			     =  ids_tensor([self.batch_size, self.seq_length]       ,   vocab_size=2       )

						return (
						    config,
						    input_ids,
						    input_mask,
						    sequence_labels,
						    token_labels,
						    choice_labels,
						    encoder_hidden_states,
						    encoder_attention_mask,
						)






			def  __lowerCAmelCase (					self       ,   __A       ,   __A       ,   __A       ,   __A       ,   __A       ,   __A       )			->     int:
						lowerCAmelCase_    :Optional[Any]			     =  TFEsmModel(config=__A       )
						lowerCAmelCase_    :Tuple			     =  {"""input_ids""": input_ids, """attention_mask""": input_mask}
						lowerCAmelCase_    :List[str]			     =  model(__A       )

						lowerCAmelCase_    :Union[str, Any]			     =  [input_ids, input_mask]
						lowerCAmelCase_    :int			     =  model(__A       )

						lowerCAmelCase_    :int			     =  model(__A       )

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






			def  __lowerCAmelCase (					self       ,   __A       ,   __A       ,   __A       ,   __A       ,   __A       ,   __A       ,   __A       ,   __A       ,   )			->     Any:
						lowerCAmelCase_    :Optional[int]			     =  True

						lowerCAmelCase_    :Tuple			     =  TFEsmModel(config=__A       )
						lowerCAmelCase_    :List[Any]			     =  {
						    """input_ids""": input_ids,
						    """attention_mask""": input_mask,
						    """encoder_hidden_states""": encoder_hidden_states,
						    """encoder_attention_mask""": encoder_attention_mask,
						}
						lowerCAmelCase_    :Dict			     =  model(__A       )

						lowerCAmelCase_    :Optional[int]			     =  [input_ids, input_mask]
						lowerCAmelCase_    :Optional[Any]			     =  model(__A       ,   encoder_hidden_states=__A       )

						# Also check the case where encoder outputs are not passed
						lowerCAmelCase_    :int			     =  model(__A       ,   attention_mask=__A       )

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






			def  __lowerCAmelCase (					self       ,   __A       ,   __A       ,   __A       ,   __A       ,   __A       ,   __A       )			->     Any:
						lowerCAmelCase_    :Optional[int]			     =  TFEsmForMaskedLM(config=__A       )
						lowerCAmelCase_    :Tuple			     =  model([input_ids, input_mask]       )
						self.parent.assertEqual(result.logits.shape       ,   (self.batch_size, self.seq_length, self.vocab_size)       )






			def  __lowerCAmelCase (					self       ,   __A       ,   __A       ,   __A       ,   __A       ,   __A       ,   __A       )			->     Optional[int]:
						lowerCAmelCase_    :Dict			     =  self.num_labels
						lowerCAmelCase_    :List[str]			     =  TFEsmForTokenClassification(config=__A       )
						lowerCAmelCase_    :str			     =  {"""input_ids""": input_ids, """attention_mask""": input_mask}
						lowerCAmelCase_    :Tuple			     =  model(__A       )
						self.parent.assertEqual(result.logits.shape       ,   (self.batch_size, self.seq_length, self.num_labels)       )






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




@require_tf
class 			_SCREAMING_SNAKE_CASE					(			A__	, A__	, unittest.TestCase							):
			UpperCAmelCase_  :Any				 =						(
			    (
			        TFEsmModel,
			        TFEsmForMaskedLM,
			        TFEsmForSequenceClassification,
			        TFEsmForTokenClassification,
			    )
			    if is_tf_available()
			    else ()
			)
			UpperCAmelCase_  :int				 =						(
			    {
			        "feature-extraction": TFEsmModel,
			        "fill-mask": TFEsmForMaskedLM,
			        "text-classification": TFEsmForSequenceClassification,
			        "token-classification": TFEsmForTokenClassification,
			        "zero-shot": TFEsmForSequenceClassification,
			    }
			    if is_tf_available()
			    else {}
			)
			UpperCAmelCase_  :Any				 =						False
			UpperCAmelCase_  :List[str]				 =						False






			def  __lowerCAmelCase (					self       )			->     List[Any]:
						lowerCAmelCase_    :str			     =  TFEsmModelTester(self       )
						lowerCAmelCase_    :List[Any]			     =  ConfigTester(self       ,   config_class=__A       ,   hidden_size=37       )






			def  __lowerCAmelCase (					self       )			->     Optional[int]:
						self.config_tester.run_common_tests()






			def  __lowerCAmelCase (					self       )			->     Optional[Any]:
						lowerCAmelCase_    :str			     =  self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_model(*__A       )






			def  __lowerCAmelCase (					self       )			->     List[str]:
						lowerCAmelCase_    :Dict			     =  self.model_tester.prepare_config_and_inputs_for_decoder()
						self.model_tester.create_and_check_model_as_decoder(*__A       )






			def  __lowerCAmelCase (					self       )			->     Optional[Any]:
						lowerCAmelCase_    :Optional[Any]			     =  self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_masked_lm(*__A       )






			def  __lowerCAmelCase (					self       )			->     List[str]:
						lowerCAmelCase_    :Optional[int]			     =  self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_token_classification(*__A       )






			@slow
			def  __lowerCAmelCase (					self       )			->     Optional[int]:
						for model_name in TF_ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
									lowerCAmelCase_    :Union[str, Any]			     =  TFEsmModel.from_pretrained(__A       )
									self.assertIsNotNone(__A       )






			@unittest.skip("""Protein models do not support embedding resizing."""       )
			def  __lowerCAmelCase (					self       )			->     str:
						pass






			@unittest.skip("""Protein models do not support embedding resizing."""       )
			def  __lowerCAmelCase (					self       )			->     Optional[Any]:
						pass






			def  __lowerCAmelCase (					self       )			->     str:
						lowerCAmelCase_     ,       lowerCAmelCase_    :List[str]			     =  self.model_tester.prepare_config_and_inputs_for_common()

						for model_class in self.all_model_classes:
									lowerCAmelCase_    :Union[str, Any]			     =  model_class(__A       )
									assert isinstance(model.get_input_embeddings()       ,   tf.keras.layers.Layer       )
									if model_class is TFEsmForMaskedLM:
												# Output embedding test differs from the main test because they're a matrix, not a layer
												lowerCAmelCase_    :Optional[int]			     =  model.get_bias()
												assert isinstance(__A       ,   __A       )
												for k, v in name.items():
															assert isinstance(__A       ,   tf.Variable       )
									else:
												lowerCAmelCase_    :Tuple			     =  model.get_output_embeddings()
												assert x is None
												lowerCAmelCase_    :Any			     =  model.get_bias()
												assert name is None




@require_tf
class 			_SCREAMING_SNAKE_CASE					(			unittest.TestCase							):






			@slow
			def  __lowerCAmelCase (					self       )			->     List[Any]:
						lowerCAmelCase_    :Any			     =  TFEsmForMaskedLM.from_pretrained("""facebook/esm2_t6_8M_UR50D"""       )

						lowerCAmelCase_    :Tuple			     =  tf.constant([[0, 1, 2, 3, 4, 5]]       )
						lowerCAmelCase_    :str			     =  model(__A       )[0]
						lowerCAmelCase_    :str			     =  [1, 6, 33]
						self.assertEqual(list(output.numpy().shape       )       ,   __A       )
						# compare the actual values for a slice.
						lowerCAmelCase_    :int			     =  tf.constant(
						    [
						        [
						            [8.9_2_1_5_1_8, -1_0.5_8_9_8_1_4, -6.4_6_7_1_3_0_7],
						            [-6.3_9_6_7_1_5_6, -1_3.9_1_1_3_7_7, -1.1_2_1_1_9_1_5],
						            [-7.7_8_1_2_4_7, -1_3.9_5_1_5_5_7, -3.7_4_0_5_9_2],
						        ]
						    ]       )
						self.assertTrue(numpy.allclose(output[:, :3, :3].numpy()       ,   expected_slice.numpy()       ,   atol=1E-2       )       )






			@slow
			def  __lowerCAmelCase (					self       )			->     int:
						lowerCAmelCase_    :List[str]			     =  TFEsmModel.from_pretrained("""facebook/esm2_t6_8M_UR50D"""       )

						lowerCAmelCase_    :Dict			     =  tf.constant([[0, 6, 4, 13, 5, 4, 16, 12, 11, 7, 2]]       )
						lowerCAmelCase_    :Tuple			     =  model(__A       )[0]
						# compare the actual values for a slice.
						lowerCAmelCase_    :int			     =  tf.constant(
						    [
						        [
						            [0.1_4_4_4_3_0_9_2, 0.5_4_1_2_5_3_2_7, 0.3_2_4_7_7_3_9],
						            [0.3_0_3_4_0_4_8_4, 0.0_0_5_2_6_6_7_6, 0.3_1_0_7_7_7_2_2],
						            [0.3_2_2_7_8_0_4_3, -0.2_4_9_8_7_0_9_6, 0.3_4_1_4_6_2_8],
						        ]
						    ]       )
						self.assertTrue(numpy.allclose(output[:, :3, :3].numpy()       ,   expected_slice.numpy()       ,   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,
    convert_to_rgb,
    get_resize_output_image_size,
    normalize,
    rescale,
    resize,
    to_channel_dimension_format,
)
from ...image_utils import (
    OPENAI_CLIP_MEAN,
    OPENAI_CLIP_STD,
    ChannelDimension,
    ImageInput,
    PILImageResampling,
    make_list_of_images,
    to_numpy_array,
    valid_images,
)
from ...utils import TensorType, is_vision_available, logging


__A   			=logging.get_logger(__name__)


if is_vision_available():
				import PIL

class 		UpperCAmelCase__			(			__UpperCamelCase   ):




			'''simple docstring'''






			UpperCamelCase									=       ["""pixel_values"""]




			def __init__(   self							: Tuple       ,       a_							: bool = True       ,       a_							: Dict[str, int] = None       ,       a_							: PILImageResampling = PILImageResampling.BICUBIC       ,       a_							: bool = True       ,       a_							: Dict[str, int] = None       ,       a_							: bool = True       ,       a_							: Union[int, float] = 1 / 2_55       ,       a_							: bool = True       ,       a_							: Optional[Union[float, List[float]]] = None       ,       a_							: Optional[Union[float, List[float]]] = None       ,       a_							: bool = True       ,       **a_							: str       ,       ):
					'''simple docstring'''





					super().__init__(**a_ )
					__UpperCAmelCase    :     List[Any]													=	size if size is not None else {'''shortest_edge''': 2_24}
					__UpperCAmelCase    :     List[str]													=	get_size_dict(a_       ,       default_to_square=a_ )
					__UpperCAmelCase    :     int													=	crop_size if crop_size is not None else {'''height''': 2_24, '''width''': 2_24}
					__UpperCAmelCase    :     Optional[int]													=	get_size_dict(a_       ,       default_to_square=a_       ,       param_name='''crop_size''' )

					__UpperCAmelCase    :     int													=	do_resize
					__UpperCAmelCase    :     Union[str, Any]													=	size
					__UpperCAmelCase    :     Union[str, Any]													=	resample
					__UpperCAmelCase    :     Any													=	do_center_crop
					__UpperCAmelCase    :     Any													=	crop_size
					__UpperCAmelCase    :     Any													=	do_rescale
					__UpperCAmelCase    :     Dict													=	rescale_factor
					__UpperCAmelCase    :     Union[str, Any]													=	do_normalize
					__UpperCAmelCase    :     Any													=	image_mean if image_mean is not None else OPENAI_CLIP_MEAN
					__UpperCAmelCase    :     int													=	image_std if image_std is not None else OPENAI_CLIP_STD
					__UpperCAmelCase    :     List[str]													=	do_convert_rgb




			def  snake_case__     (   self							: Optional[Any]       ,       a_							: np.ndarray       ,       a_							: Dict[str, int]       ,       a_							: PILImageResampling = PILImageResampling.BICUBIC       ,       a_							: Optional[Union[str, ChannelDimension]] = None       ,       **a_							: Optional[int]       ,       ):
					'''simple docstring'''





					__UpperCAmelCase    :     Dict													=	get_size_dict(a_       ,       default_to_square=a_ )
					if "shortest_edge" not in size:
							raise ValueError(F'The `size` parameter must contain the key `shortest_edge`. Got {size.keys()}' )
					__UpperCAmelCase    :     Optional[int]													=	get_resize_output_image_size(a_       ,       size=size['''shortest_edge''']       ,       default_to_square=a_ )
					return resize(a_       ,       size=a_       ,       resample=a_       ,       data_format=a_       ,       **a_ )




			def  snake_case__     (   self							: Union[str, Any]       ,       a_							: np.ndarray       ,       a_							: Dict[str, int]       ,       a_							: Optional[Union[str, ChannelDimension]] = None       ,       **a_							: Dict       ,       ):
					'''simple docstring'''





					__UpperCAmelCase    :     Optional[Any]													=	get_size_dict(a_ )
					if "height" not in size or "width" not in size:
							raise ValueError(F'The `size` parameter must contain the keys (height, width). Got {size.keys()}' )
					return center_crop(a_       ,       size=(size['''height'''], size['''width'''])       ,       data_format=a_       ,       **a_ )




			def  snake_case__     (   self							: Union[str, Any]       ,       a_							: np.ndarray       ,       a_							: Union[int, float]       ,       a_							: Optional[Union[str, ChannelDimension]] = None       ,       **a_							: List[str]       ,       ):
					'''simple docstring'''





					return rescale(a_       ,       scale=a_       ,       data_format=a_       ,       **a_ )




			def  snake_case__     (   self							: Optional[Any]       ,       a_							: np.ndarray       ,       a_							: Union[float, List[float]]       ,       a_							: Union[float, List[float]]       ,       a_							: Optional[Union[str, ChannelDimension]] = None       ,       **a_							: Dict       ,       ):
					'''simple docstring'''





					return normalize(a_       ,       mean=a_       ,       std=a_       ,       data_format=a_       ,       **a_ )




			def  snake_case__     (   self							: Any       ,       a_							: ImageInput       ,       a_							: bool = None       ,       a_							: Dict[str, int] = None       ,       a_							: PILImageResampling = None       ,       a_							: bool = None       ,       a_							: int = None       ,       a_							: bool = None       ,       a_							: float = None       ,       a_							: bool = None       ,       a_							: Optional[Union[float, List[float]]] = None       ,       a_							: Optional[Union[float, List[float]]] = None       ,       a_							: bool = None       ,       a_							: Optional[Union[str, TensorType]] = None       ,       a_							: Optional[ChannelDimension] = ChannelDimension.FIRST       ,       **a_							: Dict       ,       ):
					'''simple docstring'''





					__UpperCAmelCase    :     Tuple													=	do_resize if do_resize is not None else self.do_resize
					__UpperCAmelCase    :     Optional[Any]													=	size if size is not None else self.size
					__UpperCAmelCase    :     Dict													=	get_size_dict(a_       ,       param_name='''size'''       ,       default_to_square=a_ )
					__UpperCAmelCase    :     int													=	resample if resample is not None else self.resample
					__UpperCAmelCase    :     Optional[Any]													=	do_center_crop if do_center_crop is not None else self.do_center_crop
					__UpperCAmelCase    :     Any													=	crop_size if crop_size is not None else self.crop_size
					__UpperCAmelCase    :     Dict													=	get_size_dict(a_       ,       param_name='''crop_size'''       ,       default_to_square=a_ )
					__UpperCAmelCase    :     int													=	do_rescale if do_rescale is not None else self.do_rescale
					__UpperCAmelCase    :     Optional[int]													=	rescale_factor if rescale_factor is not None else self.rescale_factor
					__UpperCAmelCase    :     List[str]													=	do_normalize if do_normalize is not None else self.do_normalize
					__UpperCAmelCase    :     Optional[int]													=	image_mean if image_mean is not None else self.image_mean
					__UpperCAmelCase    :     Tuple													=	image_std if image_std is not None else self.image_std
					__UpperCAmelCase    :     List[str]													=	do_convert_rgb if do_convert_rgb is not None else self.do_convert_rgb

					__UpperCAmelCase    :     List[str]													=	make_list_of_images(a_ )

					if not valid_images(a_ ):
							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:
							raise ValueError('''Size 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.''' )

					# PIL RGBA images are converted to RGB
					if do_convert_rgb:
							__UpperCAmelCase    :     Optional[Any]													=	[convert_to_rgb(a_ ) for image in images]

					# All transformations expect numpy arrays.
					__UpperCAmelCase    :     int													=	[to_numpy_array(a_ ) for image in images]

					if do_resize:
							__UpperCAmelCase    :     int													=	[self.resize(image=a_       ,       size=a_       ,       resample=a_ ) for image in images]

					if do_center_crop:
							__UpperCAmelCase    :     List[str]													=	[self.center_crop(image=a_       ,       size=a_ ) for image in images]

					if do_rescale:
							__UpperCAmelCase    :     Dict													=	[self.rescale(image=a_       ,       scale=a_ ) for image in images]

					if do_normalize:
							__UpperCAmelCase    :     Optional[int]													=	[self.normalize(image=a_       ,       mean=a_       ,       std=a_ ) for image in images]

					__UpperCAmelCase    :     Optional[int]													=	[to_channel_dimension_format(a_       ,       a_ ) for image in images]

					__UpperCAmelCase    :     Union[str, Any]													=	{'''pixel_values''': images}
					return BatchFeature(data=a_       ,       tensor_type=a_ )





