




'''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      A       :







							def __init__(    self			:    Optional[int]				,     lowerCAmelCase_			:    Tuple				,     )     ->				Tuple:


														"""simple docstring"""


														_a      	=							parent
														_a      	=							13
														_a      	=							7
														_a      	=							True
														_a      	=							True
														_a      	=							True
														_a      	=							99
														_a      	=							32
														_a      	=							2
														_a      	=							4
														_a      	=							37
														_a      	=							'''gelu'''
														_a      	=							0.1
														_a      	=							0.1
														_a      	=							5_12
														_a      	=							16
														_a      	=							2
														_a      	=							0.0_2
														_a      	=							3
														_a      	=							4
														_a      	=							None







							def   __lowerCAmelCase     (    self			:    int   )     ->				Any:


														"""simple docstring"""


														_a      	=							ids_tensor([self.batch_size, self.seq_length]				,     self.vocab_size   )

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

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

														_a      	=							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			:    Any   )     ->				Dict:


														"""simple docstring"""


														(
														    (
														    _a
														)	,				(
														    _a
														)	,				(
														    _a
														)	,				(
														    _a
														)	,				(
														    _a
														)	,				(
														    _a
														)	,				
														)      	=							self.prepare_config_and_inputs()

														_a      	=							True
														_a      	=							floats_tensor([self.batch_size, self.seq_length, self.hidden_size]   )
														_a      	=							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			:    int				,     lowerCAmelCase_			:    Any				,     lowerCAmelCase_			:    List[str]				,     lowerCAmelCase_			:    Optional[int]				,     lowerCAmelCase_			:    Dict				,     lowerCAmelCase_			:    List[Any]				,     lowerCAmelCase_			:    int   )     ->				Optional[Any]:


														"""simple docstring"""


														_a      	=							TFEsmModel(config=lowerCAmelCase_   )
														_a      	=							{'''input_ids''': input_ids, '''attention_mask''': input_mask}
														_a      	=							model(lowerCAmelCase_   )

														_a      	=							[input_ids, input_mask]
														_a      	=							model(lowerCAmelCase_   )

														_a      	=							model(lowerCAmelCase_   )

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







							def   __lowerCAmelCase     (    self			:    List[str]				,     lowerCAmelCase_			:    Optional[Any]				,     lowerCAmelCase_			:    List[Any]				,     lowerCAmelCase_			:    Tuple				,     lowerCAmelCase_			:    List[Any]				,     lowerCAmelCase_			:    Any				,     lowerCAmelCase_			:    List[str]				,     lowerCAmelCase_			:    Dict				,     lowerCAmelCase_			:    str				,     )     ->				Optional[int]:


														"""simple docstring"""


														_a      	=							True

														_a      	=							TFEsmModel(config=lowerCAmelCase_   )
														_a      	=							{
														    '''input_ids''': input_ids,
														    '''attention_mask''': input_mask,
														    '''encoder_hidden_states''': encoder_hidden_states,
														    '''encoder_attention_mask''': encoder_attention_mask,
														}
														_a      	=							model(lowerCAmelCase_   )

														_a      	=							[input_ids, input_mask]
														_a      	=							model(lowerCAmelCase_				,     encoder_hidden_states=lowerCAmelCase_   )

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

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







							def   __lowerCAmelCase     (    self			:    Optional[Any]				,     lowerCAmelCase_			:    Dict				,     lowerCAmelCase_			:    Union[str, Any]				,     lowerCAmelCase_			:    Optional[Any]				,     lowerCAmelCase_			:    List[str]				,     lowerCAmelCase_			:    List[Any]				,     lowerCAmelCase_			:    str   )     ->				List[Any]:


														"""simple docstring"""


														_a      	=							TFEsmForMaskedLM(config=lowerCAmelCase_   )
														_a      	=							model([input_ids, input_mask]   )
														self.parent.assertEqual(result.logits.shape				,     (self.batch_size, self.seq_length, self.vocab_size)   )







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


														"""simple docstring"""


														_a      	=							self.num_labels
														_a      	=							TFEsmForTokenClassification(config=lowerCAmelCase_   )
														_a      	=							{'''input_ids''': input_ids, '''attention_mask''': input_mask}
														_a      	=							model(lowerCAmelCase_   )
														self.parent.assertEqual(result.logits.shape				,     (self.batch_size, self.seq_length, self.num_labels)   )







							def   __lowerCAmelCase     (    self			:    int   )     ->				Tuple:


														"""simple docstring"""


														_a      	=							self.prepare_config_and_inputs()
														(
														    (
														    _a
														)	,				(
														    _a
														)	,				(
														    _a
														)	,				(
														    _a
														)	,				(
														    _a
														)	,				(
														    _a
														)	,				
														)      	=							config_and_inputs
														_a      	=							{'''input_ids''': input_ids, '''attention_mask''': input_mask}
														return config, inputs_dict



@require_tf
class      A       (						_a	,_a	,unittest.TestCase       ):
							lowercase_							     =							(
							    (
							        TFEsmModel,
							        TFEsmForMaskedLM,
							        TFEsmForSequenceClassification,
							        TFEsmForTokenClassification,
							    )
							    if is_tf_available()
							    else ()
							)
							lowercase_							     =							(
							    {
							        'feature-extraction': TFEsmModel,
							        'fill-mask': TFEsmForMaskedLM,
							        'text-classification': TFEsmForSequenceClassification,
							        'token-classification': TFEsmForTokenClassification,
							        'zero-shot': TFEsmForSequenceClassification,
							    }
							    if is_tf_available()
							    else {}
							)
							lowercase_							     =							False
							lowercase_							     =							False







							def   __lowerCAmelCase     (    self			:    List[str]   )     ->				str:


														"""simple docstring"""


														_a      	=							TFEsmModelTester(self   )
														_a      	=							ConfigTester(self				,     config_class=lowerCAmelCase_				,     hidden_size=37   )







							def   __lowerCAmelCase     (    self			:    int   )     ->				Union[str, Any]:


														"""simple docstring"""


														self.config_tester.run_common_tests()







							def   __lowerCAmelCase     (    self			:    str   )     ->				str:


														"""simple docstring"""


														_a      	=							self.model_tester.prepare_config_and_inputs()
														self.model_tester.create_and_check_model(*lowerCAmelCase_   )







							def   __lowerCAmelCase     (    self			:    Optional[Any]   )     ->				Union[str, Any]:


														"""simple docstring"""


														_a      	=							self.model_tester.prepare_config_and_inputs_for_decoder()
														self.model_tester.create_and_check_model_as_decoder(*lowerCAmelCase_   )







							def   __lowerCAmelCase     (    self			:    Optional[Any]   )     ->				Union[str, Any]:


														"""simple docstring"""


														_a      	=							self.model_tester.prepare_config_and_inputs()
														self.model_tester.create_and_check_for_masked_lm(*lowerCAmelCase_   )







							def   __lowerCAmelCase     (    self			:    Optional[Any]   )     ->				Optional[int]:


														"""simple docstring"""


														_a      	=							self.model_tester.prepare_config_and_inputs()
														self.model_tester.create_and_check_for_token_classification(*lowerCAmelCase_   )







							@slow
							def   __lowerCAmelCase     (    self			:    List[str]   )     ->				Dict:


														"""simple docstring"""


														for model_name in TF_ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																					_a      	=							TFEsmModel.from_pretrained(lowerCAmelCase_   )
																					self.assertIsNotNone(lowerCAmelCase_   )







							@unittest.skip('''Protein models do not support embedding resizing.'''   )
							def   __lowerCAmelCase     (    self			:    List[Any]   )     ->				Tuple:


														"""simple docstring"""


														pass







							@unittest.skip('''Protein models do not support embedding resizing.'''   )
							def   __lowerCAmelCase     (    self			:    int   )     ->				Tuple:


														"""simple docstring"""


														pass







							def   __lowerCAmelCase     (    self			:    Union[str, Any]   )     ->				str:


														"""simple docstring"""


														_a	,				_a      	=							self.model_tester.prepare_config_and_inputs_for_common()

														for model_class in self.all_model_classes:
																					_a      	=							model_class(lowerCAmelCase_   )
																					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
																												_a      	=							model.get_bias()
																												assert isinstance(lowerCAmelCase_				,     lowerCAmelCase_   )
																												for k, v in name.items():
																																			assert isinstance(lowerCAmelCase_				,     tf.Variable   )
																					else:
																												_a      	=							model.get_output_embeddings()
																												assert x is None
																												_a      	=							model.get_bias()
																												assert name is None



@require_tf
class      A       (						unittest.TestCase       ):







							@slow
							def   __lowerCAmelCase     (    self			:    Optional[int]   )     ->				Dict:


														"""simple docstring"""


														_a      	=							TFEsmForMaskedLM.from_pretrained('''facebook/esm2_t6_8M_UR50D'''   )

														_a      	=							tf.constant([[0, 1, 2, 3, 4, 5]]   )
														_a      	=							model(lowerCAmelCase_   )[0]
														_a      	=							[1, 6, 33]
														self.assertEqual(list(output.numpy().shape   )				,     lowerCAmelCase_   )
														# compare the actual values for a slice.
														_a      	=							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			:    Any   )     ->				List[Any]:


														"""simple docstring"""


														_a      	=							TFEsmModel.from_pretrained('''facebook/esm2_t6_8M_UR50D'''   )

														_a      	=							tf.constant([[0, 6, 4, 13, 5, 4, 16, 12, 11, 7, 2]]   )
														_a      	=							model(lowerCAmelCase_   )[0]
														# compare the actual values for a slice.
														_a      	=							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   )   )





def 						_lowerCAmelCase							(     _lowerCAmelCase = 100			)    ->						int:


						'''simple docstring'''






						__snake_case				=		n * (n + 1) * (2 * n + 1) / 6
						__snake_case				=		(n * (n + 1) / 2) ** 2
						return int(square_of_sum - sum_of_squares			)


if __name__ == "__main__":
	print(f'''{solution() = }''')





