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



	"""simple docstring"""




	def __init__(       self  :			Tuple       ,							SCREAMING_SNAKE_CASE_  :			Union[str, Any]       ,							)       ->     Any:
					lowercase_       =					parent
					lowercase_       =					1_3
					lowercase_       =					7
					lowercase_       =					True
					lowercase_       =					True
					lowercase_       =					True
					lowercase_       =					9_9
					lowercase_       =					3_2
					lowercase_       =					2
					lowercase_       =					4
					lowercase_       =					3_7
					lowercase_       =					'''gelu'''
					lowercase_       =					0.1
					lowercase_       =					0.1
					lowercase_       =					5_1_2
					lowercase_       =					1_6
					lowercase_       =					2
					lowercase_       =					0.02
					lowercase_       =					3
					lowercase_       =					4
					lowercase_       =					None
	def 						_lowercase  (       self  :			Optional[int]       )       ->     Tuple:
					lowercase_       =					ids_tensor([self.batch_size, self.seq_length]       ,							self.vocab_size       )

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

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

					lowercase_       =					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 						_lowercase  (       self  :			Any       )       ->     Optional[Any]:
					(
					    (
					    lowercase_
					)  ,	(
					    lowercase_
					)  ,	(
					    lowercase_
					)  ,	(
					    lowercase_
					)  ,	(
					    lowercase_
					)  ,	(
					    lowercase_
					)  ,	
					)       =					self.prepare_config_and_inputs()

					lowercase_       =					True
					lowercase_       =					floats_tensor([self.batch_size, self.seq_length, self.hidden_size]       )
					lowercase_       =					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 						_lowercase  (       self  :			Tuple       ,							SCREAMING_SNAKE_CASE_  :			Any       ,							SCREAMING_SNAKE_CASE_  :			List[str]       ,							SCREAMING_SNAKE_CASE_  :			str       ,							SCREAMING_SNAKE_CASE_  :			List[str]       ,							SCREAMING_SNAKE_CASE_  :			Optional[Any]       ,							SCREAMING_SNAKE_CASE_  :			str       )       ->     List[str]:
					lowercase_       =					TFEsmModel(config=SCREAMING_SNAKE_CASE_       )
					lowercase_       =					{'''input_ids''': input_ids, '''attention_mask''': input_mask}
					lowercase_       =					model(SCREAMING_SNAKE_CASE_       )

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

					lowercase_       =					model(SCREAMING_SNAKE_CASE_       )

					self.parent.assertEqual(result.last_hidden_state.shape       ,							(self.batch_size, self.seq_length, self.hidden_size)       )
	def 						_lowercase  (       self  :			List[Any]       ,							SCREAMING_SNAKE_CASE_  :			Optional[int]       ,							SCREAMING_SNAKE_CASE_  :			List[Any]       ,							SCREAMING_SNAKE_CASE_  :			Dict       ,							SCREAMING_SNAKE_CASE_  :			Optional[Any]       ,							SCREAMING_SNAKE_CASE_  :			List[Any]       ,							SCREAMING_SNAKE_CASE_  :			int       ,							SCREAMING_SNAKE_CASE_  :			int       ,							SCREAMING_SNAKE_CASE_  :			str       ,							)       ->     Union[str, Any]:
					lowercase_       =					True

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

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

					# Also check the case where encoder outputs are not passed
					lowercase_       =					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 						_lowercase  (       self  :			Optional[int]       ,							SCREAMING_SNAKE_CASE_  :			List[str]       ,							SCREAMING_SNAKE_CASE_  :			str       ,							SCREAMING_SNAKE_CASE_  :			Dict       ,							SCREAMING_SNAKE_CASE_  :			Optional[Any]       ,							SCREAMING_SNAKE_CASE_  :			Any       ,							SCREAMING_SNAKE_CASE_  :			Any       )       ->     Any:
					lowercase_       =					TFEsmForMaskedLM(config=SCREAMING_SNAKE_CASE_       )
					lowercase_       =					model([input_ids, input_mask]       )
					self.parent.assertEqual(result.logits.shape       ,							(self.batch_size, self.seq_length, self.vocab_size)       )
	def 						_lowercase  (       self  :			List[Any]       ,							SCREAMING_SNAKE_CASE_  :			List[Any]       ,							SCREAMING_SNAKE_CASE_  :			Optional[int]       ,							SCREAMING_SNAKE_CASE_  :			Union[str, Any]       ,							SCREAMING_SNAKE_CASE_  :			Optional[int]       ,							SCREAMING_SNAKE_CASE_  :			Optional[Any]       ,							SCREAMING_SNAKE_CASE_  :			str       )       ->     int:
					lowercase_       =					self.num_labels
					lowercase_       =					TFEsmForTokenClassification(config=SCREAMING_SNAKE_CASE_       )
					lowercase_       =					{'''input_ids''': input_ids, '''attention_mask''': input_mask}
					lowercase_       =					model(SCREAMING_SNAKE_CASE_       )
					self.parent.assertEqual(result.logits.shape       ,							(self.batch_size, self.seq_length, self.num_labels)       )




	def 						_lowercase  (       self  :			Tuple       )       ->     str:
					lowercase_       =					self.prepare_config_and_inputs()
					(
					    (
					    lowercase_
					)  ,	(
					    lowercase_
					)  ,	(
					    lowercase_
					)  ,	(
					    lowercase_
					)  ,	(
					    lowercase_
					)  ,	(
					    lowercase_
					)  ,	
					)       =					config_and_inputs
					lowercase_       =					{'''input_ids''': input_ids, '''attention_mask''': input_mask}
					return config, inputs_dict

@require_tf
class  lowercase__(				UpperCAmelCase ,						UpperCAmelCase ,						unittest.TestCase							):



	"""simple docstring"""




	a			:str         =	(
	    (
	        TFEsmModel,
	        TFEsmForMaskedLM,
	        TFEsmForSequenceClassification,
	        TFEsmForTokenClassification,
	    )
	    if is_tf_available()
	    else ()
	)
	a			:Union[str, Any]         =	(
	    {
	        'feature-extraction': TFEsmModel,
	        'fill-mask': TFEsmForMaskedLM,
	        'text-classification': TFEsmForSequenceClassification,
	        'token-classification': TFEsmForTokenClassification,
	        'zero-shot': TFEsmForSequenceClassification,
	    }
	    if is_tf_available()
	    else {}
	)
	a			:Union[str, Any]         =	False
	a			:Dict         =	False
	def 						_lowercase  (       self  :			List[str]       )       ->     List[Any]:
					lowercase_       =					TFEsmModelTester(self       )
					lowercase_       =					ConfigTester(self       ,							config_class=SCREAMING_SNAKE_CASE_       ,							hidden_size=3_7       )
	def 						_lowercase  (       self  :			Optional[int]       )       ->     str:
					self.config_tester.run_common_tests()
	def 						_lowercase  (       self  :			str       )       ->     int:
					lowercase_       =					self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_model(*SCREAMING_SNAKE_CASE_       )
	def 						_lowercase  (       self  :			Any       )       ->     Union[str, Any]:
					lowercase_       =					self.model_tester.prepare_config_and_inputs_for_decoder()
					self.model_tester.create_and_check_model_as_decoder(*SCREAMING_SNAKE_CASE_       )
	def 						_lowercase  (       self  :			str       )       ->     Optional[int]:
					lowercase_       =					self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_for_masked_lm(*SCREAMING_SNAKE_CASE_       )
	def 						_lowercase  (       self  :			Optional[Any]       )       ->     List[str]:
					lowercase_       =					self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_for_token_classification(*SCREAMING_SNAKE_CASE_       )
	@slow
	def 						_lowercase  (       self  :			int       )       ->     List[str]:
					for model_name in TF_ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
									lowercase_       =					TFEsmModel.from_pretrained(SCREAMING_SNAKE_CASE_       )
									self.assertIsNotNone(SCREAMING_SNAKE_CASE_       )
	@unittest.skip('''Protein models do not support embedding resizing.'''       )
	def 						_lowercase  (       self  :			str       )       ->     Tuple:
					pass
	@unittest.skip('''Protein models do not support embedding resizing.'''       )
	def 						_lowercase  (       self  :			List[Any]       )       ->     Any:
					pass




	def 						_lowercase  (       self  :			List[str]       )       ->     Any:
					lowercase_  ,	lowercase_       =					self.model_tester.prepare_config_and_inputs_for_common()

					for model_class in self.all_model_classes:
									lowercase_       =					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
													lowercase_       =					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:
													lowercase_       =					model.get_output_embeddings()
													assert x is None
													lowercase_       =					model.get_bias()
													assert name is None

@require_tf
class  lowercase__(				unittest.TestCase							):



	"""simple docstring"""




	@slow
	def 						_lowercase  (       self  :			str       )       ->     Optional[int]:
					lowercase_       =					TFEsmForMaskedLM.from_pretrained('''facebook/esm2_t6_8M_UR50D'''       )

					lowercase_       =					tf.constant([[0, 1, 2, 3, 4, 5]]       )
					lowercase_       =					model(SCREAMING_SNAKE_CASE_       )[0]
					lowercase_       =					[1, 6, 3_3]
					self.assertEqual(list(output.numpy().shape       )       ,							SCREAMING_SNAKE_CASE_       )
					# compare the actual values for a slice.
					lowercase_       =					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 						_lowercase  (       self  :			int       )       ->     Union[str, Any]:
					lowercase_       =					TFEsmModel.from_pretrained('''facebook/esm2_t6_8M_UR50D'''       )

					lowercase_       =					tf.constant([[0, 6, 4, 1_3, 5, 4, 1_6, 1_2, 1_1, 7, 2]]       )
					lowercase_       =					model(SCREAMING_SNAKE_CASE_       )[0]
					# compare the actual values for a slice.
					lowercase_       =					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       )       )






'''simple docstring'''







import warnings

from ...configuration_utils import PretrainedConfig
from ...utils import logging


lowercase    :							int      =					logging.get_logger(__name__)

lowercase    :							Union[str, Any]      =					{
    "xlnet-base-cased": "https://huggingface.co/xlnet-base-cased/resolve/main/config.json",
    "xlnet-large-cased": "https://huggingface.co/xlnet-large-cased/resolve/main/config.json",
}


class 						__UpperCAmelCase      (					_lowerCamelCase						):
							__lowercase      		=       """xlnet"""
							__lowercase      		=       ["""mems"""]
							__lowercase      		=       {
							    """n_token""": """vocab_size""",  # Backward compatibility
							    """hidden_size""": """d_model""",
							    """num_attention_heads""": """n_head""",
							    """num_hidden_layers""": """n_layer""",
							}


							def __init__(       self    ,   lowerCAmelCase_=3_20_00    ,   lowerCAmelCase_=10_24    ,   lowerCAmelCase_=24    ,   lowerCAmelCase_=16    ,   lowerCAmelCase_=40_96    ,   lowerCAmelCase_="gelu"    ,   lowerCAmelCase_=True    ,   lowerCAmelCase_="bi"    ,   lowerCAmelCase_=0.02    ,   lowerCAmelCase_=1E-12    ,   lowerCAmelCase_=0.1    ,   lowerCAmelCase_=5_12    ,   lowerCAmelCase_=None    ,   lowerCAmelCase_=True    ,   lowerCAmelCase_=False    ,   lowerCAmelCase_=False    ,   lowerCAmelCase_=-1    ,   lowerCAmelCase_=False    ,   lowerCAmelCase_="last"    ,   lowerCAmelCase_=True    ,   lowerCAmelCase_="tanh"    ,   lowerCAmelCase_=0.1    ,   lowerCAmelCase_=5    ,   lowerCAmelCase_=5    ,   lowerCAmelCase_=5    ,   lowerCAmelCase_=1    ,   lowerCAmelCase_=2    ,   **lowerCAmelCase_    ,   ):





								"""simple docstring"""






								_snake_case   						=		vocab_size
								_snake_case   						=		d_model
								_snake_case   						=		n_layer
								_snake_case   						=		n_head
								if d_model % n_head != 0:
									raise ValueError(F'\'d_model % n_head\' ({d_model % n_head}) should be equal to 0'  )
								if "d_head" in kwargs:
									if kwargs["d_head"] != d_model // n_head:
										raise ValueError(
										    F'`d_head` ({kwargs["d_head"]}) should be equal to `d_model // n_head` ({d_model // n_head})'  )
								_snake_case   						=		d_model // n_head
								_snake_case   						=		ff_activation
								_snake_case   						=		d_inner
								_snake_case   						=		untie_r
								_snake_case   						=		attn_type

								_snake_case   						=		initializer_range
								_snake_case   						=		layer_norm_eps

								_snake_case   						=		dropout
								_snake_case   						=		mem_len
								_snake_case   						=		reuse_len
								_snake_case   						=		bi_data
								_snake_case   						=		clamp_len
								_snake_case   						=		same_length

								_snake_case   						=		summary_type
								_snake_case   						=		summary_use_proj
								_snake_case   						=		summary_activation
								_snake_case   						=		summary_last_dropout
								_snake_case   						=		start_n_top
								_snake_case   						=		end_n_top

								_snake_case   						=		bos_token_id
								_snake_case   						=		pad_token_id
								_snake_case   						=		eos_token_id

								if "use_cache" in kwargs:
									warnings.warn(
									    'The `use_cache` argument is deprecated and will be removed in a future version, use `use_mems_eval`'
									    ' instead.'    ,   lowerCAmelCase_    ,   )
									_snake_case   						=		kwargs['use_cache']

								_snake_case   						=		use_mems_eval
								_snake_case   						=		use_mems_train
								super().__init__(pad_token_id=lowerCAmelCase_    ,   bos_token_id=lowerCAmelCase_    ,   eos_token_id=lowerCAmelCase_    ,   **lowerCAmelCase_  )


							@property
							def  lowerCamelCase				(       self  ):





								"""simple docstring"""






								logger.info(F'The model {self.model_type} is one of the few models that has no sequence length limit.'  )
								return -1




							@max_position_embeddings.setter
							def  lowerCamelCase				(       self    ,   lowerCAmelCase_  ):





								"""simple docstring"""






								raise NotImplementedError(
								    F'The model {self.model_type} is one of the few models that has no sequence length limit.'  )

