





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







	def __init__(				self			:   str , lowerCamelCase			:   Any , ):




				'''simple docstring'''
				a__     = parent
				a__     = 1_3
				a__     = 7
				a__     = True
				a__     = True
				a__     = True
				a__     = 9_9
				a__     = 3_2
				a__     = 2
				a__     = 4
				a__     = 3_7
				a__     = "gelu"
				a__     = 0.1
				a__     = 0.1
				a__     = 5_1_2
				a__     = 1_6
				a__     = 2
				a__     = 0.02
				a__     = 3
				a__     = 4
				a__     = None







	def 						__a   (				self			:   Tuple		):




				'''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 						__a   (				self			:   Any		):




				'''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 						__a   (				self			:   List[Any] , lowerCamelCase			:   Optional[Any] , lowerCamelCase			:   int , lowerCamelCase			:   List[Any] , lowerCamelCase			:   Optional[Any] , lowerCamelCase			:   List[str] , lowerCamelCase			:   Dict		):




				'''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 						__a   (				self			:   Union[str, Any] , lowerCamelCase			:   Dict , lowerCamelCase			:   Any , lowerCamelCase			:   Any , lowerCamelCase			:   Dict , lowerCamelCase			:   Dict , lowerCamelCase			:   int , lowerCamelCase			:   List[Any] , lowerCamelCase			:   Tuple , ):




				'''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 						__a   (				self			:   Optional[int] , lowerCamelCase			:   Dict , lowerCamelCase			:   Union[str, Any] , lowerCamelCase			:   Optional[int] , lowerCamelCase			:   Dict , lowerCamelCase			:   Tuple , lowerCamelCase			:   Optional[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 						__a   (				self			:   Tuple , lowerCamelCase			:   Union[str, Any] , lowerCamelCase			:   Union[str, Any] , lowerCamelCase			:   str , lowerCamelCase			:   Optional[int] , lowerCamelCase			:   List[Any] , lowerCamelCase			:   str		):




				'''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 						__a   (				self			:   str		):




				'''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       UpperCamelCase__  (		__lowerCAmelCase	,__lowerCAmelCase	,unittest.TestCase     ):
	lowerCAmelCase__					:			Any								=						(
	    (
	        TFEsmModel,
	        TFEsmForMaskedLM,
	        TFEsmForSequenceClassification,
	        TFEsmForTokenClassification,
	    )
	    if is_tf_available()
	    else ()
	)
	lowerCAmelCase__					:			Dict								=						(
	    {
	        "feature-extraction": TFEsmModel,
	        "fill-mask": TFEsmForMaskedLM,
	        "text-classification": TFEsmForSequenceClassification,
	        "token-classification": TFEsmForTokenClassification,
	        "zero-shot": TFEsmForSequenceClassification,
	    }
	    if is_tf_available()
	    else {}
	)
	lowerCAmelCase__					:			Dict								=						False
	lowerCAmelCase__					:			List[Any]								=						False







	def 						__a   (				self			:   List[str]		):




				'''simple docstring'''
				a__     = TFEsmModelTester(self		)
				a__     = ConfigTester(self , config_class=lowerCamelCase , hidden_size=3_7		)







	def 						__a   (				self			:   Optional[Any]		):




				'''simple docstring'''
				self.config_tester.run_common_tests()







	def 						__a   (				self			:   Tuple		):




				'''simple docstring'''
				a__     = self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_model(*lowerCamelCase		)







	def 						__a   (				self			:   int		):




				'''simple docstring'''
				a__     = self.model_tester.prepare_config_and_inputs_for_decoder()
				self.model_tester.create_and_check_model_as_decoder(*lowerCamelCase		)







	def 						__a   (				self			:   Optional[int]		):




				'''simple docstring'''
				a__     = self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_for_masked_lm(*lowerCamelCase		)







	def 						__a   (				self			:   Any		):




				'''simple docstring'''
				a__     = self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_for_token_classification(*lowerCamelCase		)







	@slow
	def 						__a   (				self			:   List[str]		):




				'''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 						__a   (				self			:   List[Any]		):




				'''simple docstring'''
				pass







	@unittest.skip("Protein models do not support embedding resizing."		)
	def 						__a   (				self			:   Optional[Any]		):




				'''simple docstring'''
				pass







	def 						__a   (				self			:   Optional[Any]		):




				'''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       UpperCamelCase__  (		unittest.TestCase     ):







	@slow
	def 						__a   (				self			:   List[str]		):




				'''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, 3_3]
				self.assertEqual(list(output.numpy().shape		) , lowerCamelCase		)
				# compare the actual values for a slice.
				a__     = tf.constant(
				    [
				        [
				            [8.921518, -10.589814, -6.4671307],
				            [-6.3967156, -13.911377, -1.1211915],
				            [-7.781247, -13.951557, -3.740592],
				        ]
				    ]		)
				self.assertTrue(numpy.allclose(output[:, :3, :3].numpy() , expected_slice.numpy() , atol=1e-2		)		)







	@slow
	def 						__a   (				self			:   Any		):




				'''simple docstring'''
				a__     = TFEsmModel.from_pretrained("facebook/esm2_t6_8M_UR50D"		)

				a__     = tf.constant([[0, 6, 4, 1_3, 5, 4, 1_6, 1_2, 1_1, 7, 2]]		)
				a__     = model(lowerCamelCase		)[0]
				# compare the actual values for a slice.
				a__     = tf.constant(
				    [
				        [
				            [0.14443092, 0.54125327, 0.3247739],
				            [0.30340484, 0.00526676, 0.31077722],
				            [0.32278043, -0.24987096, 0.3414628],
				        ]
				    ]		)
				self.assertTrue(numpy.allclose(output[:, :3, :3].numpy() , expected_slice.numpy() , atol=1e-4		)		)










'''simple docstring'''


# Copyright 2021 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import argparse

from ...utils.dataclasses import (
    ComputeEnvironment,
    DistributedType,
    DynamoBackend,
    PrecisionType,
    SageMakerDistributedType,
)
from ..menu import BulletMenu


lowerCAmelCase_   :   List[Any]       =					[
    "EAGER",
    "AOT_EAGER",
    "INDUCTOR",
    "NVFUSER",
    "AOT_NVFUSER",
    "AOT_CUDAGRAPHS",
    "OFI",
    "FX2TRT",
    "ONNXRT",
    "IPEX",
]







def    _lowerCamelCase  (__lowerCamelCase  :     Tuple					,  __lowerCamelCase  :     Tuple=None					,  __lowerCamelCase  :     Tuple=None					,  __lowerCamelCase  :     List[Any]=None     )   ->    Any:
			a__     = True
			while ask_again:
						a__     = input(__lowerCamelCase     )
						try:
									if default is not None and len(__lowerCamelCase     ) == 0:
												return default
									return convert_value(__lowerCamelCase     ) if convert_value is not None else result
						except Exception:
									if error_message is not None:
												print(__lowerCamelCase     )







def    _lowerCamelCase  (__lowerCamelCase  :     Dict					,  __lowerCamelCase  :     List[Any]=[]					,  __lowerCamelCase  :     str=None					,  __lowerCamelCase  :     List[Any]=0     )   ->    List[Any]:
			a__     = BulletMenu(__lowerCamelCase					,  __lowerCamelCase     )
			a__     = menu.run(default_choice=__lowerCamelCase     )
			return convert_value(__lowerCamelCase     ) if convert_value is not None else result







def    _lowerCamelCase  (__lowerCamelCase  :     Tuple     )   ->    Optional[Any]:
			a__     = int(__lowerCamelCase     )
			return ComputeEnvironment(["LOCAL_MACHINE", "AMAZON_SAGEMAKER"][value]     )







def    _lowerCamelCase  (__lowerCamelCase  :     List[str]     )   ->    Optional[int]:
			a__     = int(__lowerCamelCase     )
			return DistributedType(["NO", "MULTI_CPU", "MULTI_XPU", "MULTI_GPU", "MULTI_NPU", "TPU"][value]     )







def    _lowerCamelCase  (__lowerCamelCase  :     List[Any]     )   ->    Optional[Any]:
			a__     = int(__lowerCamelCase     )
			return DynamoBackend(DYNAMO_BACKENDS[value]     ).value







def    _lowerCamelCase  (__lowerCamelCase  :     str     )   ->    str:
			a__     = int(__lowerCamelCase     )
			return PrecisionType(["no", "fp16", "bf16", "fp8"][value]     )







def    _lowerCamelCase  (__lowerCamelCase  :     int     )   ->    Optional[Any]:
			a__     = int(__lowerCamelCase     )
			return SageMakerDistributedType(["NO", "DATA_PARALLEL", "MODEL_PARALLEL"][value]     )







def    _lowerCamelCase  (__lowerCamelCase  :     Any     )   ->    List[Any]:
			return {"yes": True, "no": False}[value.lower()]






class       UpperCamelCase__  (		argparse.RawDescriptionHelpFormatter     ):







	def 						__a   (				self			:   Dict , lowerCamelCase			:   int , lowerCamelCase			:   Optional[Any] , lowerCamelCase			:   List[Any] , lowerCamelCase			:   Dict		):




				'''simple docstring'''
				a__     = super()._format_usage(lowerCamelCase , lowerCamelCase , lowerCamelCase , lowerCamelCase		)
				a__     = usage.replace("<command> [<args>] " , ""		)
				return usage



