



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						:





							'''simple docstring'''

							def __init__(							self ,					__UpperCAmelCase ,					):







											"""simple docstring"""



											a__	:							List[str]      						= parent
											a__	:							List[Any]      						= 13
											a__	:							Any      						= 7
											a__	:							str      						= True
											a__	:							Dict      						= True
											a__	:							Tuple      						= True
											a__	:							List[str]      						= 99
											a__	:							str      						= 32
											a__	:							str      						= 2
											a__	:							List[str]      						= 4
											a__	:							str      						= 37
											a__	:							Dict      						= "gelu"
											a__	:							Dict      						= 0.1
											a__	:							int      						= 0.1
											a__	:							str      						= 512
											a__	:							Dict      						= 16
											a__	:							List[str]      						= 2
											a__	:							Optional[int]      						= 0.0_2
											a__	:							List[str]      						= 3
											a__	:							str      						= 4
											a__	:							Optional[Any]      						= None



							def 		_A	(							self				):







											"""simple docstring"""



											a__	:							Tuple      						= ids_tensor([self.batch_size, self.seq_length] ,					self.vocab_size				)

											a__	:							Any      						= None
											if self.use_input_mask:
															a__	:							Dict      						= random_attention_mask([self.batch_size, self.seq_length]				)

											a__	:							Any      						= None
											a__	:							Dict      						= None
											a__	:							List[Any]      						= None
											if self.use_labels:
															a__	:							int      						= ids_tensor([self.batch_size] ,					self.type_sequence_label_size				)
															a__	:							Tuple      						= ids_tensor([self.batch_size, self.seq_length] ,					self.num_labels				)
															a__	:							List[str]      						= ids_tensor([self.batch_size] ,					self.num_choices				)

											a__	:							Union[str, 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 		_A	(							self				):







											"""simple docstring"""



											(
											    (
											    a__
											)						,   (
											    a__
											)						,   (
											    a__
											)						,   (
											    a__
											)						,   (
											    a__
											)						,   (
											    a__
											)						,   
											)	:							List[str]      						= self.prepare_config_and_inputs()

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







											"""simple docstring"""



											a__	:							Any      						= TFEsmModel(config=__UpperCAmelCase				)
											a__	:							str      						= {"input_ids": input_ids, "attention_mask": input_mask}
											a__	:							str      						= model(__UpperCAmelCase				)

											a__	:							Dict      						= [input_ids, input_mask]
											a__	:							List[str]      						= model(__UpperCAmelCase				)

											a__	:							Union[str, Any]      						= model(__UpperCAmelCase				)

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



							def 		_A	(							self ,					__UpperCAmelCase ,					__UpperCAmelCase ,					__UpperCAmelCase ,					__UpperCAmelCase ,					__UpperCAmelCase ,					__UpperCAmelCase ,					__UpperCAmelCase ,					__UpperCAmelCase ,					):







											"""simple docstring"""



											a__	:							Dict      						= True

											a__	:							Any      						= TFEsmModel(config=__UpperCAmelCase				)
											a__	:							Tuple      						= {
											    "input_ids": input_ids,
											    "attention_mask": input_mask,
											    "encoder_hidden_states": encoder_hidden_states,
											    "encoder_attention_mask": encoder_attention_mask,
											}
											a__	:							str      						= model(__UpperCAmelCase				)

											a__	:							Any      						= [input_ids, input_mask]
											a__	:							Optional[Any]      						= model(__UpperCAmelCase ,					encoder_hidden_states=__UpperCAmelCase				)

											# Also check the case where encoder outputs are not passed
											a__	:							Tuple      						= model(__UpperCAmelCase ,					attention_mask=__UpperCAmelCase				)

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



							def 		_A	(							self ,					__UpperCAmelCase ,					__UpperCAmelCase ,					__UpperCAmelCase ,					__UpperCAmelCase ,					__UpperCAmelCase ,					__UpperCAmelCase				):







											"""simple docstring"""



											a__	:							Dict      						= TFEsmForMaskedLM(config=__UpperCAmelCase				)
											a__	:							List[str]      						= model([input_ids, input_mask]				)
											self.parent.assertEqual(result.logits.shape ,					(self.batch_size, self.seq_length, self.vocab_size)				)



							def 		_A	(							self ,					__UpperCAmelCase ,					__UpperCAmelCase ,					__UpperCAmelCase ,					__UpperCAmelCase ,					__UpperCAmelCase ,					__UpperCAmelCase				):







											"""simple docstring"""



											a__	:							Union[str, Any]      						= self.num_labels
											a__	:							str      						= TFEsmForTokenClassification(config=__UpperCAmelCase				)
											a__	:							str      						= {"input_ids": input_ids, "attention_mask": input_mask}
											a__	:							Optional[Any]      						= model(__UpperCAmelCase				)
											self.parent.assertEqual(result.logits.shape ,					(self.batch_size, self.seq_length, self.num_labels)				)




							def 		_A	(							self				):







											"""simple docstring"""



											a__	:							List[str]      						= self.prepare_config_and_inputs()
											(
											    (
											    a__
											)						,   (
											    a__
											)						,   (
											    a__
											)						,   (
											    a__
											)						,   (
											    a__
											)						,   (
											    a__
											)						,   
											)	:							Dict      						= config_and_inputs
											a__	:							Union[str, Any]      						= {"input_ids": input_ids, "attention_mask": input_mask}
											return config, inputs_dict





@require_tf
class 				_a						(       SCREAMING_SNAKE_CASE       ,  SCREAMING_SNAKE_CASE       ,  unittest.TestCase			):





							'''simple docstring'''

							A				:List[str] =       (
							    (
							        TFEsmModel,
							        TFEsmForMaskedLM,
							        TFEsmForSequenceClassification,
							        TFEsmForTokenClassification,
							    )
							    if is_tf_available()
							    else ()
							)
							A				:List[Any] =       (
							    {
							        "feature-extraction": TFEsmModel,
							        "fill-mask": TFEsmForMaskedLM,
							        "text-classification": TFEsmForSequenceClassification,
							        "token-classification": TFEsmForTokenClassification,
							        "zero-shot": TFEsmForSequenceClassification,
							    }
							    if is_tf_available()
							    else {}
							)
							A				:List[str] =       False
							A				:List[Any] =       False



							def 		_A	(							self				):







											"""simple docstring"""



											a__	:							List[Any]      						= TFEsmModelTester(self				)
											a__	:							str      						= ConfigTester(self ,					config_class=__UpperCAmelCase ,					hidden_size=37				)



							def 		_A	(							self				):







											"""simple docstring"""



											self.config_tester.run_common_tests()



							def 		_A	(							self				):







											"""simple docstring"""



											a__	:							Dict      						= self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_model(*__UpperCAmelCase				)



							def 		_A	(							self				):







											"""simple docstring"""



											a__	:							List[str]      						= self.model_tester.prepare_config_and_inputs_for_decoder()
											self.model_tester.create_and_check_model_as_decoder(*__UpperCAmelCase				)



							def 		_A	(							self				):







											"""simple docstring"""



											a__	:							str      						= self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_for_masked_lm(*__UpperCAmelCase				)



							def 		_A	(							self				):







											"""simple docstring"""



											a__	:							Optional[int]      						= self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_for_token_classification(*__UpperCAmelCase				)



							@slow
							def 		_A	(							self				):







											"""simple docstring"""



											for model_name in TF_ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
															a__	:							List[str]      						= TFEsmModel.from_pretrained(__UpperCAmelCase				)
															self.assertIsNotNone(__UpperCAmelCase				)



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







											"""simple docstring"""



											pass



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







											"""simple docstring"""



											pass




							def 		_A	(							self				):







											"""simple docstring"""



											a__						,   a__	:							List[str]      						= self.model_tester.prepare_config_and_inputs_for_common()

											for model_class in self.all_model_classes:
															a__	:							Union[str, Any]      						= model_class(__UpperCAmelCase				)
															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__	:							Tuple      						= model.get_bias()
																			assert isinstance(__UpperCAmelCase ,					__UpperCAmelCase				)
																			for k, v in name.items():
																							assert isinstance(__UpperCAmelCase ,					tf.Variable				)
															else:
																			a__	:							Union[str, Any]      						= model.get_output_embeddings()
																			assert x is None
																			a__	:							Any      						= model.get_bias()
																			assert name is None





@require_tf
class 				_a						(       unittest.TestCase			):





							'''simple docstring'''

							@slow
							def 		_A	(							self				):







											"""simple docstring"""



											a__	:							List[str]      						= TFEsmForMaskedLM.from_pretrained("facebook/esm2_t6_8M_UR50D"				)

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







											"""simple docstring"""



											a__	:							List[Any]      						= TFEsmModel.from_pretrained("facebook/esm2_t6_8M_UR50D"				)

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







import operator as op


lowerCamelCase		       =      """scaler.pt"""
lowerCamelCase		       =      """pytorch_model"""
lowerCamelCase		       =      """random_states"""
lowerCamelCase		       =      """optimizer"""
lowerCamelCase		       =      """scheduler"""
lowerCamelCase		       =      """pytorch_model.bin"""
lowerCamelCase		       =      """pytorch_model.bin.index.json"""
lowerCamelCase		       =      """model.safetensors"""
lowerCamelCase		       =      """model.safetensors.index.json"""
lowerCamelCase		       =      """1.10.2"""
lowerCamelCase		       =      """py38"""
lowerCamelCase		       =      """4.17.0"""
lowerCamelCase		       =      ["""ml.p3.16xlarge""", """ml.p3dn.24xlarge""", """ml.p4dn.24xlarge"""]
lowerCamelCase		       =      ["""FULL_SHARD""", """SHARD_GRAD_OP""", """NO_SHARD""", """HYBRID_SHARD""", """HYBRID_SHARD_ZERO2"""]
lowerCamelCase		       =      ["""TRANSFORMER_BASED_WRAP""", """SIZE_BASED_WRAP""", """NO_WRAP"""]
lowerCamelCase		       =      ["""BACKWARD_PRE""", """BACKWARD_POST""", """NO_PREFETCH"""]
lowerCamelCase		       =      ["""FULL_STATE_DICT""", """LOCAL_STATE_DICT""", """SHARDED_STATE_DICT"""]
lowerCamelCase		       =      """2.0.1"""
lowerCamelCase		       =      ["""pdsh""", """standard""", """openmpi""", """mvapich"""]
lowerCamelCase		       =      ["""default""", """reduce-overhead""", """max-autotune"""]

lowerCamelCase		       =      {""">""": op.gt, """>=""": op.ge, """==""": op.eq, """!=""": op.ne, """<=""": op.le, """<""": op.lt}

# These are the args for `torch.distributed.launch` for pytorch < 1.9
lowerCamelCase		       =      [
    """nnodes""",
    """nproc_per_node""",
    """rdzv_backend""",
    """rdzv_endpoint""",
    """rdzv_id""",
    """rdzv_conf""",
    """standalone""",
    """max_restarts""",
    """monitor_interval""",
    """start_method""",
    """role""",
    """module""",
    """m""",
    """no_python""",
    """run_path""",
    """log_dir""",
    """r""",
    """redirects""",
    """t""",
    """tee""",
    """node_rank""",
    """master_addr""",
    """master_port""",
]

lowerCamelCase		       =      ["""DEEPSPEED""", """MULTI_GPU""", """FSDP""", """MEGATRON_LM"""]
lowerCamelCase		       =      ["""DEEPSPEED""", """MULTI_XPU""", """FSDP"""]


