



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			,	UpperCAmelCase__			,	UpperCAmelCase__=13			,	UpperCAmelCase__=7			,	UpperCAmelCase__=False			,	UpperCAmelCase__=True			,	UpperCAmelCase__=False			,	UpperCAmelCase__=True			,	UpperCAmelCase__=33			,	UpperCAmelCase__=32			,	UpperCAmelCase__=5			,	UpperCAmelCase__=4			,	UpperCAmelCase__=37			,	UpperCAmelCase__="gelu"			,	UpperCAmelCase__=0.1			,	UpperCAmelCase__=0.1			,	UpperCAmelCase__=512			,	UpperCAmelCase__=16			,	UpperCAmelCase__=2			,	UpperCAmelCase__=0.02			,	UpperCAmelCase__=3			,	UpperCAmelCase__=4			,	UpperCAmelCase__=None			,	):
							A__							=       parent
							A__							=       batch_size
							A__							=       seq_length
							A__							=       is_training
							A__							=       use_input_mask
							A__							=       use_token_type_ids
							A__							=       use_labels
							A__							=       vocab_size
							A__							=       hidden_size
							A__							=       num_hidden_layers
							A__							=       num_attention_heads
							A__							=       intermediate_size
							A__							=       hidden_act
							A__							=       hidden_dropout_prob
							A__							=       attention_probs_dropout_prob
							A__							=       max_position_embeddings
							A__							=       type_vocab_size
							A__							=       type_sequence_label_size
							A__							=       initializer_range
							A__							=       num_labels
							A__							=       num_choices
							A__							=       scope



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

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



	def 			__A							(       self		):
							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			,	UpperCAmelCase__			,	UpperCAmelCase__			,	UpperCAmelCase__			,	UpperCAmelCase__			,	UpperCAmelCase__			,	UpperCAmelCase__		):
							A__							=       EsmModel(config=UpperCAmelCase__		)
							model.to(UpperCAmelCase__		)
							model.eval()
							A__							=       model(UpperCAmelCase__			,	attention_mask=UpperCAmelCase__		)
							A__							=       model(UpperCAmelCase__		)
							A__							=       model(UpperCAmelCase__		)

							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			,	UpperCAmelCase__			,	UpperCAmelCase__			,	UpperCAmelCase__			,	UpperCAmelCase__			,	UpperCAmelCase__			,	UpperCAmelCase__		):
							A__							=       EsmForMaskedLM(config=UpperCAmelCase__		)
							model.to(UpperCAmelCase__		)
							model.eval()
							A__							=       model(UpperCAmelCase__			,	attention_mask=UpperCAmelCase__			,	labels=UpperCAmelCase__		)
							self.parent.assertEqual(result.logits.shape			,	(self.batch_size, self.seq_length, self.vocab_size)		)



	def 			__A							(       self			,	UpperCAmelCase__			,	UpperCAmelCase__			,	UpperCAmelCase__			,	UpperCAmelCase__			,	UpperCAmelCase__			,	UpperCAmelCase__		):
							A__							=       self.num_labels
							A__							=       EsmForTokenClassification(config=UpperCAmelCase__		)
							model.to(UpperCAmelCase__		)
							model.eval()
							A__							=       model(UpperCAmelCase__			,	attention_mask=UpperCAmelCase__			,	labels=UpperCAmelCase__		)
							self.parent.assertEqual(result.logits.shape			,	(self.batch_size, self.seq_length, self.num_labels)		)



	def 			__A							(       self		):
							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_torch
class      UpperCamelCase    (						_UpperCAmelCase       ,    _UpperCAmelCase       ,    unittest.TestCase     ):
	lowerCAmelCase	:   Union[str, Any]				=	False

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



	def 			__A							(       self		):
							A__							=       EsmModelTester(self		)
							A__							=       ConfigTester(self			,	config_class=UpperCAmelCase__			,	hidden_size=37		)



	def 			__A							(       self		):
							self.config_tester.run_common_tests()



	def 			__A							(       self		):
							A__							=       self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_model(*UpperCAmelCase__		)



	def 			__A							(       self		):
							A__							=       self.model_tester.prepare_config_and_inputs()
							for type in ["absolute", "relative_key", "relative_key_query"]:
													A__							=       type
													self.model_tester.create_and_check_model(*UpperCAmelCase__		)



	def 			__A							(       self		):
							A__							=       self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_masked_lm(*UpperCAmelCase__		)



	def 			__A							(       self		):
							A__							=       self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_token_classification(*UpperCAmelCase__		)



	@slow
	def 			__A							(       self		):
							for model_name in ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
													A__							=       EsmModel.from_pretrained(UpperCAmelCase__		)
													self.assertIsNotNone(UpperCAmelCase__		)



	def 			__A							(       self		):
							A__							=       self.model_tester.prepare_config_and_inputs()[0]
							A__							=       EsmEmbeddings(config=UpperCAmelCase__		)

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



	def 			__A							(       self		):
							A__							=       self.model_tester.prepare_config_and_inputs()[0]
							A__							=       EsmEmbeddings(config=UpperCAmelCase__		)

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



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



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



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






@require_torch
class      UpperCamelCase    (						_UpperCAmelCase     ):



	@slow
	def 			__A							(       self		):
							with torch.no_grad():
													A__							=       EsmForMaskedLM.from_pretrained("facebook/esm2_t6_8M_UR50D"		)
													model.eval()
													A__							=       torch.tensor([[0, 1, 2, 3, 4, 5]]		)
													A__							=       model(UpperCAmelCase__		)[0]

													A__							=       33

													A__							=       torch.Size((1, 6, vocab_size)		)
													self.assertEqual(output.shape			,	UpperCAmelCase__		)

													A__							=       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]			,	UpperCAmelCase__			,	atol=1e-4		)		)



	@slow
	def 			__A							(       self		):
							with torch.no_grad():
													A__							=       EsmModel.from_pretrained("facebook/esm2_t6_8M_UR50D"		)
													model.eval()

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










import math
from enum import Enum
from typing import Optional, Union

from torch.optim import Optimizer
from torch.optim.lr_scheduler import LambdaLR

from .utils import logging


UpperCAmelCase_							:      Optional[Any]      =  logging.get_logger(__name__)






class      UpperCamelCase    (						_UpperCAmelCase     ):
	lowerCAmelCase	:   List[str]				=	"""linear"""
	lowerCAmelCase	:   int				=	"""cosine"""
	lowerCAmelCase	:   Dict				=	"""cosine_with_restarts"""
	lowerCAmelCase	:   Optional[Any]				=	"""polynomial"""
	lowerCAmelCase	:   Dict				=	"""constant"""
	lowerCAmelCase	:   Any				=	"""constant_with_warmup"""
	lowerCAmelCase	:   Union[str, Any]				=	"""piecewise_constant"""





def  UpperCamelCase						(       _A			:				Optimizer			,							_A			:				int = -1    )->       Dict:





						"""simple docstring"""





						return LambdaLR(_A			,							lambda _A    : 1			,							last_epoch=_A    )
def  UpperCamelCase						(       _A			:				Optimizer			,							_A			:				int			,							_A			:				int = -1    )->       Optional[Any]:





						"""simple docstring"""





						def lr_lambda(_A			:				int    ):
												if current_step < num_warmup_steps:
																		return float(_A    ) / float(max(1.0			,							_A    )    )
												return 1.0

						return LambdaLR(_A			,							_A			,							last_epoch=_A    )
def  UpperCamelCase						(       _A			:				Optimizer			,							_A			:				str			,							_A			:				int = -1    )->       Dict:





						"""simple docstring"""





						A__							=       {}
						A__							=       step_rules.split(","    )
						for rule_str in rule_list[:-1]:
												A__						,       A__							=       rule_str.split(":"    )
												A__							=       int(_A    )
												A__							=       float(_A    )
												A__							=       value
						A__							=       float(rule_list[-1]    )

						def create_rules_function(_A			:				Any			,							_A			:				Optional[int]    ):
												def rule_func(_A			:				int    ) -> float:
																		A__							=       sorted(rules_dict.keys()    )
																		for i, sorted_step in enumerate(_A    ):
																								if steps < sorted_step:
																														return rules_dict[sorted_steps[i]]
																		return last_lr_multiple

												return rule_func

						A__							=       create_rules_function(_A			,							_A    )

						return LambdaLR(_A			,							_A			,							last_epoch=_A    )
def  UpperCamelCase						(       _A			:				Any			,							_A			:				Union[str, Any]			,							_A			:				str			,							_A			:				str=-1    )->       Tuple:





						"""simple docstring"""





						def lr_lambda(_A			:				int    ):
												if current_step < num_warmup_steps:
																		return float(_A    ) / float(max(1			,							_A    )    )
												return max(
												    0.0			,							float(num_training_steps - current_step    ) / float(max(1			,							num_training_steps - num_warmup_steps    )    )    )

						return LambdaLR(_A			,							_A			,							_A    )
def  UpperCamelCase						(       _A			:				Optimizer			,							_A			:				int			,							_A			:				int			,							_A			:				float = 0.5			,							_A			:				int = -1    )->       Any:





						"""simple docstring"""





						def lr_lambda(_A			:				Tuple    ):
												if current_step < num_warmup_steps:
																		return float(_A    ) / float(max(1			,							_A    )    )
												A__							=       float(current_step - num_warmup_steps    ) / float(max(1			,							num_training_steps - num_warmup_steps    )    )
												return max(0.0			,							0.5 * (1.0 + math.cos(math.pi * float(_A    ) * 2.0 * progress    ))    )

						return LambdaLR(_A			,							_A			,							_A    )
def  UpperCamelCase						(       _A			:				Optimizer			,							_A			:				int			,							_A			:				int			,							_A			:				int = 1			,							_A			:				int = -1    )->       Any:





						"""simple docstring"""





						def lr_lambda(_A			:				Tuple    ):
												if current_step < num_warmup_steps:
																		return float(_A    ) / float(max(1			,							_A    )    )
												A__							=       float(current_step - num_warmup_steps    ) / float(max(1			,							num_training_steps - num_warmup_steps    )    )
												if progress >= 1.0:
																		return 0.0
												return max(0.0			,							0.5 * (1.0 + math.cos(math.pi * ((float(_A    ) * progress) % 1.0)    ))    )

						return LambdaLR(_A			,							_A			,							_A    )
def  UpperCamelCase						(       _A			:				Union[str, Any]			,							_A			:				Union[str, Any]			,							_A			:				List[str]			,							_A			:				Tuple=1E-7			,							_A			:				Dict=1.0			,							_A			:				Union[str, Any]=-1    )->       Any:





						"""simple docstring"""





						A__							=       optimizer.defaults["lr"]
						if not (lr_init > lr_end):
												raise ValueError(f"""lr_end ({lr_end}) must be be smaller than initial lr ({lr_init})"""    )

						def lr_lambda(_A			:				int    ):
												if current_step < num_warmup_steps:
																		return float(_A    ) / float(max(1			,							_A    )    )
												elif current_step > num_training_steps:
																		return lr_end / lr_init  # as LambdaLR multiplies by lr_init
												else:
																		A__							=       lr_init - lr_end
																		A__							=       num_training_steps - num_warmup_steps
																		A__							=       1 - (current_step - num_warmup_steps) / decay_steps
																		A__							=       lr_range * pct_remaining**power + lr_end
																		return decay / lr_init  # as LambdaLR multiplies by lr_init

						return LambdaLR(_A			,							_A			,							_A    )


UpperCAmelCase_							:      Any      =  {
    SchedulerType.LINEAR: get_linear_schedule_with_warmup,
    SchedulerType.COSINE: get_cosine_schedule_with_warmup,
    SchedulerType.COSINE_WITH_RESTARTS: get_cosine_with_hard_restarts_schedule_with_warmup,
    SchedulerType.POLYNOMIAL: get_polynomial_decay_schedule_with_warmup,
    SchedulerType.CONSTANT: get_constant_schedule,
    SchedulerType.CONSTANT_WITH_WARMUP: get_constant_schedule_with_warmup,
    SchedulerType.PIECEWISE_CONSTANT: get_piecewise_constant_schedule,
}
def  UpperCamelCase						(       _A			:				Union[str, SchedulerType]			,							_A			:				Optimizer			,							_A			:				Optional[str] = None			,							_A			:				Optional[int] = None			,							_A			:				Optional[int] = None			,							_A			:				int = 1			,							_A			:				float = 1.0			,							_A			:				int = -1			,							)->       Union[str, Any]:





						"""simple docstring"""





						A__							=       SchedulerType(_A    )
						A__							=       TYPE_TO_SCHEDULER_FUNCTION[name]
						if name == SchedulerType.CONSTANT:
												return schedule_func(_A			,							last_epoch=_A    )

						if name == SchedulerType.PIECEWISE_CONSTANT:
												return schedule_func(_A			,							step_rules=_A			,							last_epoch=_A    )

						# All other schedulers require `num_warmup_steps`
						if num_warmup_steps is None:
												raise ValueError(f"""{name} requires `num_warmup_steps`, please provide that argument."""    )

						if name == SchedulerType.CONSTANT_WITH_WARMUP:
												return schedule_func(_A			,							num_warmup_steps=_A			,							last_epoch=_A    )

						# All other schedulers require `num_training_steps`
						if num_training_steps is None:
												raise ValueError(f"""{name} requires `num_training_steps`, please provide that argument."""    )

						if name == SchedulerType.COSINE_WITH_RESTARTS:
												return schedule_func(
												    _A			,							num_warmup_steps=_A			,							num_training_steps=_A			,							num_cycles=_A			,							last_epoch=_A			,							)

						if name == SchedulerType.POLYNOMIAL:
												return schedule_func(
												    _A			,							num_warmup_steps=_A			,							num_training_steps=_A			,							power=_A			,							last_epoch=_A			,							)

						return schedule_func(
						    _A			,							num_warmup_steps=_A			,							num_training_steps=_A			,							last_epoch=_A    )





