






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




							def __init__(   self     :   Union[str, Any] ,							SCREAMING_SNAKE_CASE_     :   List[Any] ,							):
											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__      			=						512
											lowerCAmelCase__      			=						16
											lowerCAmelCase__      			=						2
											lowerCAmelCase__      			=						0.02
											lowerCAmelCase__      			=						3
											lowerCAmelCase__      			=						4
											lowerCAmelCase__      			=						None




							def 						__snake_case					(   self     :   List[str]	):
											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 						__snake_case					(   self     :   Dict	):
											(
											    (
											    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 						__snake_case					(   self     :   Dict ,							SCREAMING_SNAKE_CASE_     :   List[str] ,							SCREAMING_SNAKE_CASE_     :   Tuple ,							SCREAMING_SNAKE_CASE_     :   Tuple ,							SCREAMING_SNAKE_CASE_     :   Union[str, Any] ,							SCREAMING_SNAKE_CASE_     :   Dict ,							SCREAMING_SNAKE_CASE_     :   Any	):
											lowerCAmelCase__      			=						TFEsmModel(config=SCREAMING_SNAKE_CASE_	)
											lowerCAmelCase__      			=						{'''input_ids''': input_ids, '''attention_mask''': input_mask}
											lowerCAmelCase__      			=						model(SCREAMING_SNAKE_CASE_	)

											lowerCAmelCase__      			=						[input_ids, input_mask]
											lowerCAmelCase__      			=						model(SCREAMING_SNAKE_CASE_	)

											lowerCAmelCase__      			=						model(SCREAMING_SNAKE_CASE_	)

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




							def 						__snake_case					(   self     :   Union[str, Any] ,							SCREAMING_SNAKE_CASE_     :   int ,							SCREAMING_SNAKE_CASE_     :   List[str] ,							SCREAMING_SNAKE_CASE_     :   List[Any] ,							SCREAMING_SNAKE_CASE_     :   Tuple ,							SCREAMING_SNAKE_CASE_     :   int ,							SCREAMING_SNAKE_CASE_     :   Optional[int] ,							SCREAMING_SNAKE_CASE_     :   List[Any] ,							SCREAMING_SNAKE_CASE_     :   Optional[int] ,							):
											lowerCAmelCase__      			=						True

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

											lowerCAmelCase__      			=						[input_ids, input_mask]
											lowerCAmelCase__      			=						model(SCREAMING_SNAKE_CASE_ ,							encoder_hidden_states=SCREAMING_SNAKE_CASE_	)

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

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




							def 						__snake_case					(   self     :   Any ,							SCREAMING_SNAKE_CASE_     :   Optional[int] ,							SCREAMING_SNAKE_CASE_     :   Union[str, Any] ,							SCREAMING_SNAKE_CASE_     :   Any ,							SCREAMING_SNAKE_CASE_     :   Union[str, Any] ,							SCREAMING_SNAKE_CASE_     :   Tuple ,							SCREAMING_SNAKE_CASE_     :   Tuple	):
											lowerCAmelCase__      			=						TFEsmForMaskedLM(config=SCREAMING_SNAKE_CASE_	)
											lowerCAmelCase__      			=						model([input_ids, input_mask]	)
											self.parent.assertEqual(result.logits.shape ,							(self.batch_size, self.seq_length, self.vocab_size)	)




							def 						__snake_case					(   self     :   Union[str, Any] ,							SCREAMING_SNAKE_CASE_     :   Dict ,							SCREAMING_SNAKE_CASE_     :   Dict ,							SCREAMING_SNAKE_CASE_     :   List[str] ,							SCREAMING_SNAKE_CASE_     :   Optional[int] ,							SCREAMING_SNAKE_CASE_     :   str ,							SCREAMING_SNAKE_CASE_     :   Optional[Any]	):
											lowerCAmelCase__      			=						self.num_labels
											lowerCAmelCase__      			=						TFEsmForTokenClassification(config=SCREAMING_SNAKE_CASE_	)
											lowerCAmelCase__      			=						{'''input_ids''': input_ids, '''attention_mask''': input_mask}
											lowerCAmelCase__      			=						model(SCREAMING_SNAKE_CASE_	)
											self.parent.assertEqual(result.logits.shape ,							(self.batch_size, self.seq_length, self.num_labels)	)




							def 						__snake_case					(   self     :   List[Any]	):
											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  lowerCAmelCase_       (				snake_case__     ,    snake_case__     ,    unittest.TestCase					):
							UpperCamelCase_		:Optional[int]		      =    (
							    (
							        TFEsmModel,
							        TFEsmForMaskedLM,
							        TFEsmForSequenceClassification,
							        TFEsmForTokenClassification,
							    )
							    if is_tf_available()
							    else ()
							)
							UpperCamelCase_		:List[Any]		      =    (
							    {
							        'feature-extraction': TFEsmModel,
							        'fill-mask': TFEsmForMaskedLM,
							        'text-classification': TFEsmForSequenceClassification,
							        'token-classification': TFEsmForTokenClassification,
							        'zero-shot': TFEsmForSequenceClassification,
							    }
							    if is_tf_available()
							    else {}
							)
							UpperCamelCase_		:Union[str, Any]		      =    False
							UpperCamelCase_		:Union[str, Any]		      =    False




							def 						__snake_case					(   self     :   Union[str, Any]	):
											lowerCAmelCase__      			=						TFEsmModelTester(self	)
											lowerCAmelCase__      			=						ConfigTester(self ,							config_class=SCREAMING_SNAKE_CASE_ ,							hidden_size=37	)




							def 						__snake_case					(   self     :   Optional[Any]	):
											self.config_tester.run_common_tests()




							def 						__snake_case					(   self     :   Dict	):
											lowerCAmelCase__      			=						self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_model(*SCREAMING_SNAKE_CASE_	)




							def 						__snake_case					(   self     :   Optional[Any]	):
											lowerCAmelCase__      			=						self.model_tester.prepare_config_and_inputs_for_decoder()
											self.model_tester.create_and_check_model_as_decoder(*SCREAMING_SNAKE_CASE_	)




							def 						__snake_case					(   self     :   str	):
											lowerCAmelCase__      			=						self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_for_masked_lm(*SCREAMING_SNAKE_CASE_	)




							def 						__snake_case					(   self     :   Union[str, Any]	):
											lowerCAmelCase__      			=						self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_for_token_classification(*SCREAMING_SNAKE_CASE_	)




							@slow
							def 						__snake_case					(   self     :   Optional[Any]	):
											for model_name in TF_ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
															lowerCAmelCase__      			=						TFEsmModel.from_pretrained(SCREAMING_SNAKE_CASE_	)
															self.assertIsNotNone(SCREAMING_SNAKE_CASE_	)




							@unittest.skip('''Protein models do not support embedding resizing.'''	)
							def 						__snake_case					(   self     :   Any	):
											pass




							@unittest.skip('''Protein models do not support embedding resizing.'''	)
							def 						__snake_case					(   self     :   Dict	):
											pass




							def 						__snake_case					(   self     :   Union[str, Any]	):
											lowerCAmelCase__			,		lowerCAmelCase__      			=						self.model_tester.prepare_config_and_inputs_for_common()

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



@require_tf
class  lowerCAmelCase_       (				unittest.TestCase					):




							@slow
							def 						__snake_case					(   self     :   Optional[int]	):
											lowerCAmelCase__      			=						TFEsmForMaskedLM.from_pretrained('''facebook/esm2_t6_8M_UR50D'''	)

											lowerCAmelCase__      			=						tf.constant([[0, 1, 2, 3, 4, 5]]	)
											lowerCAmelCase__      			=						model(SCREAMING_SNAKE_CASE_	)[0]
											lowerCAmelCase__      			=						[1, 6, 33]
											self.assertEqual(list(output.numpy().shape	) ,							SCREAMING_SNAKE_CASE_	)
											# compare the actual values for a slice.
											lowerCAmelCase__      			=						tf.constant(
											    [
											        [
											            [8.921_518, -10.589_814, -6.4_671_307],
											            [-6.3_967_156, -13.911_377, -1.1_211_915],
											            [-7.781_247, -13.951_557, -3.740_592],
											        ]
											    ]	)
											self.assertTrue(numpy.allclose(output[:, :3, :3].numpy() ,							expected_slice.numpy() ,							atol=1e-2	)	)




							@slow
							def 						__snake_case					(   self     :   int	):
											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(SCREAMING_SNAKE_CASE_	)[0]
											# compare the actual values for a slice.
											lowerCAmelCase__      			=						tf.constant(
											    [
											        [
											            [0.14_443_092, 0.54_125_327, 0.3_247_739],
											            [0.30_340_484, 0.00_526_676, 0.31_077_722],
											            [0.32_278_043, -0.24_987_096, 0.3_414_628],
											        ]
											    ]	)
											self.assertTrue(numpy.allclose(output[:, :3, :3].numpy() ,							expected_slice.numpy() ,							atol=1e-4	)	)













from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_flax_available,
    is_tf_available,
    is_tokenizers_available,
    is_torch_available,
)


_UpperCAmelCase			:		List[str]   				=    {
    "configuration_electra": ["ELECTRA_PRETRAINED_CONFIG_ARCHIVE_MAP", "ElectraConfig", "ElectraOnnxConfig"],
    "tokenization_electra": ["ElectraTokenizer"],
}

try:
			if not is_tokenizers_available():
						raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
			pass
else:
			_UpperCAmelCase			:		Dict   				=    ["ElectraTokenizerFast"]

try:
			if not is_torch_available():
						raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
			pass
else:
			_UpperCAmelCase			:		Optional[Any]   				=    [
			    "ELECTRA_PRETRAINED_MODEL_ARCHIVE_LIST",
			    "ElectraForCausalLM",
			    "ElectraForMaskedLM",
			    "ElectraForMultipleChoice",
			    "ElectraForPreTraining",
			    "ElectraForQuestionAnswering",
			    "ElectraForSequenceClassification",
			    "ElectraForTokenClassification",
			    "ElectraModel",
			    "ElectraPreTrainedModel",
			    "load_tf_weights_in_electra",
			]

try:
			if not is_tf_available():
						raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
			pass
else:
			_UpperCAmelCase			:		Union[str, Any]   				=    [
			    "TF_ELECTRA_PRETRAINED_MODEL_ARCHIVE_LIST",
			    "TFElectraForMaskedLM",
			    "TFElectraForMultipleChoice",
			    "TFElectraForPreTraining",
			    "TFElectraForQuestionAnswering",
			    "TFElectraForSequenceClassification",
			    "TFElectraForTokenClassification",
			    "TFElectraModel",
			    "TFElectraPreTrainedModel",
			]

try:
			if not is_flax_available():
						raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
			pass
else:
			_UpperCAmelCase			:		Optional[int]   				=    [
			    "FlaxElectraForCausalLM",
			    "FlaxElectraForMaskedLM",
			    "FlaxElectraForMultipleChoice",
			    "FlaxElectraForPreTraining",
			    "FlaxElectraForQuestionAnswering",
			    "FlaxElectraForSequenceClassification",
			    "FlaxElectraForTokenClassification",
			    "FlaxElectraModel",
			    "FlaxElectraPreTrainedModel",
			]


if TYPE_CHECKING:
			from .configuration_electra import ELECTRA_PRETRAINED_CONFIG_ARCHIVE_MAP, ElectraConfig, ElectraOnnxConfig
			from .tokenization_electra import ElectraTokenizer

			try:
						if not is_tokenizers_available():
									raise OptionalDependencyNotAvailable()
			except OptionalDependencyNotAvailable:
						pass
			else:
						from .tokenization_electra_fast import ElectraTokenizerFast

			try:
						if not is_torch_available():
									raise OptionalDependencyNotAvailable()
			except OptionalDependencyNotAvailable:
						pass
			else:
						from .modeling_electra import (
						    ELECTRA_PRETRAINED_MODEL_ARCHIVE_LIST,
						    ElectraForCausalLM,
						    ElectraForMaskedLM,
						    ElectraForMultipleChoice,
						    ElectraForPreTraining,
						    ElectraForQuestionAnswering,
						    ElectraForSequenceClassification,
						    ElectraForTokenClassification,
						    ElectraModel,
						    ElectraPreTrainedModel,
						    load_tf_weights_in_electra,
						)

			try:
						if not is_tf_available():
									raise OptionalDependencyNotAvailable()
			except OptionalDependencyNotAvailable:
						pass
			else:
						from .modeling_tf_electra import (
						    TF_ELECTRA_PRETRAINED_MODEL_ARCHIVE_LIST,
						    TFElectraForMaskedLM,
						    TFElectraForMultipleChoice,
						    TFElectraForPreTraining,
						    TFElectraForQuestionAnswering,
						    TFElectraForSequenceClassification,
						    TFElectraForTokenClassification,
						    TFElectraModel,
						    TFElectraPreTrainedModel,
						)

			try:
						if not is_flax_available():
									raise OptionalDependencyNotAvailable()
			except OptionalDependencyNotAvailable:
						pass
			else:
						from .modeling_flax_electra import (
						    FlaxElectraForCausalLM,
						    FlaxElectraForMaskedLM,
						    FlaxElectraForMultipleChoice,
						    FlaxElectraForPreTraining,
						    FlaxElectraForQuestionAnswering,
						    FlaxElectraForSequenceClassification,
						    FlaxElectraForTokenClassification,
						    FlaxElectraModel,
						    FlaxElectraPreTrainedModel,
						)

else:
			import sys

			_UpperCAmelCase			:		Any   				=    _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)





