

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	:						str		,	lowerCAmelCase	:						Any		,	lowerCAmelCase	:						Optional[Any]=13		,	lowerCAmelCase	:						List[str]=7		,	lowerCAmelCase	:						Dict=False		,	lowerCAmelCase	:						Dict=True		,	lowerCAmelCase	:						Tuple=False		,	lowerCAmelCase	:						int=True		,	lowerCAmelCase	:						Any=33		,	lowerCAmelCase	:						int=32		,	lowerCAmelCase	:						int=5		,	lowerCAmelCase	:						List[str]=4		,	lowerCAmelCase	:						Union[str, Any]=37		,	lowerCAmelCase	:						List[str]="gelu"		,	lowerCAmelCase	:						Any=0.1		,	lowerCAmelCase	:						List[str]=0.1		,	lowerCAmelCase	:						Tuple=512		,	lowerCAmelCase	:						Dict=16		,	lowerCAmelCase	:						Tuple=2		,	lowerCAmelCase	:						Optional[Any]=0.02		,	lowerCAmelCase	:						List[Any]=3		,	lowerCAmelCase	:						int=4		,	lowerCAmelCase	:						Optional[int]=None		,	):


		'''simple docstring'''







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




	def       __A					(		self	:						Any							):


		'''simple docstring'''







		UpperCAmelCase_       				=			ids_tensor([self.batch_size, self.seq_length]		,	self.vocab_size							)

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

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

		UpperCAmelCase_       				=			self.get_config()

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




	def       __A					(		self	:						Optional[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	:						int		,	lowerCAmelCase	:						str		,	lowerCAmelCase	:						str		,	lowerCAmelCase	:						Dict		,	lowerCAmelCase	:						Union[str, Any]		,	lowerCAmelCase	:						Any		,	lowerCAmelCase	:						List[Any]							):


		'''simple docstring'''







		UpperCAmelCase_       				=			EsmModel(config=lowerCAmelCase							)
		model.to(lowerCAmelCase							)
		model.eval()
		UpperCAmelCase_       				=			model(lowerCAmelCase		,	attention_mask=lowerCAmelCase							)
		UpperCAmelCase_       				=			model(lowerCAmelCase							)
		UpperCAmelCase_       				=			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	:						Tuple		,	lowerCAmelCase	:						Optional[Any]		,	lowerCAmelCase	:						List[str]		,	lowerCAmelCase	:						int		,	lowerCAmelCase	:						str		,	lowerCAmelCase	:						Tuple		,	lowerCAmelCase	:						Tuple							):


		'''simple docstring'''







		UpperCAmelCase_       				=			EsmForMaskedLM(config=lowerCAmelCase							)
		model.to(lowerCAmelCase							)
		model.eval()
		UpperCAmelCase_       				=			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	:						Optional[Any]		,	lowerCAmelCase	:						str		,	lowerCAmelCase	:						Any		,	lowerCAmelCase	:						int		,	lowerCAmelCase	:						List[str]		,	lowerCAmelCase	:						List[str]		,	lowerCAmelCase	:						str							):


		'''simple docstring'''







		UpperCAmelCase_       				=			self.num_labels
		UpperCAmelCase_       				=			EsmForTokenClassification(config=lowerCAmelCase							)
		model.to(lowerCAmelCase							)
		model.eval()
		UpperCAmelCase_       				=			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	:						Any							):


		'''simple docstring'''







		UpperCAmelCase_       				=			self.prepare_config_and_inputs()
		(
		    (
		    UpperCAmelCase_   
		)				,   (
		    UpperCAmelCase_   
		)				,   (
		    UpperCAmelCase_   
		)				,   (
		    UpperCAmelCase_   
		)				,   (
		    UpperCAmelCase_   
		)				,   (
		    UpperCAmelCase_   
		)				,      
		)       				=			config_and_inputs
		UpperCAmelCase_       				=			{"input_ids": input_ids, "attention_mask": input_mask}
		return config, inputs_dict


@require_torch
class 						__UpperCamelCase    (					lowercase     ,     lowercase     ,     unittest.TestCase						):
	SCREAMING_SNAKE_CASE__    				=     False

	SCREAMING_SNAKE_CASE__    				=     (
	    (
	        EsmForMaskedLM,
	        EsmModel,
	        EsmForSequenceClassification,
	        EsmForTokenClassification,
	    )
	    if is_torch_available()
	    else ()
	)
	SCREAMING_SNAKE_CASE__    				=     ()
	SCREAMING_SNAKE_CASE__    				=     (
	    {
	        'feature-extraction': EsmModel,
	        'fill-mask': EsmForMaskedLM,
	        'text-classification': EsmForSequenceClassification,
	        'token-classification': EsmForTokenClassification,
	        'zero-shot': EsmForSequenceClassification,
	    }
	    if is_torch_available()
	    else {}
	)
	SCREAMING_SNAKE_CASE__    				=     True




	def       __A					(		self	:						List[str]							):


		'''simple docstring'''







		UpperCAmelCase_       				=			EsmModelTester(self							)
		UpperCAmelCase_       				=			ConfigTester(self		,	config_class=lowerCAmelCase		,	hidden_size=37							)




	def       __A					(		self	:						Union[str, Any]							):


		'''simple docstring'''







		self.config_tester.run_common_tests()




	def       __A					(		self	:						List[str]							):


		'''simple docstring'''







		UpperCAmelCase_       				=			self.model_tester.prepare_config_and_inputs()
		self.model_tester.create_and_check_model(*lowerCAmelCase							)




	def       __A					(		self	:						Optional[Any]							):


		'''simple docstring'''







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




	def       __A					(		self	:						int							):


		'''simple docstring'''







		UpperCAmelCase_       				=			self.model_tester.prepare_config_and_inputs()
		self.model_tester.create_and_check_for_masked_lm(*lowerCAmelCase							)




	def       __A					(		self	:						Tuple							):


		'''simple docstring'''







		UpperCAmelCase_       				=			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 ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
			UpperCAmelCase_       				=			EsmModel.from_pretrained(lowerCAmelCase							)
			self.assertIsNotNone(lowerCAmelCase							)




	def       __A					(		self	:						int							):


		'''simple docstring'''







		UpperCAmelCase_       				=			self.model_tester.prepare_config_and_inputs()[0]
		UpperCAmelCase_       				=			EsmEmbeddings(config=lowerCAmelCase							)

		UpperCAmelCase_       				=			torch.as_tensor([[12, 31, 13, model.padding_idx]]							)
		UpperCAmelCase_       				=			torch.as_tensor(
		    [
		        [
		            0 + model.padding_idx + 1,
		            1 + model.padding_idx + 1,
		            2 + model.padding_idx + 1,
		            model.padding_idx,
		        ]
		    ]							)
		UpperCAmelCase_       				=			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	:						Dict							):


		'''simple docstring'''







		UpperCAmelCase_       				=			self.model_tester.prepare_config_and_inputs()[0]
		UpperCAmelCase_       				=			EsmEmbeddings(config=lowerCAmelCase							)

		UpperCAmelCase_       				=			torch.empty(2		,	4		,	30							)
		UpperCAmelCase_       				=			[
		    0 + embeddings.padding_idx + 1,
		    1 + embeddings.padding_idx + 1,
		    2 + embeddings.padding_idx + 1,
		    3 + embeddings.padding_idx + 1,
		]
		UpperCAmelCase_       				=			torch.as_tensor([expected_single_positions, expected_single_positions]							)
		UpperCAmelCase_       				=			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							):


		'''simple docstring'''







		pass




	@unittest.skip("Esm does not support embedding resizing"							)
	def       __A					(		self	:						int							):


		'''simple docstring'''







		pass




	@unittest.skip("Will be fixed soon by reducing the size of the model used for common tests."							)
	def       __A					(		self	:						str							):


		'''simple docstring'''







		pass


@require_torch
class 						__UpperCamelCase    (					lowercase						):




	@slow
	def       __A					(		self	:						Optional[int]							):


		'''simple docstring'''







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

			UpperCAmelCase_       				=			33

			UpperCAmelCase_       				=			torch.Size((1, 6, vocab_size)							)
			self.assertEqual(output.shape		,	lowerCAmelCase							)

			UpperCAmelCase_       				=			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	:						Optional[Any]							):


		'''simple docstring'''







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

			UpperCAmelCase_       				=			torch.tensor([[0, 6, 4, 13, 5, 4, 16, 12, 11, 7, 2]]							)
			UpperCAmelCase_       				=			model(lowerCAmelCase							)[0]
			# compare the actual values for a slice.
			UpperCAmelCase_       				=			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							)							)


from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_torch_available,
)


_a:    List[str] 							=			{
    """configuration_encodec""": [
        """ENCODEC_PRETRAINED_CONFIG_ARCHIVE_MAP""",
        """EncodecConfig""",
    ],
    """feature_extraction_encodec""": ["""EncodecFeatureExtractor"""],
}

try:
				if not is_torch_available():
								raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
				pass
else:
				_a:    List[Any] 							=			[
				    """ENCODEC_PRETRAINED_MODEL_ARCHIVE_LIST""",
				    """EncodecModel""",
				    """EncodecPreTrainedModel""",
				]

if TYPE_CHECKING:
				from .configuration_encodec import (
				    ENCODEC_PRETRAINED_CONFIG_ARCHIVE_MAP,
				    EncodecConfig,
				)
				from .feature_extraction_encodec import EncodecFeatureExtractor

				try:
								if not is_torch_available():
												raise OptionalDependencyNotAvailable()
				except OptionalDependencyNotAvailable:
								pass
				else:
								from .modeling_encodec import (
								    ENCODEC_PRETRAINED_MODEL_ARCHIVE_LIST,
								    EncodecModel,
								    EncodecPreTrainedModel,
								)

else:
				import sys

				_a:    int 							=			_LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)