



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		)		)










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


UpperCAmelCase_							:      Dict      =  logging.get_logger(__name__)

UpperCAmelCase_							:      List[Any]      =  {
    "EleutherAI/gpt-neox-20b": "https://huggingface.co/EleutherAI/gpt-neox-20b/resolve/main/config.json",
    # See all GPTNeoX models at https://huggingface.co/models?filter=gpt_neox
}






class      UpperCamelCase    (						_UpperCAmelCase     ):
	lowerCAmelCase	:   int				=	"""gpt_neox"""



	def __init__(       self			,	UpperCAmelCase__=50_432			,	UpperCAmelCase__=6_144			,	UpperCAmelCase__=44			,	UpperCAmelCase__=64			,	UpperCAmelCase__=24_576			,	UpperCAmelCase__="gelu"			,	UpperCAmelCase__=0.25			,	UpperCAmelCase__=10_000			,	UpperCAmelCase__=0.0			,	UpperCAmelCase__=0.0			,	UpperCAmelCase__=0.1			,	UpperCAmelCase__=2_048			,	UpperCAmelCase__=0.02			,	UpperCAmelCase__=1e-5			,	UpperCAmelCase__=True			,	UpperCAmelCase__=0			,	UpperCAmelCase__=2			,	UpperCAmelCase__=False			,	UpperCAmelCase__=True			,	UpperCAmelCase__=None			,	**UpperCAmelCase__			,	):
							super().__init__(bos_token_id=UpperCAmelCase__			,	eos_token_id=UpperCAmelCase__			,	**UpperCAmelCase__		)
							A__							=       vocab_size
							A__							=       max_position_embeddings
							A__							=       hidden_size
							A__							=       num_hidden_layers
							A__							=       num_attention_heads
							A__							=       intermediate_size
							A__							=       hidden_act
							A__							=       rotary_pct
							A__							=       rotary_emb_base
							A__							=       attention_dropout
							A__							=       hidden_dropout
							A__							=       classifier_dropout
							A__							=       initializer_range
							A__							=       layer_norm_eps
							A__							=       use_cache
							A__							=       tie_word_embeddings
							A__							=       use_parallel_residual
							A__							=       rope_scaling
							self._rope_scaling_validation()

							if self.hidden_size % self.num_attention_heads != 0:
													raise ValueError(
													    "The hidden size is not divisble by the number of attention heads! Make sure to update them!"		)



	def 			__A							(       self		):
							if self.rope_scaling is None:
													return

							if not isinstance(self.rope_scaling			,	UpperCAmelCase__		) or len(self.rope_scaling		) != 2:
													raise ValueError(
													    "`rope_scaling` must be a dictionary with with two fields, `name` and `factor`, "
													    F"""got {self.rope_scaling}"""		)
							A__							=       self.rope_scaling.get("type"			,	UpperCAmelCase__		)
							A__							=       self.rope_scaling.get("factor"			,	UpperCAmelCase__		)
							if rope_scaling_type is None or rope_scaling_type not in ["linear", "dynamic"]:
													raise ValueError(
													    F"""`rope_scaling`'s name field must be one of ['linear', 'dynamic'], got {rope_scaling_type}"""		)
							if rope_scaling_factor is None or not isinstance(UpperCAmelCase__			,	UpperCAmelCase__		) or rope_scaling_factor <= 1.0:
													raise ValueError(F"""`rope_scaling`'s factor field must be an float > 1, got {rope_scaling_factor}"""		)





