


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




	def __init__( self						:		int		, _SCREAMING_SNAKE_CASE						:		Tuple		, _SCREAMING_SNAKE_CASE						:		Tuple=13		, _SCREAMING_SNAKE_CASE						:		Dict=7		, _SCREAMING_SNAKE_CASE						:		List[Any]=False		, _SCREAMING_SNAKE_CASE						:		Union[str, Any]=True		, _SCREAMING_SNAKE_CASE						:		str=False		, _SCREAMING_SNAKE_CASE						:		Any=True		, _SCREAMING_SNAKE_CASE						:		int=33		, _SCREAMING_SNAKE_CASE						:		Any=32		, _SCREAMING_SNAKE_CASE						:		Tuple=5		, _SCREAMING_SNAKE_CASE						:		List[str]=4		, _SCREAMING_SNAKE_CASE						:		Optional[Any]=37		, _SCREAMING_SNAKE_CASE						:		List[Any]="gelu"		, _SCREAMING_SNAKE_CASE						:		int=0.1		, _SCREAMING_SNAKE_CASE						:		Optional[Any]=0.1		, _SCREAMING_SNAKE_CASE						:		Optional[Any]=512		, _SCREAMING_SNAKE_CASE						:		Dict=16		, _SCREAMING_SNAKE_CASE						:		Dict=2		, _SCREAMING_SNAKE_CASE						:		Dict=0.02		, _SCREAMING_SNAKE_CASE						:		Optional[Any]=3		, _SCREAMING_SNAKE_CASE						:		str=4		, _SCREAMING_SNAKE_CASE						:		int=None		, )->	Dict:
			lowerCAmelCase__       :      Dict   					= parent
			lowerCAmelCase__       :      List[Any]   					= batch_size
			lowerCAmelCase__       :      Tuple   					= seq_length
			lowerCAmelCase__       :      Optional[Any]   					= is_training
			lowerCAmelCase__       :      str   					= use_input_mask
			lowerCAmelCase__       :      Optional[int]   					= use_token_type_ids
			lowerCAmelCase__       :      List[str]   					= use_labels
			lowerCAmelCase__       :      Union[str, Any]   					= vocab_size
			lowerCAmelCase__       :      List[Any]   					= hidden_size
			lowerCAmelCase__       :      Optional[Any]   					= num_hidden_layers
			lowerCAmelCase__       :      Dict   					= num_attention_heads
			lowerCAmelCase__       :      Dict   					= intermediate_size
			lowerCAmelCase__       :      List[Any]   					= hidden_act
			lowerCAmelCase__       :      str   					= hidden_dropout_prob
			lowerCAmelCase__       :      str   					= attention_probs_dropout_prob
			lowerCAmelCase__       :      Dict   					= max_position_embeddings
			lowerCAmelCase__       :      Optional[int]   					= type_vocab_size
			lowerCAmelCase__       :      Optional[int]   					= type_sequence_label_size
			lowerCAmelCase__       :      List[Any]   					= initializer_range
			lowerCAmelCase__       :      Union[str, Any]   					= num_labels
			lowerCAmelCase__       :      int   					= num_choices
			lowerCAmelCase__       :      Union[str, Any]   					= scope




	def 							UpperCAmelCase__( self						:		Union[str, Any]		)->	List[Any]:
			lowerCAmelCase__       :      int   					= ids_tensor([self.batch_size, self.seq_length]		, self.vocab_size		)

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

			lowerCAmelCase__       :      Optional[Any]   					= None
			lowerCAmelCase__       :      List[Any]   					= None
			lowerCAmelCase__       :      Any   					= None
			if self.use_labels:
					lowerCAmelCase__       :      Optional[Any]   					= ids_tensor([self.batch_size]		, self.type_sequence_label_size		)
					lowerCAmelCase__       :      int   					= ids_tensor([self.batch_size, self.seq_length]		, self.num_labels		)
					lowerCAmelCase__       :      List[Any]   					= ids_tensor([self.batch_size]		, self.num_choices		)

			lowerCAmelCase__       :      str   					= self.get_config()

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




	def 							UpperCAmelCase__( self						:		List[Any]		)->	Tuple:
			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 							UpperCAmelCase__( self						:		Dict		, _SCREAMING_SNAKE_CASE						:		str		, _SCREAMING_SNAKE_CASE						:		Any		, _SCREAMING_SNAKE_CASE						:		str		, _SCREAMING_SNAKE_CASE						:		Tuple		, _SCREAMING_SNAKE_CASE						:		Union[str, Any]		, _SCREAMING_SNAKE_CASE						:		Union[str, Any]		)->	str:
			lowerCAmelCase__       :      str   					= EsmModel(config=__A		)
			model.to(__A		)
			model.eval()
			lowerCAmelCase__       :      Dict   					= model(__A		, attention_mask=__A		)
			lowerCAmelCase__       :      Any   					= model(__A		)
			lowerCAmelCase__       :      List[Any]   					= model(__A		)

			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 							UpperCAmelCase__( self						:		Union[str, Any]		, _SCREAMING_SNAKE_CASE						:		List[Any]		, _SCREAMING_SNAKE_CASE						:		List[str]		, _SCREAMING_SNAKE_CASE						:		List[Any]		, _SCREAMING_SNAKE_CASE						:		Optional[int]		, _SCREAMING_SNAKE_CASE						:		Optional[Any]		, _SCREAMING_SNAKE_CASE						:		Tuple		)->	List[str]:
			lowerCAmelCase__       :      List[Any]   					= EsmForMaskedLM(config=__A		)
			model.to(__A		)
			model.eval()
			lowerCAmelCase__       :      Optional[int]   					= model(__A		, attention_mask=__A		, labels=__A		)
			self.parent.assertEqual(result.logits.shape		, (self.batch_size, self.seq_length, self.vocab_size)		)




	def 							UpperCAmelCase__( self						:		int		, _SCREAMING_SNAKE_CASE						:		str		, _SCREAMING_SNAKE_CASE						:		int		, _SCREAMING_SNAKE_CASE						:		List[Any]		, _SCREAMING_SNAKE_CASE						:		List[str]		, _SCREAMING_SNAKE_CASE						:		Any		, _SCREAMING_SNAKE_CASE						:		Optional[Any]		)->	int:
			lowerCAmelCase__       :      List[Any]   					= self.num_labels
			lowerCAmelCase__       :      Optional[int]   					= EsmForTokenClassification(config=__A		)
			model.to(__A		)
			model.eval()
			lowerCAmelCase__       :      Dict   					= model(__A		, attention_mask=__A		, labels=__A		)
			self.parent.assertEqual(result.logits.shape		, (self.batch_size, self.seq_length, self.num_labels)		)




	def 							UpperCAmelCase__( self						:		Optional[int]		)->	List[Any]:
			lowerCAmelCase__       :      Any   					= self.prepare_config_and_inputs()
			(
			    lowerCAmelCase__	
			)       :      Optional[int]   					= config_and_inputs
			lowerCAmelCase__       :      Optional[Any]   					= {"""input_ids""": input_ids, """attention_mask""": input_mask}
			return config, inputs_dict





@require_torch
class  _a							(    A__					,						A__					,						unittest.TestCase):
	_a  :  Dict				   =							False

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




	def 							UpperCAmelCase__( self						:		Dict		)->	str:
			lowerCAmelCase__       :      Any   					= EsmModelTester(self		)
			lowerCAmelCase__       :      List[str]   					= ConfigTester(self		, config_class=__A		, hidden_size=37		)




	def 							UpperCAmelCase__( self						:		Optional[Any]		)->	List[Any]:
			self.config_tester.run_common_tests()




	def 							UpperCAmelCase__( self						:		Union[str, Any]		)->	str:
			lowerCAmelCase__       :      Union[str, Any]   					= self.model_tester.prepare_config_and_inputs()
			self.model_tester.create_and_check_model(*__A		)




	def 							UpperCAmelCase__( self						:		Dict		)->	Dict:
			lowerCAmelCase__       :      Optional[int]   					= self.model_tester.prepare_config_and_inputs()
			for type in ["absolute", "relative_key", "relative_key_query"]:
					lowerCAmelCase__       :      List[Any]   					= type
					self.model_tester.create_and_check_model(*__A		)




	def 							UpperCAmelCase__( self						:		int		)->	Tuple:
			lowerCAmelCase__       :      Any   					= self.model_tester.prepare_config_and_inputs()
			self.model_tester.create_and_check_for_masked_lm(*__A		)




	def 							UpperCAmelCase__( self						:		Any		)->	Any:
			lowerCAmelCase__       :      Optional[int]   					= self.model_tester.prepare_config_and_inputs()
			self.model_tester.create_and_check_for_token_classification(*__A		)




	@slow
	def 							UpperCAmelCase__( self						:		str		)->	Union[str, Any]:
			for model_name in ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
					lowerCAmelCase__       :      Any   					= EsmModel.from_pretrained(__A		)
					self.assertIsNotNone(__A		)




	def 							UpperCAmelCase__( self						:		List[Any]		)->	str:
			lowerCAmelCase__       :      Union[str, Any]   					= self.model_tester.prepare_config_and_inputs()[0]
			lowerCAmelCase__       :      Dict   					= EsmEmbeddings(config=__A		)

			lowerCAmelCase__       :      Optional[Any]   					= torch.as_tensor([[12, 31, 13, model.padding_idx]]		)
			lowerCAmelCase__       :      str   					= torch.as_tensor(
			    [
			        [
			            0 + model.padding_idx + 1,
			            1 + model.padding_idx + 1,
			            2 + model.padding_idx + 1,
			            model.padding_idx,
			        ]
			    ]		)
			lowerCAmelCase__       :      List[Any]   					= create_position_ids_from_input_ids(__A		, model.padding_idx		)
			self.assertEqual(position_ids.shape		, expected_positions.shape		)
			self.assertTrue(torch.all(torch.eq(__A		, __A		)		)		)




	def 							UpperCAmelCase__( self						:		List[str]		)->	List[str]:
			lowerCAmelCase__       :      Tuple   					= self.model_tester.prepare_config_and_inputs()[0]
			lowerCAmelCase__       :      List[str]   					= EsmEmbeddings(config=__A		)

			lowerCAmelCase__       :      int   					= torch.empty(2		, 4		, 30		)
			lowerCAmelCase__       :      Tuple   					= [
			    0 + embeddings.padding_idx + 1,
			    1 + embeddings.padding_idx + 1,
			    2 + embeddings.padding_idx + 1,
			    3 + embeddings.padding_idx + 1,
			]
			lowerCAmelCase__       :      Dict   					= torch.as_tensor([expected_single_positions, expected_single_positions]		)
			lowerCAmelCase__       :      List[str]   					= embeddings.create_position_ids_from_inputs_embeds(__A		)
			self.assertEqual(position_ids.shape		, expected_positions.shape		)
			self.assertTrue(torch.all(torch.eq(__A		, __A		)		)		)




	@unittest.skip('''Esm does not support embedding resizing'''		)
	def 							UpperCAmelCase__( self						:		Union[str, Any]		)->	Dict:
			pass




	@unittest.skip('''Esm does not support embedding resizing'''		)
	def 							UpperCAmelCase__( self						:		List[str]		)->	Dict:
			pass




	@unittest.skip('''Will be fixed soon by reducing the size of the model used for common tests.'''		)
	def 							UpperCAmelCase__( self						:		List[Any]		)->	str:
			pass





@require_torch
class  _a							(    A__):




	@slow
	def 							UpperCAmelCase__( self						:		int		)->	List[str]:
			with torch.no_grad():
					lowerCAmelCase__       :      int   					= EsmForMaskedLM.from_pretrained('''facebook/esm2_t6_8M_UR50D'''		)
					model.eval()
					lowerCAmelCase__       :      List[Any]   					= torch.tensor([[0, 1, 2, 3, 4, 5]]		)
					lowerCAmelCase__       :      Optional[int]   					= model(__A		)[0]

					lowerCAmelCase__       :      Optional[Any]   					= 33

					lowerCAmelCase__       :      Union[str, Any]   					= torch.Size((1, 6, vocab_size)		)
					self.assertEqual(output.shape		, __A		)

					lowerCAmelCase__       :      List[str]   					= torch.tensor(
					    [[[8.9215, -10.5898, -6.4671], [-6.3967, -13.9114, -1.1212], [-7.7812, -13.9516, -3.7406]]]		)
					self.assertTrue(torch.allclose(output[:, :3, :3]		, __A		, atol=1E-4		)		)




	@slow
	def 							UpperCAmelCase__( self						:		Optional[int]		)->	List[str]:
			with torch.no_grad():
					lowerCAmelCase__       :      Any   					= EsmModel.from_pretrained('''facebook/esm2_t6_8M_UR50D'''		)
					model.eval()

					lowerCAmelCase__       :      Dict   					= torch.tensor([[0, 6, 4, 13, 5, 4, 16, 12, 11, 7, 2]]		)
					lowerCAmelCase__       :      Union[str, Any]   					= model(__A		)[0]
					# compare the actual values for a slice.
					lowerCAmelCase__       :      List[str]   					= torch.tensor(
					    [[[0.1444, 0.5413, 0.3248], [0.3034, 0.0053, 0.3108], [0.3228, -0.2499, 0.3415]]]		)
					self.assertTrue(torch.allclose(output[:, :3, :3]		, __A		, atol=1E-4		)		)








import timeit

import numpy as np

import datasets
from datasets.arrow_writer import ArrowWriter
from datasets.features.features import _ArrayXD






def 					lowerCamelCase_   (	_a	):




		"""simple docstring"""






		def wrapper(*_a     ,       **_a	):
				lowerCAmelCase__       :      List[str]   					= timeit.default_timer()
				lowerCAmelCase__       :      List[Any]   					= func(*_a     ,       **_a	)
				lowerCAmelCase__       :      Any   					= timeit.default_timer() - starttime
				return delta

		lowerCAmelCase__       :      Any   					= func.__name__

		return wrapper






def 					lowerCamelCase_   (	_a     ,       _a=100     ,       _a=None	):




		"""simple docstring"""






		lowerCAmelCase__       :      str   					= []
		lowerCAmelCase__       :      str   					= seq_shapes or {}
		for i in range(_a	):
				lowerCAmelCase__       :      List[str]   					= {}
				for col_id, (k, v) in enumerate(features.items()	):
						if isinstance(_a     ,       _ArrayXD	):
								lowerCAmelCase__       :      List[str]   					= np.random.rand(*v.shape	).astype(v.dtype	)
						elif isinstance(_a     ,       datasets.Value	):
								if v.dtype == "string":
										lowerCAmelCase__       :      Dict   					= '''The small grey turtle was surprisingly fast when challenged.'''
								else:
										lowerCAmelCase__       :      Any   					= np.random.randint(10     ,       size=1	).astype(v.dtype	).item()
						elif isinstance(_a     ,       datasets.Sequence	):
								while isinstance(_a     ,       datasets.Sequence	):
										lowerCAmelCase__       :      Optional[int]   					= v.feature
								lowerCAmelCase__       :      str   					= seq_shapes[k]
								lowerCAmelCase__       :      Any   					= np.random.rand(*_a	).astype(v.dtype	)
						lowerCAmelCase__       :      int   					= data

				dummy_data.append((i, example)	)

		return dummy_data






def 					lowerCamelCase_   (	_a     ,       _a     ,       _a=100     ,       _a=None	):




		"""simple docstring"""






		lowerCAmelCase__       :      Optional[Any]   					= generate_examples(_a     ,       num_examples=_a     ,       seq_shapes=_a	)

		with ArrowWriter(features=_a     ,       path=_a	) as writer:
				for key, record in dummy_data:
						lowerCAmelCase__       :      Optional[int]   					= features.encode_example(_a	)
						writer.write(_a	)

				lowerCAmelCase__					,     lowerCAmelCase__       :      List[str]   					= writer.finalize()

		if not num_final_examples == num_examples:
				raise ValueError(
				    f'Error writing the dataset, wrote {num_final_examples} examples but should have written {num_examples}.'	)

		lowerCAmelCase__       :      List[Any]   					= datasets.Dataset.from_file(filename=_a     ,       info=datasets.DatasetInfo(features=_a	)	)

		return dataset




