






import unittest

from transformers import EsmConfig, is_torch_available
from transformers.testing_utils import TestCasePlus, require_torch, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
						import torch

						from transformers import EsmForMaskedLM, EsmForSequenceClassification, EsmForTokenClassification, EsmModel
						from transformers.models.esm.modeling_esm import (
						    ESM_PRETRAINED_MODEL_ARCHIVE_LIST,
						    EsmEmbeddings,
						    create_position_ids_from_input_ids,
						)


class 							__UpperCamelCase     :







	def __init__(	self   :				Union[str, Any]							,					_lowerCAmelCase   :				Any							,					_lowerCAmelCase   :				Optional[Any]=13							,					_lowerCAmelCase   :				List[Any]=7							,					_lowerCAmelCase   :				Optional[Any]=False							,					_lowerCAmelCase   :				Optional[Any]=True							,					_lowerCAmelCase   :				Any=False							,					_lowerCAmelCase   :				Union[str, Any]=True							,					_lowerCAmelCase   :				Optional[int]=33							,					_lowerCAmelCase   :				Any=32							,					_lowerCAmelCase   :				List[Any]=5							,					_lowerCAmelCase   :				Tuple=4							,					_lowerCAmelCase   :				List[Any]=37							,					_lowerCAmelCase   :				Any="gelu"							,					_lowerCAmelCase   :				Optional[Any]=0.1							,					_lowerCAmelCase   :				Optional[Any]=0.1							,					_lowerCAmelCase   :				str=512							,					_lowerCAmelCase   :				Union[str, Any]=16							,					_lowerCAmelCase   :				List[Any]=2							,					_lowerCAmelCase   :				Any=0.02							,					_lowerCAmelCase   :				List[Any]=3							,					_lowerCAmelCase   :				List[Any]=4							,					_lowerCAmelCase   :				Any=None							,					)		->			Optional[int]:
		"""simple docstring"""







		__lowercase   							=	parent
		__lowercase   							=	batch_size
		__lowercase   							=	seq_length
		__lowercase   							=	is_training
		__lowercase   							=	use_input_mask
		__lowercase   							=	use_token_type_ids
		__lowercase   							=	use_labels
		__lowercase   							=	vocab_size
		__lowercase   							=	hidden_size
		__lowercase   							=	num_hidden_layers
		__lowercase   							=	num_attention_heads
		__lowercase   							=	intermediate_size
		__lowercase   							=	hidden_act
		__lowercase   							=	hidden_dropout_prob
		__lowercase   							=	attention_probs_dropout_prob
		__lowercase   							=	max_position_embeddings
		__lowercase   							=	type_vocab_size
		__lowercase   							=	type_sequence_label_size
		__lowercase   							=	initializer_range
		__lowercase   							=	num_labels
		__lowercase   							=	num_choices
		__lowercase   							=	scope







	def 			_a (	self   :				List[str]      )		->			Any:
		"""simple docstring"""







		__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   							=	self.get_config()

		return config, input_ids, input_mask, sequence_labels, token_labels, choice_labels







	def 			_a (	self   :				List[Any]      )		->			Union[str, Any]:
		"""simple docstring"""







		return EsmConfig(
		    vocab_size=self.vocab_size							,					hidden_size=self.hidden_size							,					pad_token_id=1							,					num_hidden_layers=self.num_hidden_layers							,					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							,					)







	def 			_a (	self   :				Optional[Any]							,					_lowerCAmelCase   :				Dict							,					_lowerCAmelCase   :				List[str]							,					_lowerCAmelCase   :				str							,					_lowerCAmelCase   :				Any							,					_lowerCAmelCase   :				str							,					_lowerCAmelCase   :				str      )		->			int:
		"""simple docstring"""







		__lowercase   							=	EsmModel(config=__lowerCAmelCase      )
		model.to(__lowerCAmelCase      )
		model.eval()
		__lowercase   							=	model(__lowerCAmelCase							,					attention_mask=__lowerCAmelCase      )
		__lowercase   							=	model(__lowerCAmelCase      )
		__lowercase   							=	model(__lowerCAmelCase      )

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







	def 			_a (	self   :				List[Any]							,					_lowerCAmelCase   :				int							,					_lowerCAmelCase   :				List[Any]							,					_lowerCAmelCase   :				int							,					_lowerCAmelCase   :				List[str]							,					_lowerCAmelCase   :				Dict							,					_lowerCAmelCase   :				Tuple      )		->			Any:
		"""simple docstring"""







		__lowercase   							=	EsmForMaskedLM(config=__lowerCAmelCase      )
		model.to(__lowerCAmelCase      )
		model.eval()
		__lowercase   							=	model(__lowerCAmelCase							,					attention_mask=__lowerCAmelCase							,					labels=__lowerCAmelCase      )
		self.parent.assertEqual(result.logits.shape							,					(self.batch_size, self.seq_length, self.vocab_size)      )







	def 			_a (	self   :				Tuple							,					_lowerCAmelCase   :				Tuple							,					_lowerCAmelCase   :				Dict							,					_lowerCAmelCase   :				Dict							,					_lowerCAmelCase   :				Optional[Any]							,					_lowerCAmelCase   :				Optional[int]							,					_lowerCAmelCase   :				Any      )		->			Tuple:
		"""simple docstring"""







		__lowercase   							=	self.num_labels
		__lowercase   							=	EsmForTokenClassification(config=__lowerCAmelCase      )
		model.to(__lowerCAmelCase      )
		model.eval()
		__lowercase   							=	model(__lowerCAmelCase							,					attention_mask=__lowerCAmelCase							,					labels=__lowerCAmelCase      )
		self.parent.assertEqual(result.logits.shape							,					(self.batch_size, self.seq_length, self.num_labels)      )







	def 			_a (	self   :				Optional[int]      )		->			Union[str, Any]:
		"""simple docstring"""







		__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_torch
class 							__UpperCamelCase     (  lowerCAmelCase__ ,      lowerCAmelCase__ ,      unittest.TestCase ):
	__snake_case						:Optional[Any] =							False

	__snake_case						:Optional[Any] =							(
	    (
	        EsmForMaskedLM,
	        EsmModel,
	        EsmForSequenceClassification,
	        EsmForTokenClassification,
	    )
	    if is_torch_available()
	    else ()
	)
	__snake_case						:List[Any] =							()
	__snake_case						:Tuple =							(
	    {
	        '''feature-extraction''': EsmModel,
	        '''fill-mask''': EsmForMaskedLM,
	        '''text-classification''': EsmForSequenceClassification,
	        '''token-classification''': EsmForTokenClassification,
	        '''zero-shot''': EsmForSequenceClassification,
	    }
	    if is_torch_available()
	    else {}
	)
	__snake_case						:List[str] =							True







	def 			_a (	self   :				Tuple      )		->			Dict:
		"""simple docstring"""







		__lowercase   							=	EsmModelTester(self      )
		__lowercase   							=	ConfigTester(self							,					config_class=__lowerCAmelCase							,					hidden_size=37      )







	def 			_a (	self   :				List[Any]      )		->			str:
		"""simple docstring"""







		self.config_tester.run_common_tests()







	def 			_a (	self   :				Dict      )		->			Tuple:
		"""simple docstring"""







		__lowercase   							=	self.model_tester.prepare_config_and_inputs()
		self.model_tester.create_and_check_model(*__lowerCAmelCase      )







	def 			_a (	self   :				Dict      )		->			Union[str, Any]:
		"""simple docstring"""







		__lowercase   							=	self.model_tester.prepare_config_and_inputs()
		for type in ["absolute", "relative_key", "relative_key_query"]:
			__lowercase   							=	type
			self.model_tester.create_and_check_model(*__lowerCAmelCase      )







	def 			_a (	self   :				Optional[int]      )		->			Dict:
		"""simple docstring"""







		__lowercase   							=	self.model_tester.prepare_config_and_inputs()
		self.model_tester.create_and_check_for_masked_lm(*__lowerCAmelCase      )







	def 			_a (	self   :				Tuple      )		->			Optional[Any]:
		"""simple docstring"""







		__lowercase   							=	self.model_tester.prepare_config_and_inputs()
		self.model_tester.create_and_check_for_token_classification(*__lowerCAmelCase      )







	@slow
	def 			_a (	self   :				Union[str, Any]      )		->			Dict:
		"""simple docstring"""







		for model_name in ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
			__lowercase   							=	EsmModel.from_pretrained(__lowerCAmelCase      )
			self.assertIsNotNone(__lowerCAmelCase      )







	def 			_a (	self   :				Tuple      )		->			List[str]:
		"""simple docstring"""







		__lowercase   							=	self.model_tester.prepare_config_and_inputs()[0]
		__lowercase   							=	EsmEmbeddings(config=__lowerCAmelCase      )

		__lowercase   							=	torch.as_tensor([[12, 31, 13, model.padding_idx]]      )
		__lowercase   							=	torch.as_tensor(
		    [
		        [
		            0 + model.padding_idx + 1,
		            1 + model.padding_idx + 1,
		            2 + model.padding_idx + 1,
		            model.padding_idx,
		        ]
		    ]      )
		__lowercase   							=	create_position_ids_from_input_ids(__lowerCAmelCase							,					model.padding_idx      )
		self.assertEqual(position_ids.shape							,					expected_positions.shape      )
		self.assertTrue(torch.all(torch.eq(__lowerCAmelCase							,					__lowerCAmelCase      )      )      )







	def 			_a (	self   :				Optional[Any]      )		->			Optional[int]:
		"""simple docstring"""







		__lowercase   							=	self.model_tester.prepare_config_and_inputs()[0]
		__lowercase   							=	EsmEmbeddings(config=__lowerCAmelCase      )

		__lowercase   							=	torch.empty(2							,					4							,					30      )
		__lowercase   							=	[
		    0 + embeddings.padding_idx + 1,
		    1 + embeddings.padding_idx + 1,
		    2 + embeddings.padding_idx + 1,
		    3 + embeddings.padding_idx + 1,
		]
		__lowercase   							=	torch.as_tensor([expected_single_positions, expected_single_positions]      )
		__lowercase   							=	embeddings.create_position_ids_from_inputs_embeds(__lowerCAmelCase      )
		self.assertEqual(position_ids.shape							,					expected_positions.shape      )
		self.assertTrue(torch.all(torch.eq(__lowerCAmelCase							,					__lowerCAmelCase      )      )      )







	@unittest.skip("""Esm does not support embedding resizing"""      )
	def 			_a (	self   :				Tuple      )		->			int:
		"""simple docstring"""







		pass







	@unittest.skip("""Esm does not support embedding resizing"""      )
	def 			_a (	self   :				List[str]      )		->			Any:
		"""simple docstring"""







		pass







	@unittest.skip("""Will be fixed soon by reducing the size of the model used for common tests."""      )
	def 			_a (	self   :				str      )		->			Optional[Any]:
		"""simple docstring"""







		pass


@require_torch
class 							__UpperCamelCase     (  lowerCAmelCase__ ):







	@slow
	def 			_a (	self   :				Any      )		->			int:
		"""simple docstring"""







		with torch.no_grad():
			__lowercase   							=	EsmForMaskedLM.from_pretrained("""facebook/esm2_t6_8M_UR50D"""      )
			model.eval()
			__lowercase   							=	torch.tensor([[0, 1, 2, 3, 4, 5]]      )
			__lowercase   							=	model(__lowerCAmelCase      )[0]

			__lowercase   							=	33

			__lowercase   							=	torch.Size((1, 6, vocab_size)      )
			self.assertEqual(output.shape							,					__lowerCAmelCase      )

			__lowercase   							=	torch.tensor(
			    [[[8.9_215, -10.5_898, -6.4_671], [-6.3_967, -13.9_114, -1.1_212], [-7.7_812, -13.9_516, -3.7_406]]]      )
			self.assertTrue(torch.allclose(output[:, :3, :3]							,					__lowerCAmelCase							,					atol=1e-4      )      )







	@slow
	def 			_a (	self   :				List[str]      )		->			Dict:
		"""simple docstring"""







		with torch.no_grad():
			__lowercase   							=	EsmModel.from_pretrained("""facebook/esm2_t6_8M_UR50D"""      )
			model.eval()

			__lowercase   							=	torch.tensor([[0, 6, 4, 13, 5, 4, 16, 12, 11, 7, 2]]      )
			__lowercase   							=	model(__lowerCAmelCase      )[0]
			# compare the actual values for a slice.
			__lowercase   							=	torch.tensor(
			    [[[0.1_444, 0.5_413, 0.3_248], [0.3_034, 0.0_053, 0.3_108], [0.3_228, -0.2_499, 0.3_415]]]      )
			self.assertTrue(torch.allclose(output[:, :3, :3]							,					__lowerCAmelCase							,					atol=1e-4      )      )


'''simple docstring'''




def 							snake_case__				(     _A:   int					)       ->      list[int]:

		'''simple docstring'''


		if length <= 0 or not isinstance(_A ,  _A					):
				raise ValueError("""Length must be a positive integer."""					)
		return [n * (2 * n - 1) for n in range(_A					)]


if __name__ == "__main__":
						print(hexagonal_numbers(length=5))
						print(hexagonal_numbers(length=1_0))




