

'''simple docstring'''





from dataclasses import asdict, dataclass
from typing import Optional

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


UpperCamelCase_    						=       logging.get_logger(__name__)

# TODO Update this
UpperCamelCase_    						=       {
    """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__      ):
				lowerCamelCase_		       =    'esm'







				def __init__(       self     :		Optional[Any]		,					UpperCAmelCase__     :		str=None		,					UpperCAmelCase__     :		Dict=None		,					UpperCAmelCase__     :		Any=None		,					UpperCAmelCase__     :		int=768		,					UpperCAmelCase__     :		Union[str, Any]=12		,					UpperCAmelCase__     :		List[Any]=12		,					UpperCAmelCase__     :		Optional[Any]=3072		,					UpperCAmelCase__     :		Dict=0.1		,					UpperCAmelCase__     :		int=0.1		,					UpperCAmelCase__     :		int=1026		,					UpperCAmelCase__     :		Tuple=0.02		,					UpperCAmelCase__     :		Optional[int]=1E-12		,					UpperCAmelCase__     :		List[Any]="absolute"		,					UpperCAmelCase__     :		int=True		,					UpperCAmelCase__     :		Dict=None		,					UpperCAmelCase__     :		str=False		,					UpperCAmelCase__     :		int=False		,					UpperCAmelCase__     :		Optional[Any]=None		,					UpperCAmelCase__     :		List[str]=None		,					**UpperCAmelCase__     :		int		,					):
										'''simple docstring'''
										super().__init__(pad_token_id=UpperCAmelCase__		,					mask_token_id=UpperCAmelCase__		,					**UpperCAmelCase__   )

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







				def  lowerCamelCase_	(       self     :		Union[str, Any]   ):
										'''simple docstring'''
										lowercase :							Union[str, Any]    		=super().to_dict()
										if isinstance(self.esmfold_config		,					UpperCAmelCase__   ):
																lowercase :							Optional[Any]    		=self.esmfold_config.to_dict()
										return output





@dataclass
class 		__SCREAMING_SNAKE_CASE							:
				lowerCamelCase_		       =    None
				lowerCamelCase_		       =    True
				lowerCamelCase_		       =    False
				lowerCamelCase_		       =    False
				lowerCamelCase_		       =    False
				lowerCamelCase_		       =    0

				lowerCamelCase_		       =    True
				lowerCamelCase_		       =    False

				lowerCamelCase_		       =    1_28
				lowerCamelCase_		       =    None







				def  lowerCamelCase_	(       self     :		Tuple   ):
										'''simple docstring'''
										if self.trunk is None:
																lowercase :							str    		=TrunkConfig()
										elif isinstance(self.trunk		,					UpperCAmelCase__   ):
																lowercase :							int    		=TrunkConfig(**self.trunk   )







				def  lowerCamelCase_	(       self     :		List[Any]   ):
										'''simple docstring'''
										lowercase :							str    		=asdict(self   )
										lowercase :							Union[str, Any]    		=self.trunk.to_dict()
										return output





@dataclass
class 		__SCREAMING_SNAKE_CASE							:
				lowerCamelCase_		       =    48
				lowerCamelCase_		       =    10_24
				lowerCamelCase_		       =    1_28
				lowerCamelCase_		       =    32
				lowerCamelCase_		       =    32
				lowerCamelCase_		       =    32
				lowerCamelCase_		       =    0
				lowerCamelCase_		       =    0
				lowerCamelCase_		       =    False
				lowerCamelCase_		       =    4
				lowerCamelCase_		       =    1_28
				lowerCamelCase_		       =    None







				def  lowerCamelCase_	(       self     :		List[str]   ):
										'''simple docstring'''
										if self.structure_module is None:
																lowercase :							Any    		=StructureModuleConfig()
										elif isinstance(self.structure_module		,					UpperCAmelCase__   ):
																lowercase :							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}.'''   )

										lowercase :							str    		=self.sequence_state_dim // self.sequence_head_width
										lowercase :							int    		=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  lowerCamelCase_	(       self     :		Any   ):
										'''simple docstring'''
										lowercase :							List[Any]    		=asdict(self   )
										lowercase :							Any    		=self.structure_module.to_dict()
										return output





@dataclass
class 		__SCREAMING_SNAKE_CASE							:
				lowerCamelCase_		       =    3_84
				lowerCamelCase_		       =    1_28
				lowerCamelCase_		       =    16
				lowerCamelCase_		       =    1_28
				lowerCamelCase_		       =    12
				lowerCamelCase_		       =    4
				lowerCamelCase_		       =    8
				lowerCamelCase_		       =    0.1
				lowerCamelCase_		       =    8
				lowerCamelCase_		       =    1
				lowerCamelCase_		       =    2
				lowerCamelCase_		       =    7
				lowerCamelCase_		       =    10
				lowerCamelCase_		       =    1E-8
				lowerCamelCase_		       =    1E5







				def  lowerCamelCase_	(       self     :		Dict   ):
										'''simple docstring'''
										return asdict(self   )






def 				_lowerCAmelCase			(   )				->  Optional[int]:
						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 os
import tempfile
import unittest

from transformers import is_torch_available
from transformers.testing_utils import require_torch


if is_torch_available():
						import torch
						from torch import nn

						from transformers import (
						    Adafactor,
						    AdamW,
						    get_constant_schedule,
						    get_constant_schedule_with_warmup,
						    get_cosine_schedule_with_warmup,
						    get_cosine_with_hard_restarts_schedule_with_warmup,
						    get_inverse_sqrt_schedule,
						    get_linear_schedule_with_warmup,
						    get_polynomial_decay_schedule_with_warmup,
						)

def  __snake_case							(	__magic_name__		,    __magic_name__=10						):







				'''simple docstring'''


				lowercase         =      []
				for _ in range(__magic_name__						):
								lrs.append(scheduler.get_lr()[0]						)
								scheduler.step()
				return lrs

def  __snake_case							(	__magic_name__		,    __magic_name__=10						):







				'''simple docstring'''


				lowercase         =      []
				for step in range(__magic_name__						):
								lrs.append(scheduler.get_lr()[0]						)
								scheduler.step()
								if step == num_steps // 2:
												with tempfile.TemporaryDirectory() as tmpdirname:
																lowercase         =      os.path.join(__magic_name__		,    "schedule.bin"						)
																torch.save(scheduler.state_dict()		,    __magic_name__						)

																lowercase         =      torch.load(__magic_name__						)
																scheduler.load_state_dict(__magic_name__						)
				return lrs





@require_torch
class   UpperCamelCase_      (      unittest.TestCase		):







						'''simple docstring'''



						def        SCREAMING_SNAKE_CASE(			self							:Optional[Any]		,    lowerCAmelCase__							:List[str]		,    lowerCAmelCase__							:Tuple		,    lowerCAmelCase__							:Dict     )  ->Tuple:
										self.assertEqual(len(lowerCAmelCase__     )		,    len(lowerCAmelCase__     )     )
										for a, b in zip(lowerCAmelCase__		,    lowerCAmelCase__     ):
														self.assertAlmostEqual(lowerCAmelCase__		,    lowerCAmelCase__		,    delta=lowerCAmelCase__     )





						def        SCREAMING_SNAKE_CASE(			self							:Optional[int]     )  ->Dict:
										lowercase         =      torch.tensor([0.1, -0.2, -0.1]		,    requires_grad=lowerCAmelCase__     )
										lowercase         =      torch.tensor([0.4, 0.2, -0.5]     )
										lowercase         =      nn.MSELoss()
										# No warmup, constant schedule, no gradient clipping
										lowercase         =      AdamW(params=[w]		,    lr=2E-1		,    weight_decay=0.0     )
										for _ in range(100     ):
														lowercase         =      criterion(lowerCAmelCase__		,    lowerCAmelCase__     )
														loss.backward()
														optimizer.step()
														w.grad.detach_()  # No zero_grad() function on simple tensors. we do it ourselves.
														w.grad.zero_()
										self.assertListAlmostEqual(w.tolist()		,    [0.4, 0.2, -0.5]		,    tol=1E-2     )





						def        SCREAMING_SNAKE_CASE(			self							:int     )  ->List[Any]:
										lowercase         =      torch.tensor([0.1, -0.2, -0.1]		,    requires_grad=lowerCAmelCase__     )
										lowercase         =      torch.tensor([0.4, 0.2, -0.5]     )
										lowercase         =      nn.MSELoss()
										# No warmup, constant schedule, no gradient clipping
										lowercase         =      Adafactor(
										    params=[w]		,    lr=1E-2		,    eps=(1E-3_0, 1E-3)		,    clip_threshold=1.0		,    decay_rate=-0.8		,    betaa=lowerCAmelCase__		,    weight_decay=0.0		,    relative_step=lowerCAmelCase__		,    scale_parameter=lowerCAmelCase__		,    warmup_init=lowerCAmelCase__		,    )
										for _ in range(1000     ):
														lowercase         =      criterion(lowerCAmelCase__		,    lowerCAmelCase__     )
														loss.backward()
														optimizer.step()
														w.grad.detach_()  # No zero_grad() function on simple tensors. we do it ourselves.
														w.grad.zero_()
										self.assertListAlmostEqual(w.tolist()		,    [0.4, 0.2, -0.5]		,    tol=1E-2     )





@require_torch
class   UpperCamelCase_      (      unittest.TestCase		):







						'''simple docstring'''



						UpperCamelCase      :  Any				   =					nn.Linear(50		,    50		) if is_torch_available() else None
						UpperCamelCase      :  Union[str, Any]				   =					AdamW(m.parameters()		,    lr=10.0		) if is_torch_available() else None
						UpperCamelCase      :  int				   =					10





						def        SCREAMING_SNAKE_CASE(			self							:Any		,    lowerCAmelCase__							:Optional[Any]		,    lowerCAmelCase__							:Optional[Any]		,    lowerCAmelCase__							:Any		,    lowerCAmelCase__							:Any=None     )  ->Dict:
										self.assertEqual(len(lowerCAmelCase__     )		,    len(lowerCAmelCase__     )     )
										for a, b in zip(lowerCAmelCase__		,    lowerCAmelCase__     ):
														self.assertAlmostEqual(lowerCAmelCase__		,    lowerCAmelCase__		,    delta=lowerCAmelCase__		,    msg=lowerCAmelCase__     )





						def        SCREAMING_SNAKE_CASE(			self							:Any     )  ->str:
										lowercase         =      {"num_warmup_steps": 2, "num_training_steps": 10}
										# schedulers doct format
										# function: (sched_args_dict, expected_learning_rates)
										lowercase         =      {
										    get_constant_schedule: ({}, [10.0] * self.num_steps),
										    get_constant_schedule_with_warmup: (
										        {"num_warmup_steps": 4},
										        [0.0, 2.5, 5.0, 7.5, 10.0, 10.0, 10.0, 10.0, 10.0, 10.0],
										    ),
										    get_linear_schedule_with_warmup: (
										        {**common_kwargs},
										        [0.0, 5.0, 10.0, 8.75, 7.5, 6.25, 5.0, 3.75, 2.5, 1.25],
										    ),
										    get_cosine_schedule_with_warmup: (
										        {**common_kwargs},
										        [0.0, 5.0, 10.0, 9.61, 8.53, 6.91, 5.0, 3.08, 1.46, 0.38],
										    ),
										    get_cosine_with_hard_restarts_schedule_with_warmup: (
										        {**common_kwargs, "num_cycles": 2},
										        [0.0, 5.0, 10.0, 8.53, 5.0, 1.46, 10.0, 8.53, 5.0, 1.46],
										    ),
										    get_polynomial_decay_schedule_with_warmup: (
										        {**common_kwargs, "power": 2.0, "lr_end": 1E-7},
										        [0.0, 5.0, 10.0, 7.6_56, 5.6_25, 3.9_06, 2.5, 1.4_06, 0.6_25, 0.1_56],
										    ),
										    get_inverse_sqrt_schedule: (
										        {"num_warmup_steps": 2},
										        [0.0, 5.0, 10.0, 8.1_65, 7.0_71, 6.3_25, 5.7_74, 5.3_45, 5.0, 4.7_14],
										    ),
										}

										for scheduler_func, data in scheds.items():
														lowercase     ,	lowercase         =      data

														lowercase         =      scheduler_func(self.optimizer		,    **lowerCAmelCase__     )
														self.assertEqual(len([scheduler.get_lr()[0]]     )		,    1     )
														lowercase         =      unwrap_schedule(lowerCAmelCase__		,    self.num_steps     )
														self.assertListAlmostEqual(
														    lowerCAmelCase__		,    lowerCAmelCase__		,    tol=1E-2		,    msg=F'''failed for {scheduler_func} in normal scheduler'''		,    )

														lowercase         =      scheduler_func(self.optimizer		,    **lowerCAmelCase__     )
														if scheduler_func.__name__ != "get_constant_schedule":
																		LambdaScheduleWrapper.wrap_scheduler(lowerCAmelCase__     )  # wrap to test picklability of the schedule
														lowercase         =      unwrap_and_save_reload_schedule(lowerCAmelCase__		,    self.num_steps     )
														self.assertListEqual(lowerCAmelCase__		,    lowerCAmelCase__		,    msg=F'''failed for {scheduler_func} in save and reload'''     )





class   UpperCamelCase_      :







						'''simple docstring'''



						def __init__(			self							:List[str]		,    lowerCAmelCase__							:int     )  ->Optional[int]:
										lowercase         =      fn





						def __call__(			self							:Optional[Any]		,    *lowerCAmelCase__							:int		,    **lowerCAmelCase__							:List[str]     )  ->List[str]:
										return self.fn(*lowerCAmelCase__		,    **lowerCAmelCase__     )





						@classmethod
						def        SCREAMING_SNAKE_CASE(			self							:List[Any]		,    lowerCAmelCase__							:Optional[Any]     )  ->str:
										lowercase         =      list(map(self		,    scheduler.lr_lambdas     )     )
