


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


A__						:   Dict				=    logging.get_logger(__name__)

A__						:   Union[str, Any]				=    {
    """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   UpperCAmelCase_    (_UpperCAmelCase    ):





    """simple docstring"""






    lowerCamelCase    :   str					=     'gpt_neo'
    lowerCamelCase    :   str					=     ['past_key_values']
    lowerCamelCase    :   Optional[int]					=     {'num_attention_heads': 'num_heads', 'num_hidden_layers': 'num_layers'}







    def __init__(				self     ,	SCREAMING_SNAKE_CASE_=5_02_57     ,	SCREAMING_SNAKE_CASE_=20_48     ,	SCREAMING_SNAKE_CASE_=20_48     ,	SCREAMING_SNAKE_CASE_=24     ,	SCREAMING_SNAKE_CASE_=[[["global", "local"], 12]]     ,	SCREAMING_SNAKE_CASE_=16     ,	SCREAMING_SNAKE_CASE_=None     ,	SCREAMING_SNAKE_CASE_=2_56     ,	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.0_2     ,	SCREAMING_SNAKE_CASE_=True     ,	SCREAMING_SNAKE_CASE_=5_02_56     ,	SCREAMING_SNAKE_CASE_=5_02_56     ,	**SCREAMING_SNAKE_CASE_     ,	)							->	Any:
         __lowerCamelCase  :      List[Any]														=						vocab_size
         __lowerCamelCase  :      Union[str, Any]														=						max_position_embeddings
         __lowerCamelCase  :      Union[str, Any]														=						hidden_size
         __lowerCamelCase  :      List[Any]														=						num_layers
         __lowerCamelCase  :      Optional[Any]														=						num_heads
         __lowerCamelCase  :      int														=						intermediate_size
         __lowerCamelCase  :      int														=						window_size
         __lowerCamelCase  :      Optional[Any]														=						activation_function
         __lowerCamelCase  :      List[str]														=						resid_dropout
         __lowerCamelCase  :      Any														=						embed_dropout
         __lowerCamelCase  :      Union[str, Any]														=						attention_dropout
         __lowerCamelCase  :      Union[str, Any]														=						classifier_dropout
         __lowerCamelCase  :      Optional[int]														=						layer_norm_epsilon
         __lowerCamelCase  :      int														=						initializer_range
         __lowerCamelCase  :      Union[str, Any]														=						use_cache

         __lowerCamelCase  :      Dict														=						bos_token_id
         __lowerCamelCase  :      List[Any]														=						eos_token_id

         __lowerCamelCase  :      List[str]														=						attention_types
         __lowerCamelCase  :      Dict														=						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  lowercase_							(				SCREAMING_SNAKE_CASE_      )							->	int:
         __lowerCamelCase  :      Dict														=						[]
         for item in attention_types:
              for _ in range(item[1]      ):
                   attentions.extend(item[0]      )
         return attentions







def      UpperCAmelCase__   (      UpperCAmelCase_       :							Tuple ,       UpperCAmelCase_       :							List[str] ,       UpperCAmelCase_       :							Optional[int] ,       UpperCAmelCase_       :							Union[str, Any]				)	->						Union[str, Any]:
     import torch

     __lowerCamelCase  :      Optional[Any]														=						input.size()
     __lowerCamelCase  :      Dict														=						len(UpperCAmelCase_				)
     __lowerCamelCase  :      Any														=						shape[dimension]

     __lowerCamelCase  :      Tuple														=						torch.arange(0 ,       UpperCAmelCase_ ,       UpperCAmelCase_				)
     __lowerCamelCase  :      Union[str, Any]														=						torch.div(sizedim - size ,       UpperCAmelCase_ ,       rounding_mode='floor'				) + 1
     __lowerCamelCase  :      int														=						torch.arange(UpperCAmelCase_				) + low_indices[:min_length][:, None]

     __lowerCamelCase  :      Tuple														=						[slice(UpperCAmelCase_				)] * rank
     __lowerCamelCase  :      Union[str, Any]														=						indices
     __lowerCamelCase  :      Optional[int]														=						input[s]

     __lowerCamelCase  :      Optional[Any]														=						list(range(0 ,       rank + 1				)				)
     perm.append(perm.pop(dimension + 1				)				)

     return sliced.permute(UpperCAmelCase_				)



def      UpperCAmelCase__   (      UpperCAmelCase_       :							List[Any] ,       UpperCAmelCase_       :							int				)	->						Dict:
     import torch

     __lowerCamelCase  :      Optional[Any]														=						torch.arange(1 ,       UpperCAmelCase_				)
     __lowerCamelCase  :      List[Any]														=						torch.remainder(UpperCAmelCase_ ,       UpperCAmelCase_				)
     __lowerCamelCase  :      List[Any]														=						remainders == 0
     __lowerCamelCase  :      Dict														=						candidates[divisor_indices]
     __lowerCamelCase  :      Dict														=						torch.max(UpperCAmelCase_				)
     return largest_divisor, torch.div(UpperCAmelCase_ ,       UpperCAmelCase_ ,       rounding_mode='floor'				)





class   UpperCAmelCase_    (_UpperCAmelCase    ):





    """simple docstring"""






    @property
    def  lowercase_							(				self      )							->	Mapping[str, Mapping[int, str]]:
         __lowerCamelCase  :      List[Any]														=						OrderedDict({'input_ids': {0: 'batch', 1: 'sequence'}}      )
         if self.use_past:
              self.fill_with_past_key_values_(SCREAMING_SNAKE_CASE_     ,	direction='inputs'      )
              __lowerCamelCase  :      Optional[int]														=						{0: 'batch', 1: 'past_sequence + sequence'}
         else:
              __lowerCamelCase  :      Any														=						{0: 'batch', 1: 'sequence'}

         return common_inputs







    @property
    def  lowercase_							(				self      )							->	int:
         return self._config.num_heads







    def  lowercase_							(				self     ,	SCREAMING_SNAKE_CASE_     ,	SCREAMING_SNAKE_CASE_ = -1     ,	SCREAMING_SNAKE_CASE_ = -1     ,	SCREAMING_SNAKE_CASE_ = False     ,	SCREAMING_SNAKE_CASE_ = None     ,	)							->	Mapping[str, Any]:
         __lowerCamelCase  :      List[Any]														=						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()
         __lowerCamelCase  :      str														=						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

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

         __lowerCamelCase  :      List[Any]														=						common_inputs['attention_mask']
         if self.use_past:
              __lowerCamelCase  :      Any														=						ordered_inputs['attention_mask'].dtype
              __lowerCamelCase  :      Union[str, Any]														=						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  lowercase_							(				self      )							->	int:
         return 13










'''simple docstring'''



from dataclasses import asdict, dataclass
from typing import Optional

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


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

# TODO Update this
A__						:   Tuple				=    {
    """facebook/esm-1b""": """https://huggingface.co/facebook/esm-1b/resolve/main/config.json""",
    # See all ESM models at https://huggingface.co/models?filter=esm
}





class   UpperCAmelCase_    (_UpperCAmelCase    ):





    """simple docstring"""






    lowerCamelCase    :   Tuple					=     'esm'







    def __init__(				self     ,	SCREAMING_SNAKE_CASE_=None     ,	SCREAMING_SNAKE_CASE_=None     ,	SCREAMING_SNAKE_CASE_=None     ,	SCREAMING_SNAKE_CASE_=7_68     ,	SCREAMING_SNAKE_CASE_=12     ,	SCREAMING_SNAKE_CASE_=12     ,	SCREAMING_SNAKE_CASE_=30_72     ,	SCREAMING_SNAKE_CASE_=0.1     ,	SCREAMING_SNAKE_CASE_=0.1     ,	SCREAMING_SNAKE_CASE_=10_26     ,	SCREAMING_SNAKE_CASE_=0.0_2     ,	SCREAMING_SNAKE_CASE_=1E-12     ,	SCREAMING_SNAKE_CASE_="absolute"     ,	SCREAMING_SNAKE_CASE_=True     ,	SCREAMING_SNAKE_CASE_=None     ,	SCREAMING_SNAKE_CASE_=False     ,	SCREAMING_SNAKE_CASE_=False     ,	SCREAMING_SNAKE_CASE_=None     ,	SCREAMING_SNAKE_CASE_=None     ,	**SCREAMING_SNAKE_CASE_     ,	)							->	List[str]:
         super().__init__(pad_token_id=SCREAMING_SNAKE_CASE_     ,	mask_token_id=SCREAMING_SNAKE_CASE_     ,	**SCREAMING_SNAKE_CASE_      )

         __lowerCamelCase  :      int														=						vocab_size
         __lowerCamelCase  :      List[Any]														=						hidden_size
         __lowerCamelCase  :      str														=						num_hidden_layers
         __lowerCamelCase  :      List[str]														=						num_attention_heads
         __lowerCamelCase  :      Any														=						intermediate_size
         __lowerCamelCase  :      Optional[Any]														=						hidden_dropout_prob
         __lowerCamelCase  :      Tuple														=						attention_probs_dropout_prob
         __lowerCamelCase  :      Optional[int]														=						max_position_embeddings
         __lowerCamelCase  :      str														=						initializer_range
         __lowerCamelCase  :      Optional[int]														=						layer_norm_eps
         __lowerCamelCase  :      List[str]														=						position_embedding_type
         __lowerCamelCase  :      int														=						use_cache
         __lowerCamelCase  :      Optional[Any]														=						emb_layer_norm_before
         __lowerCamelCase  :      Optional[Any]														=						token_dropout
         __lowerCamelCase  :      str														=						is_folding_model
         if is_folding_model:
              if esmfold_config is None:
                   logger.info('No esmfold_config supplied for folding model, using default values.'      )
                   __lowerCamelCase  :      Dict														=						EsmFoldConfig()
              elif isinstance(SCREAMING_SNAKE_CASE_     ,	SCREAMING_SNAKE_CASE_      ):
                   __lowerCamelCase  :      Optional[int]														=						EsmFoldConfig(**SCREAMING_SNAKE_CASE_      )
              __lowerCamelCase  :      List[Any]														=						esmfold_config
              if vocab_list is None:
                   logger.warning('No vocab_list supplied for folding model, assuming the ESM-2 vocabulary!'      )
                   __lowerCamelCase  :      List[str]														=						get_default_vocab_list()
              else:
                   __lowerCamelCase  :      Optional[Any]														=						vocab_list
         else:
              __lowerCamelCase  :      Dict														=						None
              __lowerCamelCase  :      Optional[Any]														=						None
         if self.esmfold_config is not None and getattr(self.esmfold_config     ,	'use_esm_attn_map'     ,	SCREAMING_SNAKE_CASE_      ):
              raise ValueError('The HuggingFace port of ESMFold does not support use_esm_attn_map at this time!'      )







    def  lowercase_							(				self      )							->	Any:
         __lowerCamelCase  :      Any														=						super().to_dict()
         if isinstance(self.esmfold_config     ,	SCREAMING_SNAKE_CASE_      ):
              __lowerCamelCase  :      int														=						self.esmfold_config.to_dict()
         return output







@dataclass
class   UpperCAmelCase_    :





    """simple docstring"""






    lowerCamelCase    :   str					=     None
    lowerCamelCase    :   bool					=     True
    lowerCamelCase    :   bool					=     False
    lowerCamelCase    :   bool					=     False
    lowerCamelCase    :   bool					=     False
    lowerCamelCase    :   float					=     0

    lowerCamelCase    :   bool					=     True
    lowerCamelCase    :   bool					=     False

    lowerCamelCase    :   int					=     1_2_8
    lowerCamelCase    :   "TrunkConfig"					=     None







    def  lowercase_							(				self      )							->	Any:
         if self.trunk is None:
              __lowerCamelCase  :      List[str]														=						TrunkConfig()
         elif isinstance(self.trunk     ,	SCREAMING_SNAKE_CASE_      ):
              __lowerCamelCase  :      Any														=						TrunkConfig(**self.trunk      )







    def  lowercase_							(				self      )							->	int:
         __lowerCamelCase  :      Optional[int]														=						asdict(self      )
         __lowerCamelCase  :      str														=						self.trunk.to_dict()
         return output







@dataclass
class   UpperCAmelCase_    :





    """simple docstring"""






    lowerCamelCase    :   int					=     4_8
    lowerCamelCase    :   int					=     1_0_2_4
    lowerCamelCase    :   int					=     1_2_8
    lowerCamelCase    :   int					=     3_2
    lowerCamelCase    :   int					=     3_2
    lowerCamelCase    :   int					=     3_2
    lowerCamelCase    :   float					=     0
    lowerCamelCase    :   float					=     0
    lowerCamelCase    :   bool					=     False
    lowerCamelCase    :   int					=     4
    lowerCamelCase    :   Optional[int]					=     1_2_8
    lowerCamelCase    :   "StructureModuleConfig"					=     None







    def  lowercase_							(				self      )							->	Optional[int]:
         if self.structure_module is None:
              __lowerCamelCase  :      Dict														=						StructureModuleConfig()
         elif isinstance(self.structure_module     ,	SCREAMING_SNAKE_CASE_      ):
              __lowerCamelCase  :      Optional[Any]														=						StructureModuleConfig(**self.structure_module      )

         if self.max_recycles <= 0:
              raise ValueError(f'`max_recycles` should be positive, got {self.max_recycles}.'      )
         if self.sequence_state_dim % self.sequence_state_dim != 0:
              raise ValueError(
                  '`sequence_state_dim` should be a round multiple of `sequence_state_dim`, got'
                  f' {self.sequence_state_dim} and {self.sequence_state_dim}.'      )
         if self.pairwise_state_dim % self.pairwise_state_dim != 0:
              raise ValueError(
                  '`pairwise_state_dim` should be a round multiple of `pairwise_state_dim`, got'
                  f' {self.pairwise_state_dim} and {self.pairwise_state_dim}.'      )

         __lowerCamelCase  :      Tuple														=						self.sequence_state_dim // self.sequence_head_width
         __lowerCamelCase  :      str														=						self.pairwise_state_dim // self.pairwise_head_width

         if self.sequence_state_dim != sequence_num_heads * self.sequence_head_width:
              raise ValueError(
                  '`sequence_state_dim` should be equal to `sequence_num_heads * sequence_head_width, got'
                  f' {self.sequence_state_dim} != {sequence_num_heads} * {self.sequence_head_width}.'      )
         if self.pairwise_state_dim != pairwise_num_heads * self.pairwise_head_width:
              raise ValueError(
                  '`pairwise_state_dim` should be equal to `pairwise_num_heads * pairwise_head_width, got'
                  f' {self.pairwise_state_dim} != {pairwise_num_heads} * {self.pairwise_head_width}.'      )
         if self.pairwise_state_dim % 2 != 0:
              raise ValueError(f'`pairwise_state_dim` should be even, got {self.pairwise_state_dim}.'      )

         if self.dropout >= 0.4:
              raise ValueError(f'`dropout` should not be greater than 0.4, got {self.dropout}.'      )







    def  lowercase_							(				self      )							->	List[Any]:
         __lowerCamelCase  :      List[str]														=						asdict(self      )
         __lowerCamelCase  :      int														=						self.structure_module.to_dict()
         return output







@dataclass
class   UpperCAmelCase_    :





    """simple docstring"""






    lowerCamelCase    :   int					=     3_8_4
    lowerCamelCase    :   int					=     1_2_8
    lowerCamelCase    :   int					=     1_6
    lowerCamelCase    :   int					=     1_2_8
    lowerCamelCase    :   int					=     1_2
    lowerCamelCase    :   int					=     4
    lowerCamelCase    :   int					=     8
    lowerCamelCase    :   float					=     0.1
    lowerCamelCase    :   int					=     8
    lowerCamelCase    :   int					=     1
    lowerCamelCase    :   int					=     2
    lowerCamelCase    :   int					=     7
    lowerCamelCase    :   int					=     1_0
    lowerCamelCase    :   float					=     1e-8
    lowerCamelCase    :   float					=     1e5







    def  lowercase_							(				self      )							->	Any:
         return asdict(self      )







def      UpperCAmelCase__   (      )	->						Optional[Any]:
     return (
         "<cls>",
         "<pad>",
         "<eos>",
         "<unk>",
         "L",
         "A",
         "G",
         "V",
         "S",
         "E",
         "R",
         "T",
         "I",
         "D",
         "P",
         "K",
         "Q",
         "N",
         "F",
         "Y",
         "M",
         "H",
         "W",
         "C",
         "X",
         "B",
         "U",
         "Z",
         "O",
         ".",
         "-",
         "<null_1>",
         "<mask>",
     )






