
'''simple docstring'''




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






         """simple docstring"""






         A       :   int					=						parent
         A       :   int					=						batch_size
         A       :   List[Any]					=						seq_length
         A       :   Union[str, Any]					=						is_training
         A       :   Optional[Any]					=						use_input_mask
         A       :   List[str]					=						use_token_type_ids
         A       :   str					=						use_labels
         A       :   Optional[Any]					=						vocab_size
         A       :   Union[str, Any]					=						hidden_size
         A       :   Tuple					=						num_hidden_layers
         A       :   Any					=						num_attention_heads
         A       :   Dict					=						intermediate_size
         A       :   Tuple					=						hidden_act
         A       :   Optional[Any]					=						hidden_dropout_prob
         A       :   List[str]					=						attention_probs_dropout_prob
         A       :   Any					=						max_position_embeddings
         A       :   str					=						type_vocab_size
         A       :   Dict					=						type_sequence_label_size
         A       :   Tuple					=						initializer_range
         A       :   Dict					=						num_labels
         A       :   Optional[Any]					=						num_choices
         A       :   List[str]					=						scope





     def 				__lowerCAmelCase							(							self	) ->				Any:






         """simple docstring"""






         A       :   int					=						ids_tensor([self.batch_size, self.seq_length]						,			self.vocab_size	)

         A       :   Dict					=						None
         if self.use_input_mask:
             A       :   Optional[Any]					=						random_attention_mask([self.batch_size, self.seq_length]	)

         A       :   Dict					=						None
         A       :   List[str]					=						None
         A       :   List[Any]					=						None
         if self.use_labels:
             A       :   Tuple					=						ids_tensor([self.batch_size]						,			self.type_sequence_label_size	)
             A       :   Dict					=						ids_tensor([self.batch_size, self.seq_length]						,			self.num_labels	)
             A       :   Optional[Any]					=						ids_tensor([self.batch_size]						,			self.num_choices	)

         A       :   Dict					=						self.get_config()

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





     def 				__lowerCAmelCase							(							self	) ->				str:






         """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 				__lowerCAmelCase							(							self						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE	) ->				Any:






         """simple docstring"""






         A       :   Any					=						EsmModel(config=SCREAMING_SNAKE_CASE	)
         model.to(SCREAMING_SNAKE_CASE	)
         model.eval()
         A       :   Optional[Any]					=						model(SCREAMING_SNAKE_CASE						,			attention_mask=SCREAMING_SNAKE_CASE	)
         A       :   int					=						model(SCREAMING_SNAKE_CASE	)
         A       :   List[str]					=						model(SCREAMING_SNAKE_CASE	)

         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 				__lowerCAmelCase							(							self						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE	) ->				Tuple:






         """simple docstring"""






         A       :   List[str]					=						EsmForMaskedLM(config=SCREAMING_SNAKE_CASE	)
         model.to(SCREAMING_SNAKE_CASE	)
         model.eval()
         A       :   Optional[Any]					=						model(SCREAMING_SNAKE_CASE						,			attention_mask=SCREAMING_SNAKE_CASE						,			labels=SCREAMING_SNAKE_CASE	)
         self.parent.assertEqual(result.logits.shape						,			(self.batch_size, self.seq_length, self.vocab_size)	)





     def 				__lowerCAmelCase							(							self						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE	) ->				Optional[Any]:






         """simple docstring"""






         A       :   Union[str, Any]					=						self.num_labels
         A       :   Union[str, Any]					=						EsmForTokenClassification(config=SCREAMING_SNAKE_CASE	)
         model.to(SCREAMING_SNAKE_CASE	)
         model.eval()
         A       :   List[Any]					=						model(SCREAMING_SNAKE_CASE						,			attention_mask=SCREAMING_SNAKE_CASE						,			labels=SCREAMING_SNAKE_CASE	)
         self.parent.assertEqual(result.logits.shape						,			(self.batch_size, self.seq_length, self.num_labels)	)





     def 				__lowerCAmelCase							(							self	) ->				Dict:






         """simple docstring"""






         A       :   Optional[int]					=						self.prepare_config_and_inputs()
         (
             (
             A
         ),  (
             A
         ),  (
             A
         ),  (
             A
         ),  (
             A
         ),  (
             A
         ),  
         )       :   Optional[int]					=						config_and_inputs
         A       :   Optional[Any]					=						{'''input_ids''': input_ids, '''attention_mask''': input_mask}
         return config, inputs_dict





@require_torch
class  A				(							__snake_case    ,			__snake_case    ,			unittest.TestCase						):
     __magic_name__						=		False

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





     def 				__lowerCAmelCase							(							self	) ->				Optional[int]:






         """simple docstring"""






         A       :   Any					=						EsmModelTester(self	)
         A       :   int					=						ConfigTester(self						,			config_class=SCREAMING_SNAKE_CASE						,			hidden_size=37	)





     def 				__lowerCAmelCase							(							self	) ->				Optional[Any]:






         """simple docstring"""






         self.config_tester.run_common_tests()





     def 				__lowerCAmelCase							(							self	) ->				Optional[Any]:






         """simple docstring"""






         A       :   List[str]					=						self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_model(*SCREAMING_SNAKE_CASE	)





     def 				__lowerCAmelCase							(							self	) ->				Dict:






         """simple docstring"""






         A       :   Tuple					=						self.model_tester.prepare_config_and_inputs()
         for type in ["absolute", "relative_key", "relative_key_query"]:
             A       :   Optional[Any]					=						type
             self.model_tester.create_and_check_model(*SCREAMING_SNAKE_CASE	)





     def 				__lowerCAmelCase							(							self	) ->				Optional[int]:






         """simple docstring"""






         A       :   Optional[Any]					=						self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_masked_lm(*SCREAMING_SNAKE_CASE	)





     def 				__lowerCAmelCase							(							self	) ->				Tuple:






         """simple docstring"""






         A       :   List[Any]					=						self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_token_classification(*SCREAMING_SNAKE_CASE	)





     @slow
     def 				__lowerCAmelCase							(							self	) ->				Dict:






         """simple docstring"""






         for model_name in ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
             A       :   Any					=						EsmModel.from_pretrained(SCREAMING_SNAKE_CASE	)
             self.assertIsNotNone(SCREAMING_SNAKE_CASE	)





     def 				__lowerCAmelCase							(							self	) ->				Union[str, Any]:






         """simple docstring"""






         A       :   List[Any]					=						self.model_tester.prepare_config_and_inputs()[0]
         A       :   Optional[Any]					=						EsmEmbeddings(config=SCREAMING_SNAKE_CASE	)

         A       :   Any					=						torch.as_tensor([[12, 31, 13, model.padding_idx]]	)
         A       :   Any					=						torch.as_tensor(
             [
                 [
                     0 + model.padding_idx + 1,
                     1 + model.padding_idx + 1,
                     2 + model.padding_idx + 1,
                     model.padding_idx,
                 ]
             ]	)
         A       :   int					=						create_position_ids_from_input_ids(SCREAMING_SNAKE_CASE						,			model.padding_idx	)
         self.assertEqual(position_ids.shape						,			expected_positions.shape	)
         self.assertTrue(torch.all(torch.eq(SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE	)	)	)





     def 				__lowerCAmelCase							(							self	) ->				str:






         """simple docstring"""






         A       :   List[Any]					=						self.model_tester.prepare_config_and_inputs()[0]
         A       :   int					=						EsmEmbeddings(config=SCREAMING_SNAKE_CASE	)

         A       :   Optional[Any]					=						torch.empty(2						,			4						,			30	)
         A       :   Any					=						[
             0 + embeddings.padding_idx + 1,
             1 + embeddings.padding_idx + 1,
             2 + embeddings.padding_idx + 1,
             3 + embeddings.padding_idx + 1,
         ]
         A       :   Any					=						torch.as_tensor([expected_single_positions, expected_single_positions]	)
         A       :   int					=						embeddings.create_position_ids_from_inputs_embeds(SCREAMING_SNAKE_CASE	)
         self.assertEqual(position_ids.shape						,			expected_positions.shape	)
         self.assertTrue(torch.all(torch.eq(SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE	)	)	)





     @unittest.skip('''Esm does not support embedding resizing'''	)
     def 				__lowerCAmelCase							(							self	) ->				int:






         """simple docstring"""






         pass





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






         """simple docstring"""






         pass





     @unittest.skip('''Will be fixed soon by reducing the size of the model used for common tests.'''	)
     def 				__lowerCAmelCase							(							self	) ->				Tuple:






         """simple docstring"""






         pass





@require_torch
class  A				(							__snake_case						):





     @slow
     def 				__lowerCAmelCase							(							self	) ->				Any:






         """simple docstring"""






         with torch.no_grad():
             A       :   List[str]					=						EsmForMaskedLM.from_pretrained('''facebook/esm2_t6_8M_UR50D'''	)
             model.eval()
             A       :   List[Any]					=						torch.tensor([[0, 1, 2, 3, 4, 5]]	)
             A       :   Union[str, Any]					=						model(SCREAMING_SNAKE_CASE	)[0]

             A       :   Optional[Any]					=						33

             A       :   List[str]					=						torch.Size((1, 6, vocab_size)	)
             self.assertEqual(output.shape						,			SCREAMING_SNAKE_CASE	)

             A       :   Tuple					=						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]						,			SCREAMING_SNAKE_CASE						,			atol=1e-4	)	)





     @slow
     def 				__lowerCAmelCase							(							self	) ->				str:






         """simple docstring"""






         with torch.no_grad():
             A       :   str					=						EsmModel.from_pretrained('''facebook/esm2_t6_8M_UR50D'''	)
             model.eval()

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






'''simple docstring'''




from collections import OrderedDict
from typing import Any, Mapping, Optional

from ... import PreTrainedTokenizer, TensorType, is_torch_available
from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfigWithPast
from ...utils import logging


lowercase			:				List[str]    		=			logging.get_logger(__name__)

lowercase			:				str    		=			{
    'EleutherAI/gpt-neo-1.3B': 'https://huggingface.co/EleutherAI/gpt-neo-1.3B/resolve/main/config.json',
    # See all GPTNeo models at https://huggingface.co/models?filter=gpt_neo
}





class  A				(							__snake_case						):
     __magic_name__						=		'''gpt_neo'''
     __magic_name__						=		['''past_key_values''']
     __magic_name__						=		{'''num_attention_heads''': '''num_heads''', '''num_hidden_layers''': '''num_layers'''}





     def __init__(							self						,			SCREAMING_SNAKE_CASE=50257						,			SCREAMING_SNAKE_CASE=2048						,			SCREAMING_SNAKE_CASE=2048						,			SCREAMING_SNAKE_CASE=24						,			SCREAMING_SNAKE_CASE=[[["global", "local"], 12]]						,			SCREAMING_SNAKE_CASE=16						,			SCREAMING_SNAKE_CASE=None						,			SCREAMING_SNAKE_CASE=256						,			SCREAMING_SNAKE_CASE="gelu_new"						,			SCREAMING_SNAKE_CASE=0.0						,			SCREAMING_SNAKE_CASE=0.0						,			SCREAMING_SNAKE_CASE=0.0						,			SCREAMING_SNAKE_CASE=0.1						,			SCREAMING_SNAKE_CASE=1e-5						,			SCREAMING_SNAKE_CASE=0.02						,			SCREAMING_SNAKE_CASE=True						,			SCREAMING_SNAKE_CASE=50256						,			SCREAMING_SNAKE_CASE=50256						,			**SCREAMING_SNAKE_CASE						,			) ->				int:






         """simple docstring"""






         A       :   Union[str, Any]					=						vocab_size
         A       :   Optional[Any]					=						max_position_embeddings
         A       :   Dict					=						hidden_size
         A       :   Optional[Any]					=						num_layers
         A       :   Tuple					=						num_heads
         A       :   int					=						intermediate_size
         A       :   Optional[Any]					=						window_size
         A       :   List[Any]					=						activation_function
         A       :   Union[str, Any]					=						resid_dropout
         A       :   Any					=						embed_dropout
         A       :   List[Any]					=						attention_dropout
         A       :   str					=						classifier_dropout
         A       :   List[Any]					=						layer_norm_epsilon
         A       :   str					=						initializer_range
         A       :   List[str]					=						use_cache

         A       :   Optional[int]					=						bos_token_id
         A       :   List[Any]					=						eos_token_id

         A       :   int					=						attention_types
         A       :   int					=						self.expand_attention_types_params(SCREAMING_SNAKE_CASE	)

         if len(self.attention_layers	) != self.num_layers:
             raise ValueError(
                 '''Configuration for convolutional module is incorrect. '''
                 '''It is required that `len(config.attention_layers)` == `config.num_layers` '''
                 F'but is `len(config.attention_layers) = {len(self.attention_layers	)}`, '
                 F'`config.num_layers = {self.num_layers}`. '
                 '''`config.attention_layers` is prepared using `config.attention_types`. '''
                 '''Please verify the value of `config.attention_types` argument.'''	)

         super().__init__(bos_token_id=SCREAMING_SNAKE_CASE						,			eos_token_id=SCREAMING_SNAKE_CASE						,			**SCREAMING_SNAKE_CASE	)





     @staticmethod
     def 				__lowerCAmelCase							(							SCREAMING_SNAKE_CASE	) ->				int:






         """simple docstring"""






         A       :   List[str]					=						[]
         for item in attention_types:
             for _ in range(item[1]	):
                 attentions.extend(item[0]	)
         return attentions







def 		lowerCAmelCase_				(   snake_case__		,				snake_case__		,				snake_case__		,				snake_case__						):




    '''simple docstring'''



    import torch

    A       :   Tuple					=						input.size()
    A       :   Union[str, Any]					=						len(snake_case__						)
    A       :   List[str]					=						shape[dimension]

    A       :   Union[str, Any]					=						torch.arange(0		,				snake_case__		,				snake_case__						)
    A       :   List[str]					=						torch.div(sizedim - size		,				snake_case__		,				rounding_mode='''floor'''						) + 1
    A       :   Optional[int]					=						torch.arange(snake_case__						) + low_indices[:min_length][:, None]

    A       :   str					=						[slice(snake_case__						)] * rank
    A       :   List[Any]					=						indices
    A       :   Union[str, Any]					=						input[s]

    A       :   List[str]					=						list(range(0		,				rank + 1						)						)
    perm.append(perm.pop(dimension + 1						)						)

    return sliced.permute(snake_case__						)







def 		lowerCAmelCase_				(   snake_case__		,				snake_case__						):




    '''simple docstring'''



    import torch

    A       :   List[str]					=						torch.arange(1		,				snake_case__						)
    A       :   Optional[int]					=						torch.remainder(snake_case__		,				snake_case__						)
    A       :   Optional[int]					=						remainders == 0
    A       :   Optional[Any]					=						candidates[divisor_indices]
    A       :   Optional[int]					=						torch.max(snake_case__						)
    return largest_divisor, torch.div(snake_case__		,				snake_case__		,				rounding_mode='''floor'''						)







class  A				(							__snake_case						):





     @property
     def 				__lowerCAmelCase							(							self	) ->				Mapping[str, Mapping[int, str]]:






         """simple docstring"""






         A       :   Tuple					=						OrderedDict({'''input_ids''': {0: '''batch''', 1: '''sequence'''}}	)
         if self.use_past:
             self.fill_with_past_key_values_(SCREAMING_SNAKE_CASE						,			direction='''inputs'''	)
             A       :   Optional[Any]					=						{0: '''batch''', 1: '''past_sequence + sequence'''}
         else:
             A       :   Dict					=						{0: '''batch''', 1: '''sequence'''}

         return common_inputs





     @property
     def 				__lowerCAmelCase							(							self	) ->				int:






         """simple docstring"""






         return self._config.num_heads





     def 				__lowerCAmelCase							(							self						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE = -1						,			SCREAMING_SNAKE_CASE = -1						,			SCREAMING_SNAKE_CASE = False						,			SCREAMING_SNAKE_CASE = None						,			) ->				Mapping[str, Any]:






         """simple docstring"""






         A       :   List[str]					=						super(SCREAMING_SNAKE_CASE						,			self	).generate_dummy_inputs(
             SCREAMING_SNAKE_CASE						,			batch_size=SCREAMING_SNAKE_CASE						,			seq_length=SCREAMING_SNAKE_CASE						,			is_pair=SCREAMING_SNAKE_CASE						,			framework=SCREAMING_SNAKE_CASE	)

         # We need to order the input in the way they appears in the forward()
         A       :   Any					=						OrderedDict({'''input_ids''': common_inputs['''input_ids''']}	)

         # Need to add the past_keys
         if self.use_past:
             if not is_torch_available():
                 raise ValueError('''Cannot generate dummy past_keys inputs without PyTorch installed.'''	)
             else:
                 import torch

                 A,  A       :   Dict					=						common_inputs['''input_ids'''].shape
                 # Not using the same length for past_key_values
                 A       :   str					=						seqlen + 2
                 A       :   List[Any]					=						(
                     batch,
                     self.num_attention_heads,
                     past_key_values_length,
                     self._config.hidden_size // self.num_attention_heads,
                 )
                 A       :   Any					=						[
                     (torch.zeros(SCREAMING_SNAKE_CASE	), torch.zeros(SCREAMING_SNAKE_CASE	)) for _ in range(self.num_layers	)
                 ]

         A       :   str					=						common_inputs['''attention_mask''']
         if self.use_past:
             A       :   Optional[int]					=						ordered_inputs['''attention_mask'''].dtype
             A       :   List[str]					=						torch.cat(
                 [ordered_inputs['''attention_mask'''], torch.ones(SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			dtype=SCREAMING_SNAKE_CASE	)]						,			dim=1	)

         return ordered_inputs





     @property
     def 				__lowerCAmelCase							(							self	) ->				int:






         """simple docstring"""






         return 13




