




from dataclasses import asdict, dataclass
from typing import Optional

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


UpperCAmelCase_		:					Union[str, Any]          =  logging.get_logger(__name__)

# TODO Update this
UpperCAmelCase_		:					Optional[int]          =  {
    "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 		__A   (      UpperCamelCase__      ):
						UpperCamelCase      =       """esm"""

						def __init__(		self :List[Any]     ,			__snake_case :int=None     ,			__snake_case :List[Any]=None     ,			__snake_case :Tuple=None     ,			__snake_case :List[str]=7_68     ,			__snake_case :str=12     ,			__snake_case :int=12     ,			__snake_case :str=30_72     ,			__snake_case :int=0.1     ,			__snake_case :Optional[Any]=0.1     ,			__snake_case :Optional[Any]=10_26     ,			__snake_case :Dict=0.02     ,			__snake_case :List[str]=1E-12     ,			__snake_case :Any="absolute"     ,			__snake_case :Tuple=True     ,			__snake_case :Any=None     ,			__snake_case :Dict=False     ,			__snake_case :List[str]=False     ,			__snake_case :Tuple=None     ,			__snake_case :List[str]=None     ,			**__snake_case :Any     ,			):
										'''simple docstring'''

										super().__init__(pad_token_id=__snake_case     ,			mask_token_id=__snake_case     ,			**__snake_case			)

										__magic_name__						:						Dict										=vocab_size
										__magic_name__						:						Optional[Any]										=hidden_size
										__magic_name__						:						Any										=num_hidden_layers
										__magic_name__						:						int										=num_attention_heads
										__magic_name__						:						Any										=intermediate_size
										__magic_name__						:						List[Any]										=hidden_dropout_prob
										__magic_name__						:						str										=attention_probs_dropout_prob
										__magic_name__						:						Union[str, Any]										=max_position_embeddings
										__magic_name__						:						List[str]										=initializer_range
										__magic_name__						:						Optional[Any]										=layer_norm_eps
										__magic_name__						:						Optional[int]										=position_embedding_type
										__magic_name__						:						int										=use_cache
										__magic_name__						:						Any										=emb_layer_norm_before
										__magic_name__						:						List[str]										=token_dropout
										__magic_name__						:						List[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."""			)
																		__magic_name__						:						Tuple										=EsmFoldConfig()
														elif isinstance(__snake_case     ,			__snake_case			):
																		__magic_name__						:						Any										=EsmFoldConfig(**__snake_case			)
														__magic_name__						:						Dict										=esmfold_config
														if vocab_list is None:
																		logger.warning("""No vocab_list supplied for folding model, assuming the ESM-2 vocabulary!"""			)
																		__magic_name__						:						List[Any]										=get_default_vocab_list()
														else:
																		__magic_name__						:						Any										=vocab_list
										else:
														__magic_name__						:						Optional[Any]										=None
														__magic_name__						:						Optional[Any]										=None
										if self.esmfold_config is not None and getattr(self.esmfold_config     ,			"""use_esm_attn_map"""     ,			__snake_case			):
														raise ValueError("""The HuggingFace port of ESMFold does not support use_esm_attn_map at this time!"""			)


						def      A__    (		self :Union[str, Any]			):
										'''simple docstring'''

										__magic_name__						:						Optional[Any]										=super().to_dict()
										if isinstance(self.esmfold_config     ,			__snake_case			):
														__magic_name__						:						str										=self.esmfold_config.to_dict()
										return output




@dataclass
class 		__A   :
						UpperCamelCase      =       None
						UpperCamelCase      =       True
						UpperCamelCase      =       False
						UpperCamelCase      =       False
						UpperCamelCase      =       False
						UpperCamelCase      =       0

						UpperCamelCase      =       True
						UpperCamelCase      =       False

						UpperCamelCase      =       128
						UpperCamelCase      =       None

						def      A__    (		self :Union[str, Any]			):
										'''simple docstring'''

										if self.trunk is None:
														__magic_name__						:						Dict										=TrunkConfig()
										elif isinstance(self.trunk     ,			__snake_case			):
														__magic_name__						:						Optional[int]										=TrunkConfig(**self.trunk			)


						def      A__    (		self :Union[str, Any]			):
										'''simple docstring'''

										__magic_name__						:						List[Any]										=asdict(self			)
										__magic_name__						:						str										=self.trunk.to_dict()
										return output




@dataclass
class 		__A   :
						UpperCamelCase      =       48
						UpperCamelCase      =       1024
						UpperCamelCase      =       128
						UpperCamelCase      =       32
						UpperCamelCase      =       32
						UpperCamelCase      =       32
						UpperCamelCase      =       0
						UpperCamelCase      =       0
						UpperCamelCase      =       False
						UpperCamelCase      =       4
						UpperCamelCase      =       128
						UpperCamelCase      =       None

						def      A__    (		self :Optional[Any]			):
										'''simple docstring'''

										if self.structure_module is None:
														__magic_name__						:						Dict										=StructureModuleConfig()
										elif isinstance(self.structure_module     ,			__snake_case			):
														__magic_name__						:						int										=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}."			)

										__magic_name__						:						Union[str, Any]										=self.sequence_state_dim // self.sequence_head_width
										__magic_name__						:						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      A__    (		self :Optional[int]			):
										'''simple docstring'''

										__magic_name__						:						Dict										=asdict(self			)
										__magic_name__						:						Union[str, Any]										=self.structure_module.to_dict()
										return output




@dataclass
class 		__A   :
						UpperCamelCase      =       384
						UpperCamelCase      =       128
						UpperCamelCase      =       16
						UpperCamelCase      =       128
						UpperCamelCase      =       12
						UpperCamelCase      =       4
						UpperCamelCase      =       8
						UpperCamelCase      =       0.1
						UpperCamelCase      =       8
						UpperCamelCase      =       1
						UpperCamelCase      =       2
						UpperCamelCase      =       7
						UpperCamelCase      =       10
						UpperCamelCase      =       1e-8
						UpperCamelCase      =       1e5

						def      A__    (		self :Optional[int]			):
										'''simple docstring'''

										return asdict(self			)




def 		lowerCAmelCase_ (						):
				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>",
				)










import math

import tensorflow as tf
from packaging import version




def 		lowerCAmelCase_ (						lowerCamelCase ):
				__magic_name__						:						str										=tf.convert_to_tensor(lowerCamelCase )
				__magic_name__						:						List[str]										=0.5 * (1.0 + tf.math.erf(x / tf.cast(tf.sqrt(2.0 )						,  x.dtype ) ))

				return x * cdf




def 		lowerCAmelCase_ (						lowerCamelCase ):
				__magic_name__						:						str										=tf.convert_to_tensor(lowerCamelCase )
				__magic_name__						:						Optional[Any]										=tf.cast(math.pi						,  x.dtype )
				__magic_name__						:						int										=tf.cast(0.0_4_4_7_1_5						,  x.dtype )
				__magic_name__						:						Tuple										=0.5 * (1.0 + tf.tanh(tf.sqrt(2.0 / pi ) * (x + coeff * tf.pow(lowerCamelCase						,  3 )) ))

				return x * cdf




def 		lowerCAmelCase_ (						lowerCamelCase ):
				__magic_name__						:						Any										=tf.convert_to_tensor(lowerCamelCase )

				return x * tf.tanh(tf.math.softplus(lowerCamelCase ) )




def 		lowerCAmelCase_ (						lowerCamelCase ):
				__magic_name__						:						Optional[Any]										=tf.convert_to_tensor(lowerCamelCase )
				__magic_name__						:						Union[str, Any]										=tf.cast(0.0_4_4_7_1_5						,  x.dtype )
				__magic_name__						:						Tuple										=tf.cast(0.7_9_7_8_8_4_5_6_0_8						,  x.dtype )

				return 0.5 * x * (1.0 + tf.tanh(x * coeffa * (1.0 + coeffa * x * x) ))




def 		lowerCAmelCase_ (						lowerCamelCase ):
				__magic_name__						:						List[str]										=tf.convert_to_tensor(lowerCamelCase )
				__magic_name__						:						Dict										=tf.cast(1.7_0_2						,  x.dtype )
				return x * tf.math.sigmoid(coeff * x )




def 		lowerCAmelCase_ (						lowerCamelCase ):
				return tf.clip_by_value(_gelu(lowerCamelCase )						,  -10						,  10 )




def 		lowerCAmelCase_ (						lowerCamelCase						,  lowerCamelCase=-1 ):
				__magic_name__					,			__magic_name__						:						List[Any]										=tf.split(lowerCamelCase						,  2						,  axis=lowerCamelCase )
				return a * tf.math.sigmoid(lowerCamelCase )


if version.parse(tf.version.VERSION) >= version.parse("2.4"):




							def 		lowerCAmelCase_ (						lowerCamelCase ):
											return tf.keras.activations.gelu(lowerCamelCase						,  approximate=lowerCamelCase )

							UpperCAmelCase_		:					List[str]          =  tf.keras.activations.gelu



							UpperCAmelCase_		:					Dict          =  approximate_gelu_wrap
else:
							UpperCAmelCase_		:					Dict          =  _gelu
							UpperCAmelCase_		:					str          =  _gelu_new


UpperCAmelCase_		:					Any          =  {
    "gelu": gelu,
    "gelu_10": gelu_aa,
    "gelu_fast": gelu_fast,
    "gelu_new": gelu_new,
    "glu": glu,
    "mish": mish,
    "quick_gelu": quick_gelu,
    "relu": tf.keras.activations.relu,
    "sigmoid": tf.keras.activations.sigmoid,
    "silu": tf.keras.activations.swish,
    "swish": tf.keras.activations.swish,
    "tanh": tf.keras.activations.tanh,
}




def 		lowerCAmelCase_ (						lowerCamelCase ):
				if activation_string in ACTaFN:
								return ACTaFN[activation_string]
				else:
								raise KeyError(F"function {activation_string} not found in ACT2FN mapping {list(ACTaFN.keys() )}" )




