



from typing import Mapping

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxSeqaSeqConfigWithPast
from ...utils import logging


lowercase_	       =			logging.get_logger(__name__)

lowercase_	       =			{
    '''google/umt5-small''': '''https://huggingface.co/google/umt5-small/resolve/main/config.json''',
    # See all umt5 models at https://huggingface.co/models?filter=umt5
}






class 						__a (     SCREAMING_SNAKE_CASE				):
	SCREAMING_SNAKE_CASE       =       "umt5"
	SCREAMING_SNAKE_CASE       =       ["past_key_values"]



	def __init__(						self							:    List[Any]				,	snake_case_							:    Tuple=25_01_12				,	snake_case_							:    Tuple=5_12				,	snake_case_							:    Union[str, Any]=64				,	snake_case_							:    Dict=10_24				,	snake_case_							:    Tuple=8				,	snake_case_							:    Any=None				,	snake_case_							:    List[str]=6				,	snake_case_							:    Any=32				,	snake_case_							:    Union[str, Any]=1_28				,	snake_case_							:    List[str]=0.1				,	snake_case_							:    Tuple=1e-6				,	snake_case_							:    Union[str, Any]=1.0				,	snake_case_							:    Tuple="gated-gelu"				,	snake_case_							:    Tuple=True				,	snake_case_							:    Optional[int]=True				,	snake_case_							:    Dict="T5Tokenizer"				,	snake_case_							:    List[Any]=True				,	snake_case_							:    Optional[Any]=0				,	snake_case_							:    Tuple=1				,	snake_case_							:    Union[str, Any]=0				,	**snake_case_							:    Optional[Any]				,	)->							Optional[int]:
						super().__init__(
						    is_encoder_decoder=snake_case_				,	tokenizer_class=snake_case_				,	tie_word_embeddings=snake_case_				,	pad_token_id=snake_case_				,	eos_token_id=snake_case_				,	decoder_start_token_id=snake_case_				,	**snake_case_				,	)
						__lowerCAmelCase	   =vocab_size
						__lowerCAmelCase	   =d_model
						__lowerCAmelCase	   =d_kv
						__lowerCAmelCase	   =d_ff
						__lowerCAmelCase	   =num_layers
						__lowerCAmelCase	   =(
						    num_decoder_layers if num_decoder_layers is not None else self.num_layers
						)  # default = symmetry
						__lowerCAmelCase	   =num_heads
						__lowerCAmelCase	   =relative_attention_num_buckets
						__lowerCAmelCase	   =relative_attention_max_distance
						__lowerCAmelCase	   =dropout_rate
						__lowerCAmelCase	   =layer_norm_epsilon
						__lowerCAmelCase	   =initializer_factor
						__lowerCAmelCase	   =feed_forward_proj
						__lowerCAmelCase	   =use_cache

						__lowerCAmelCase	   =self.feed_forward_proj.split("""-""")
						__lowerCAmelCase	   =act_info[-1]
						__lowerCAmelCase	   =act_info[0] == """gated"""

						if len(snake_case_) > 1 and act_info[0] != "gated" or len(snake_case_) > 2:
											raise ValueError(
											    F"""`feed_forward_proj`: {feed_forward_proj} is not a valid activation function of the dense layer."""
											    """Please make sure `feed_forward_proj` is of the format `gated-{ACT_FN}` or `{ACT_FN}`, e.g. """
											    """'gated-gelu' or 'relu'""")

						if feed_forward_proj == "gated-gelu":
											__lowerCAmelCase	   ="""gelu_new"""



	@property
	def 		UpperCamelCase							(						self							:    Any)->							Optional[int]:
						return self.d_model



	@property
	def 		UpperCamelCase							(						self							:    List[str])->							Optional[int]:
						return self.num_heads



	@property
	def 		UpperCamelCase							(						self							:    List[str])->							int:
						return self.num_layers






class 						__a (     SCREAMING_SNAKE_CASE				):



	@property
	# Copied from transformers.models.t5.configuration_t5.T5OnnxConfig.inputs
	def 		UpperCamelCase							(						self							:    List[str])->							Mapping[str, Mapping[int, str]]:
						__lowerCAmelCase	   ={
						    """input_ids""": {0: """batch""", 1: """encoder_sequence"""},
						    """attention_mask""": {0: """batch""", 1: """encoder_sequence"""},
						}
						if self.use_past:
											__lowerCAmelCase	   ="""past_encoder_sequence + sequence"""
											__lowerCAmelCase	   ={0: """batch"""}
											__lowerCAmelCase	   ={0: """batch""", 1: """past_decoder_sequence + sequence"""}
						else:
											__lowerCAmelCase	   ={0: """batch""", 1: """decoder_sequence"""}
											__lowerCAmelCase	   ={0: """batch""", 1: """decoder_sequence"""}

						if self.use_past:
											self.fill_with_past_key_values_(snake_case_				,	direction="""inputs""")

						return common_inputs



	@property
	# Copied from transformers.models.t5.configuration_t5.T5OnnxConfig.default_onnx_opset
	def 		UpperCamelCase							(						self							:    str)->							int:
						return 13



	@property
	def 		UpperCamelCase							(						self							:    int)->							float:
						return 5e-4









from dataclasses import asdict, dataclass
from typing import Optional

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


lowercase_	       =			logging.get_logger(__name__)

# TODO Update this
lowercase_	       =			{
    '''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 (     SCREAMING_SNAKE_CASE				):
	SCREAMING_SNAKE_CASE       =       "esm"



	def __init__(						self							:    Any				,	snake_case_							:    List[str]=None				,	snake_case_							:    Tuple=None				,	snake_case_							:    str=None				,	snake_case_							:    str=7_68				,	snake_case_							:    Optional[int]=12				,	snake_case_							:    Tuple=12				,	snake_case_							:    Dict=30_72				,	snake_case_							:    Optional[Any]=0.1				,	snake_case_							:    Dict=0.1				,	snake_case_							:    Optional[Any]=10_26				,	snake_case_							:    int=0.0_2				,	snake_case_							:    Any=1e-12				,	snake_case_							:    Union[str, Any]="absolute"				,	snake_case_							:    Optional[int]=True				,	snake_case_							:    Tuple=None				,	snake_case_							:    Optional[int]=False				,	snake_case_							:    int=False				,	snake_case_							:    Optional[int]=None				,	snake_case_							:    List[Any]=None				,	**snake_case_							:    Optional[int]				,	)->							Any:
						super().__init__(pad_token_id=snake_case_				,	mask_token_id=snake_case_				,	**snake_case_)

						__lowerCAmelCase	   =vocab_size
						__lowerCAmelCase	   =hidden_size
						__lowerCAmelCase	   =num_hidden_layers
						__lowerCAmelCase	   =num_attention_heads
						__lowerCAmelCase	   =intermediate_size
						__lowerCAmelCase	   =hidden_dropout_prob
						__lowerCAmelCase	   =attention_probs_dropout_prob
						__lowerCAmelCase	   =max_position_embeddings
						__lowerCAmelCase	   =initializer_range
						__lowerCAmelCase	   =layer_norm_eps
						__lowerCAmelCase	   =position_embedding_type
						__lowerCAmelCase	   =use_cache
						__lowerCAmelCase	   =emb_layer_norm_before
						__lowerCAmelCase	   =token_dropout
						__lowerCAmelCase	   =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	   =EsmFoldConfig()
											elif isinstance(snake_case_				,	snake_case_):
																__lowerCAmelCase	   =EsmFoldConfig(**snake_case_)
											__lowerCAmelCase	   =esmfold_config
											if vocab_list is None:
																logger.warning("""No vocab_list supplied for folding model, assuming the ESM-2 vocabulary!""")
																__lowerCAmelCase	   =get_default_vocab_list()
											else:
																__lowerCAmelCase	   =vocab_list
						else:
											__lowerCAmelCase	   =None
											__lowerCAmelCase	   =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 		UpperCamelCase							(						self							:    Dict)->							Union[str, Any]:
						__lowerCAmelCase	   =super().to_dict()
						if isinstance(self.esmfold_config				,	snake_case_):
											__lowerCAmelCase	   =self.esmfold_config.to_dict()
						return output






@dataclass
class 						__a :
	SCREAMING_SNAKE_CASE       =       None
	SCREAMING_SNAKE_CASE       =       True
	SCREAMING_SNAKE_CASE       =       False
	SCREAMING_SNAKE_CASE       =       False
	SCREAMING_SNAKE_CASE       =       False
	SCREAMING_SNAKE_CASE       =       0

	SCREAMING_SNAKE_CASE       =       True
	SCREAMING_SNAKE_CASE       =       False

	SCREAMING_SNAKE_CASE       =       128
	SCREAMING_SNAKE_CASE       =       None



	def 		UpperCamelCase							(						self							:    Dict)->							List[Any]:
						if self.trunk is None:
											__lowerCAmelCase	   =TrunkConfig()
						elif isinstance(self.trunk				,	snake_case_):
											__lowerCAmelCase	   =TrunkConfig(**self.trunk)



	def 		UpperCamelCase							(						self							:    Optional[Any])->							int:
						__lowerCAmelCase	   =asdict(self)
						__lowerCAmelCase	   =self.trunk.to_dict()
						return output






@dataclass
class 						__a :
	SCREAMING_SNAKE_CASE       =       48
	SCREAMING_SNAKE_CASE       =       1024
	SCREAMING_SNAKE_CASE       =       128
	SCREAMING_SNAKE_CASE       =       32
	SCREAMING_SNAKE_CASE       =       32
	SCREAMING_SNAKE_CASE       =       32
	SCREAMING_SNAKE_CASE       =       0
	SCREAMING_SNAKE_CASE       =       0
	SCREAMING_SNAKE_CASE       =       False
	SCREAMING_SNAKE_CASE       =       4
	SCREAMING_SNAKE_CASE       =       128
	SCREAMING_SNAKE_CASE       =       None



	def 		UpperCamelCase							(						self							:    Tuple)->							Optional[int]:
						if self.structure_module is None:
											__lowerCAmelCase	   =StructureModuleConfig()
						elif isinstance(self.structure_module				,	snake_case_):
											__lowerCAmelCase	   =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	   =self.sequence_state_dim // self.sequence_head_width
						__lowerCAmelCase	   =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 		UpperCamelCase							(						self							:    Dict)->							Any:
						__lowerCAmelCase	   =asdict(self)
						__lowerCAmelCase	   =self.structure_module.to_dict()
						return output






@dataclass
class 						__a :
	SCREAMING_SNAKE_CASE       =       384
	SCREAMING_SNAKE_CASE       =       128
	SCREAMING_SNAKE_CASE       =       16
	SCREAMING_SNAKE_CASE       =       128
	SCREAMING_SNAKE_CASE       =       12
	SCREAMING_SNAKE_CASE       =       4
	SCREAMING_SNAKE_CASE       =       8
	SCREAMING_SNAKE_CASE       =       0.1
	SCREAMING_SNAKE_CASE       =       8
	SCREAMING_SNAKE_CASE       =       1
	SCREAMING_SNAKE_CASE       =       2
	SCREAMING_SNAKE_CASE       =       7
	SCREAMING_SNAKE_CASE       =       10
	SCREAMING_SNAKE_CASE       =       1E-8
	SCREAMING_SNAKE_CASE       =       1E5



	def 		UpperCamelCase							(						self							:    Any)->							Tuple:
						return asdict(self)






def 					__lowerCAmelCase		(  )  ->						List[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>",
					)




