


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 	snake_case_       :
					def __init__(				self       :str      ,__snake_case       :List[str]      ,) ->   int:
												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.02
												a__					=					3
												a__					=					4
												a__					=					None
					def 			lowerCamelCase__(				self       :Dict				) ->   int:
												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       :Tuple				) ->   List[str]:
												(
												    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       :Any      ,__snake_case       :Optional[Any]      ,__snake_case       :Optional[Any]      ,__snake_case       :Union[str, Any]      ,__snake_case       :int      ,__snake_case       :Optional[int]      ,__snake_case       :Optional[int]				) ->   Any:
												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       :Optional[Any]      ,__snake_case       :Any      ,__snake_case       :List[Any]      ,__snake_case       :Optional[int]      ,__snake_case       :str      ,__snake_case       :Union[str, Any]      ,__snake_case       :Tuple      ,__snake_case       :Union[str, Any]      ,__snake_case       :str      ,) ->   Optional[int]:
												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       :Tuple      ,__snake_case       :Optional[Any]      ,__snake_case       :int      ,__snake_case       :Any      ,__snake_case       :List[Any]      ,__snake_case       :Tuple      ,__snake_case       :Dict				) ->   Union[str, Any]:
												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       :Optional[int]      ,__snake_case       :List[Any]      ,__snake_case       :Optional[Any]      ,__snake_case       :Tuple      ,__snake_case       :Union[str, Any]      ,__snake_case       :str      ,__snake_case       :Optional[Any]				) ->   Optional[int]:
												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       :List[Any]				) ->   Any:
												a__					=					self.prepare_config_and_inputs()
												(
												    a__    
												)					=					config_and_inputs
												a__					=					{'''input_ids''': input_ids, '''attention_mask''': input_mask}
												return config, inputs_dict







@require_tf
class 	snake_case_       (a__     ,						a__     ,						unittest.TestCase				):
					UpperCAmelCase__ :   str   				=      (
					    (
					        TFEsmModel,
					        TFEsmForMaskedLM,
					        TFEsmForSequenceClassification,
					        TFEsmForTokenClassification,
					    )
					    if is_tf_available()
					    else ()
					)
					UpperCAmelCase__ :   Any   				=      (
					    {
					        'feature-extraction': TFEsmModel,
					        'fill-mask': TFEsmForMaskedLM,
					        'text-classification': TFEsmForSequenceClassification,
					        'token-classification': TFEsmForTokenClassification,
					        'zero-shot': TFEsmForSequenceClassification,
					    }
					    if is_tf_available()
					    else {}
					)
					UpperCAmelCase__ :   Tuple   				=      False
					UpperCAmelCase__ :   Dict   				=      False
					def 			lowerCamelCase__(				self       :Optional[Any]				) ->   Tuple:
												a__					=					TFEsmModelTester(self				)
												a__					=					ConfigTester(self      ,config_class=_lowerCamelCase      ,hidden_size=37				)
					def 			lowerCamelCase__(				self       :Any				) ->   Optional[int]:
												self.config_tester.run_common_tests()
					def 			lowerCamelCase__(				self       :Any				) ->   Union[str, Any]:
												a__					=					self.model_tester.prepare_config_and_inputs()
												self.model_tester.create_and_check_model(*_lowerCamelCase				)
					def 			lowerCamelCase__(				self       :Union[str, Any]				) ->   int:
												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]				) ->   str:
												a__					=					self.model_tester.prepare_config_and_inputs()
												self.model_tester.create_and_check_for_masked_lm(*_lowerCamelCase				)
					def 			lowerCamelCase__(				self       :List[str]				) ->   Optional[Any]:
												a__					=					self.model_tester.prepare_config_and_inputs()
												self.model_tester.create_and_check_for_token_classification(*_lowerCamelCase				)
					@slow
					def 			lowerCamelCase__(				self       :List[str]				) ->   str:
												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       :Union[str, Any]				) ->   Tuple:
												pass
					@unittest.skip('Protein models do not support embedding resizing.'				)
					def 			lowerCamelCase__(				self       :List[Any]				) ->   Union[str, Any]:
												pass



					def 			lowerCamelCase__(				self       :List[Any]				) ->   int:
												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 	snake_case_       (unittest.TestCase				):
					@slow
					def 			lowerCamelCase__(				self       :Dict				) ->   Optional[int]:
												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.92_15_18, -10.58_98_14, -6.4_67_13_07],
												            [-6.3_96_71_56, -13.91_13_77, -1.1_21_19_15],
												            [-7.78_12_47, -13.95_15_57, -3.74_05_92],
												        ]
												    ]				)
												self.assertTrue(numpy.allclose(output[:, :3, :3].numpy()      ,expected_slice.numpy()      ,atol=1E-2				)				)



					@slow
					def 			lowerCamelCase__(				self       :Dict				) ->   Dict:
												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.14_44_30_92, 0.54_12_53_27, 0.3_24_77_39],
												            [0.30_34_04_84, 0.00_52_66_76, 0.31_07_77_22],
												            [0.32_27_80_43, -0.24_98_70_96, 0.3_41_46_28],
												        ]
												    ]				)
												self.assertTrue(numpy.allclose(output[:, :3, :3].numpy()      ,expected_slice.numpy()      ,atol=1E-4				)				)










from __future__ import annotations






def 							__lowercase       (			__lowerCAmelCase							:    float ,   __lowerCAmelCase							:    float ,   __lowerCAmelCase							:    float       ):
							if (voltage, current, resistance).count(0       ) != 1:
														raise ValueError('One and only one argument must be 0'       )
							if resistance < 0:
														raise ValueError('Resistance cannot be negative'       )
							if voltage == 0:
														return {"voltage": float(current * resistance       )}
							elif current == 0:
														return {"current": voltage / resistance}
							elif resistance == 0:
														return {"resistance": voltage / current}
							else:
														raise ValueError('Exactly one argument must be 0'       )


if __name__ == "__main__":
						import doctest

						doctest.testmod()




