






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




		'''simple docstring'''





		def __init__(      self   ,_lowerCAmelCase   ,):
								lowerCamelCase__         =  parent
								lowerCamelCase__         =  13
								lowerCamelCase__         =  7
								lowerCamelCase__         =  True
								lowerCamelCase__         =  True
								lowerCamelCase__         =  True
								lowerCamelCase__         =  99
								lowerCamelCase__         =  32
								lowerCamelCase__         =  2
								lowerCamelCase__         =  4
								lowerCamelCase__         =  37
								lowerCamelCase__         =  """gelu"""
								lowerCamelCase__         =  0.1
								lowerCamelCase__         =  0.1
								lowerCamelCase__         =  5_12
								lowerCamelCase__         =  16
								lowerCamelCase__         =  2
								lowerCamelCase__         =  0.02
								lowerCamelCase__         =  3
								lowerCamelCase__         =  4
								lowerCamelCase__         =  None






		def 			UpperCamelCase_					(      self       ):
								lowerCamelCase__         =  ids_tensor([self.batch_size, self.seq_length]   ,self.vocab_size       )

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

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

								lowerCamelCase__         =  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 			UpperCamelCase_					(      self       ):
								(
								    (
								    lowerCamelCase__
								)       ,			(
								    lowerCamelCase__
								)       ,			(
								    lowerCamelCase__
								)       ,			(
								    lowerCamelCase__
								)       ,			(
								    lowerCamelCase__
								)       ,			(
								    lowerCamelCase__
								)       ,			
								)         =  self.prepare_config_and_inputs()

								lowerCamelCase__         =  True
								lowerCamelCase__         =  floats_tensor([self.batch_size, self.seq_length, self.hidden_size]       )
								lowerCamelCase__         =  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 			UpperCamelCase_					(      self   ,_lowerCAmelCase   ,_lowerCAmelCase   ,_lowerCAmelCase   ,_lowerCAmelCase   ,_lowerCAmelCase   ,_lowerCAmelCase       ):
								lowerCamelCase__         =  TFEsmModel(config=_lowerCAmelCase       )
								lowerCamelCase__         =  {"""input_ids""": input_ids, """attention_mask""": input_mask}
								lowerCamelCase__         =  model(_lowerCAmelCase       )

								lowerCamelCase__         =  [input_ids, input_mask]
								lowerCamelCase__         =  model(_lowerCAmelCase       )

								lowerCamelCase__         =  model(_lowerCAmelCase       )

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






		def 			UpperCamelCase_					(      self   ,_lowerCAmelCase   ,_lowerCAmelCase   ,_lowerCAmelCase   ,_lowerCAmelCase   ,_lowerCAmelCase   ,_lowerCAmelCase   ,_lowerCAmelCase   ,_lowerCAmelCase   ,):
								lowerCamelCase__         =  True

								lowerCamelCase__         =  TFEsmModel(config=_lowerCAmelCase       )
								lowerCamelCase__         =  {
								    """input_ids""": input_ids,
								    """attention_mask""": input_mask,
								    """encoder_hidden_states""": encoder_hidden_states,
								    """encoder_attention_mask""": encoder_attention_mask,
								}
								lowerCamelCase__         =  model(_lowerCAmelCase       )

								lowerCamelCase__         =  [input_ids, input_mask]
								lowerCamelCase__         =  model(_lowerCAmelCase   ,encoder_hidden_states=_lowerCAmelCase       )

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

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






		def 			UpperCamelCase_					(      self   ,_lowerCAmelCase   ,_lowerCAmelCase   ,_lowerCAmelCase   ,_lowerCAmelCase   ,_lowerCAmelCase   ,_lowerCAmelCase       ):
								lowerCamelCase__         =  TFEsmForMaskedLM(config=_lowerCAmelCase       )
								lowerCamelCase__         =  model([input_ids, input_mask]       )
								self.parent.assertEqual(result.logits.shape   ,(self.batch_size, self.seq_length, self.vocab_size)       )






		def 			UpperCamelCase_					(      self   ,_lowerCAmelCase   ,_lowerCAmelCase   ,_lowerCAmelCase   ,_lowerCAmelCase   ,_lowerCAmelCase   ,_lowerCAmelCase       ):
								lowerCamelCase__         =  self.num_labels
								lowerCamelCase__         =  TFEsmForTokenClassification(config=_lowerCAmelCase       )
								lowerCamelCase__         =  {"""input_ids""": input_ids, """attention_mask""": input_mask}
								lowerCamelCase__         =  model(_lowerCAmelCase       )
								self.parent.assertEqual(result.logits.shape   ,(self.batch_size, self.seq_length, self.num_labels)       )






		def 			UpperCamelCase_					(      self       ):
								lowerCamelCase__         =  self.prepare_config_and_inputs()
								(
								    (
								    lowerCamelCase__
								)       ,			(
								    lowerCamelCase__
								)       ,			(
								    lowerCamelCase__
								)       ,			(
								    lowerCamelCase__
								)       ,			(
								    lowerCamelCase__
								)       ,			(
								    lowerCamelCase__
								)       ,			
								)         =  config_and_inputs
								lowerCamelCase__         =  {"""input_ids""": input_ids, """attention_mask""": input_mask}
								return config, inputs_dict




@require_tf
class 		UpperCamelCase__    (a   ,a   ,unittest.TestCase			):




		'''simple docstring'''





		_UpperCamelCase				  =					(
		    (
		        TFEsmModel,
		        TFEsmForMaskedLM,
		        TFEsmForSequenceClassification,
		        TFEsmForTokenClassification,
		    )
		    if is_tf_available()
		    else ()
		)
		_UpperCamelCase				  =					(
		    {
		        'feature-extraction': TFEsmModel,
		        'fill-mask': TFEsmForMaskedLM,
		        'text-classification': TFEsmForSequenceClassification,
		        'token-classification': TFEsmForTokenClassification,
		        'zero-shot': TFEsmForSequenceClassification,
		    }
		    if is_tf_available()
		    else {}
		)
		_UpperCamelCase				  =					False
		_UpperCamelCase				  =					False






		def 			UpperCamelCase_					(      self       ):
								lowerCamelCase__         =  TFEsmModelTester(self       )
								lowerCamelCase__         =  ConfigTester(self   ,config_class=_lowerCAmelCase   ,hidden_size=37       )






		def 			UpperCamelCase_					(      self       ):
								self.config_tester.run_common_tests()






		def 			UpperCamelCase_					(      self       ):
								lowerCamelCase__         =  self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_model(*_lowerCAmelCase       )






		def 			UpperCamelCase_					(      self       ):
								lowerCamelCase__         =  self.model_tester.prepare_config_and_inputs_for_decoder()
								self.model_tester.create_and_check_model_as_decoder(*_lowerCAmelCase       )






		def 			UpperCamelCase_					(      self       ):
								lowerCamelCase__         =  self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_masked_lm(*_lowerCAmelCase       )






		def 			UpperCamelCase_					(      self       ):
								lowerCamelCase__         =  self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_token_classification(*_lowerCAmelCase       )






		@slow
		def 			UpperCamelCase_					(      self       ):
								for model_name in TF_ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
														lowerCamelCase__         =  TFEsmModel.from_pretrained(_lowerCAmelCase       )
														self.assertIsNotNone(_lowerCAmelCase       )






		@unittest.skip("""Protein models do not support embedding resizing."""       )
		def 			UpperCamelCase_					(      self       ):
								pass






		@unittest.skip("""Protein models do not support embedding resizing."""       )
		def 			UpperCamelCase_					(      self       ):
								pass






		def 			UpperCamelCase_					(      self       ):
								lowerCamelCase__       ,			lowerCamelCase__         =  self.model_tester.prepare_config_and_inputs_for_common()

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




@require_tf
class 		UpperCamelCase__    (unittest.TestCase			):




		'''simple docstring'''





		@slow
		def 			UpperCamelCase_					(      self       ):
								lowerCamelCase__         =  TFEsmForMaskedLM.from_pretrained("""facebook/esm2_t6_8M_UR50D"""       )

								lowerCamelCase__         =  tf.constant([[0, 1, 2, 3, 4, 5]]       )
								lowerCamelCase__         =  model(_lowerCAmelCase       )[0]
								lowerCamelCase__         =  [1, 6, 33]
								self.assertEqual(list(output.numpy().shape       )   ,_lowerCAmelCase       )
								# compare the actual values for a slice.
								lowerCamelCase__         =  tf.constant(
								    [
								        [
								            [8.92_1518, -10.58_9814, -6.467_1307],
								            [-6.396_7156, -13.91_1377, -1.121_1915],
								            [-7.78_1247, -13.95_1557, -3.74_0592],
								        ]
								    ]       )
								self.assertTrue(numpy.allclose(output[:, :3, :3].numpy()   ,expected_slice.numpy()   ,atol=1E-2       )       )






		@slow
		def 			UpperCamelCase_					(      self       ):
								lowerCamelCase__         =  TFEsmModel.from_pretrained("""facebook/esm2_t6_8M_UR50D"""       )

								lowerCamelCase__         =  tf.constant([[0, 6, 4, 13, 5, 4, 16, 12, 11, 7, 2]]       )
								lowerCamelCase__         =  model(_lowerCAmelCase       )[0]
								# compare the actual values for a slice.
								lowerCamelCase__         =  tf.constant(
								    [
								        [
								            [0.1444_3092, 0.5412_5327, 0.324_7739],
								            [0.3034_0484, 0.0052_6676, 0.3107_7722],
								            [0.3227_8043, -0.2498_7096, 0.341_4628],
								        ]
								    ]       )
								self.assertTrue(numpy.allclose(output[:, :3, :3].numpy()   ,expected_slice.numpy()   ,atol=1E-4       )       )









'''simple docstring'''







def 						A__			(				):
						return [
						    a * b * (1000 - a - b)
						    for a in range(1   ,    999						)
						    for b in range(__lowerCAmelCase   ,    999						)
						    if (a * a + b * b == (1000 - a - b) ** 2)
						][0]


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

