


from dataclasses import asdict, dataclass
from typing import Optional

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


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

# TODO Update this
UpperCAmelCase_   :				Union[str, Any]						      = {
    '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 				SCREAMING_SNAKE_CASE__						(       lowercase__    ):
      snake_case__       :  Dict       =		'''esm'''
      def __init__( self       : List[str]  ,							SCREAMING_SNAKE_CASE__       : List[str]=None  ,							SCREAMING_SNAKE_CASE__       : Optional[Any]=None  ,							SCREAMING_SNAKE_CASE__       : str=None  ,							SCREAMING_SNAKE_CASE__       : Tuple=7_6_8  ,							SCREAMING_SNAKE_CASE__       : Dict=1_2  ,							SCREAMING_SNAKE_CASE__       : Any=1_2  ,							SCREAMING_SNAKE_CASE__       : List[Any]=3_0_7_2  ,							SCREAMING_SNAKE_CASE__       : Dict=0.1  ,							SCREAMING_SNAKE_CASE__       : Union[str, Any]=0.1  ,							SCREAMING_SNAKE_CASE__       : List[str]=1_0_2_6  ,							SCREAMING_SNAKE_CASE__       : Any=0.02  ,							SCREAMING_SNAKE_CASE__       : Dict=1E-12  ,							SCREAMING_SNAKE_CASE__       : Optional[int]="absolute"  ,							SCREAMING_SNAKE_CASE__       : int=True  ,							SCREAMING_SNAKE_CASE__       : List[str]=None  ,							SCREAMING_SNAKE_CASE__       : Tuple=False  ,							SCREAMING_SNAKE_CASE__       : int=False  ,							SCREAMING_SNAKE_CASE__       : Union[str, Any]=None  ,							SCREAMING_SNAKE_CASE__       : Dict=None  ,							**SCREAMING_SNAKE_CASE__       : str  ,							)	->    Any:
             super().__init__(pad_token_id=SCREAMING_SNAKE_CASE__  ,							mask_token_id=SCREAMING_SNAKE_CASE__  ,							**SCREAMING_SNAKE_CASE__  )

             a_      :	List[Any]							=						vocab_size
             a_      :	List[Any]							=						hidden_size
             a_      :	str							=						num_hidden_layers
             a_      :	Optional[int]							=						num_attention_heads
             a_      :	Tuple							=						intermediate_size
             a_      :	Any							=						hidden_dropout_prob
             a_      :	List[str]							=						attention_probs_dropout_prob
             a_      :	Optional[int]							=						max_position_embeddings
             a_      :	List[Any]							=						initializer_range
             a_      :	Union[str, Any]							=						layer_norm_eps
             a_      :	Any							=						position_embedding_type
             a_      :	Tuple							=						use_cache
             a_      :	List[Any]							=						emb_layer_norm_before
             a_      :	Tuple							=						token_dropout
             a_      :	Optional[int]							=						is_folding_model
             if is_folding_model:
                    if esmfold_config is None:
                           logger.info('No esmfold_config supplied for folding model, using default values.'  )
                           a_      :	Optional[int]							=						EsmFoldConfig()
                    elif isinstance(SCREAMING_SNAKE_CASE__  ,							SCREAMING_SNAKE_CASE__  ):
                           a_      :	Tuple							=						EsmFoldConfig(**SCREAMING_SNAKE_CASE__  )
                    a_      :	str							=						esmfold_config
                    if vocab_list is None:
                           logger.warning('No vocab_list supplied for folding model, assuming the ESM-2 vocabulary!'  )
                           a_      :	Union[str, Any]							=						get_default_vocab_list()
                    else:
                           a_      :	Union[str, Any]							=						vocab_list
             else:
                    a_      :	Dict							=						None
                    a_      :	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  SCREAMING_SNAKE_CASE							( self       : int  )	->    Optional[Any]:
             a_      :	Optional[Any]							=						super().to_dict()
             if isinstance(self.esmfold_config  ,							SCREAMING_SNAKE_CASE__  ):
                    a_      :	List[Any]							=						self.esmfold_config.to_dict()
             return output
@dataclass
class 				SCREAMING_SNAKE_CASE__						:
      snake_case__       :  str       =		None
      snake_case__       :  bool       =		True
      snake_case__       :  bool       =		False
      snake_case__       :  bool       =		False
      snake_case__       :  bool       =		False
      snake_case__       :  float       =		0

      snake_case__       :  bool       =		True
      snake_case__       :  bool       =		False

      snake_case__       :  int       =		128
      snake_case__       :  "TrunkConfig"       =		None
      def  SCREAMING_SNAKE_CASE							( self       : Union[str, Any]  )	->    Tuple:
             if self.trunk is None:
                    a_      :	List[Any]							=						TrunkConfig()
             elif isinstance(self.trunk  ,							SCREAMING_SNAKE_CASE__  ):
                    a_      :	Optional[int]							=						TrunkConfig(**self.trunk  )

      def  SCREAMING_SNAKE_CASE							( self       : int  )	->    Any:
             a_      :	str							=						asdict(self  )
             a_      :	Optional[int]							=						self.trunk.to_dict()
             return output
@dataclass
class 				SCREAMING_SNAKE_CASE__						:
      snake_case__       :  int       =		48
      snake_case__       :  int       =		1_024
      snake_case__       :  int       =		128
      snake_case__       :  int       =		32
      snake_case__       :  int       =		32
      snake_case__       :  int       =		32
      snake_case__       :  float       =		0
      snake_case__       :  float       =		0
      snake_case__       :  bool       =		False
      snake_case__       :  int       =		4
      snake_case__       :  Optional[int]       =		128
      snake_case__       :  "StructureModuleConfig"       =		None
      def  SCREAMING_SNAKE_CASE							( self       : str  )	->    List[str]:
             if self.structure_module is None:
                    a_      :	Optional[Any]							=						StructureModuleConfig()
             elif isinstance(self.structure_module  ,							SCREAMING_SNAKE_CASE__  ):
                    a_      :	Union[str, 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}."""  )

             a_      :	Optional[int]							=						self.sequence_state_dim // self.sequence_head_width
             a_      :	Optional[Any]							=						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  SCREAMING_SNAKE_CASE							( self       : Tuple  )	->    List[str]:
             a_      :	Tuple							=						asdict(self  )
             a_      :	Tuple							=						self.structure_module.to_dict()
             return output
@dataclass
class 				SCREAMING_SNAKE_CASE__						:
      snake_case__       :  int       =		384
      snake_case__       :  int       =		128
      snake_case__       :  int       =		16
      snake_case__       :  int       =		128
      snake_case__       :  int       =		12
      snake_case__       :  int       =		4
      snake_case__       :  int       =		8
      snake_case__       :  float       =		0.1
      snake_case__       :  int       =		8
      snake_case__       :  int       =		1
      snake_case__       :  int       =		2
      snake_case__       :  int       =		7
      snake_case__       :  int       =		10
      snake_case__       :  float       =		1e-8
      snake_case__       :  float       =		1e5
      def  SCREAMING_SNAKE_CASE							( self       : Optional[Any]  )	->    Optional[int]:
             return asdict(self  )







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

       """simple docstring"""

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






def  SCREAMING_SNAKE_CASE_		(    __A				:    str       )			->	list:

       """simple docstring"""

       return [
           txt[:a] + txt[a].upper() + txt[a + 1 :]
           for a in range(len(__A       )       )
           if txt[a].isalpha()
       ]


if __name__ == "__main__":
   __import__('doctest').testmod()


