





"""simple docstring"""

import torch

from diffusers import DPMSolverSDEScheduler
from diffusers.utils import torch_device
from diffusers.utils.testing_utils import require_torchsde

from .test_schedulers import SchedulerCommonTest


@require_torchsde
class       lowerCAmelCase__				(   lowercase__  ):
 lowercase__     :  Dict				=   (DPMSolverSDEScheduler,)
 lowercase__     :  List[Any]				=   10







 def  lowercase_       (   self		,	**UpperCamelCase__					):

    '''simple docstring'''


    A__											=					{
        '''num_train_timesteps''': 11_00,
        '''beta_start''': 0.0001,
        '''beta_end''': 0.02,
        '''beta_schedule''': '''linear''',
        '''noise_sampler_seed''': 0,
    }

    config.update(**UpperCAmelCase__					)
    return config







 def  lowercase_       (   self					):

    '''simple docstring'''


    for timesteps in [10, 50, 1_00, 10_00]:
       self.check_over_configs(num_train_timesteps=UpperCAmelCase__					)







 def  lowercase_       (   self					):

    '''simple docstring'''


    for beta_start, beta_end in zip([0.0_0001, 0.0001, 0.001]		,	[0.0002, 0.002, 0.02]					):
       self.check_over_configs(beta_start=UpperCAmelCase__		,	beta_end=UpperCAmelCase__					)







 def  lowercase_       (   self					):

    '''simple docstring'''


    for schedule in ["linear", "scaled_linear"]:
       self.check_over_configs(beta_schedule=UpperCAmelCase__					)







 def  lowercase_       (   self					):

    '''simple docstring'''


    for prediction_type in ["epsilon", "v_prediction"]:
       self.check_over_configs(prediction_type=UpperCAmelCase__					)







 def  lowercase_       (   self					):

    '''simple docstring'''


    A__											=					self.scheduler_classes[0]
    A__											=					self.get_scheduler_config()
    A__											=					scheduler_class(**UpperCAmelCase__					)

    scheduler.set_timesteps(self.num_inference_steps					)

    A__											=					self.dummy_model()
    A__											=					self.dummy_sample_deter * scheduler.init_noise_sigma
    A__											=					sample.to(UpperCAmelCase__					)

    for i, t in enumerate(scheduler.timesteps					):
       A__											=					scheduler.scale_model_input(UpperCAmelCase__		,	UpperCAmelCase__					)

       A__											=					model(UpperCAmelCase__		,	UpperCAmelCase__					)

       A__											=					scheduler.step(UpperCAmelCase__		,	UpperCAmelCase__		,	UpperCAmelCase__					)
       A__											=					output.prev_sample

    A__											=					torch.sum(torch.abs(UpperCAmelCase__					)					)
    A__											=					torch.mean(torch.abs(UpperCAmelCase__					)					)

    if torch_device in ["mps"]:
       assert abs(result_sum.item() - 1_67.47_82_10_44_92_18_75					) < 1e-2
       assert abs(result_mean.item() - 0.2178_7059_6456_5277					) < 1e-3
    elif torch_device in ["cuda"]:
       assert abs(result_sum.item() - 1_71.59_35_21_11_81_64_06					) < 1e-2
       assert abs(result_mean.item() - 0.2_2342_9068_9229_9652					) < 1e-3
    else:
       assert abs(result_sum.item() - 1_62.52_38_34_22_85_15_62					) < 1e-2
       assert abs(result_mean.item() - 0.211_6195_7085_1326					) < 1e-3







 def  lowercase_       (   self					):

    '''simple docstring'''


    A__											=					self.scheduler_classes[0]
    A__											=					self.get_scheduler_config(prediction_type="v_prediction"					)
    A__											=					scheduler_class(**UpperCAmelCase__					)

    scheduler.set_timesteps(self.num_inference_steps					)

    A__											=					self.dummy_model()
    A__											=					self.dummy_sample_deter * scheduler.init_noise_sigma
    A__											=					sample.to(UpperCAmelCase__					)

    for i, t in enumerate(scheduler.timesteps					):
       A__											=					scheduler.scale_model_input(UpperCAmelCase__		,	UpperCAmelCase__					)

       A__											=					model(UpperCAmelCase__		,	UpperCAmelCase__					)

       A__											=					scheduler.step(UpperCAmelCase__		,	UpperCAmelCase__		,	UpperCAmelCase__					)
       A__											=					output.prev_sample

    A__											=					torch.sum(torch.abs(UpperCAmelCase__					)					)
    A__											=					torch.mean(torch.abs(UpperCAmelCase__					)					)

    if torch_device in ["mps"]:
       assert abs(result_sum.item() - 1_24.77_14_92_00_43_94_53					) < 1e-2
       assert abs(result_mean.item() - 0.1_6226_2890_1481_6284					) < 1e-3
    elif torch_device in ["cuda"]:
       assert abs(result_sum.item() - 1_28.1_66_33_60_59_57_03					) < 1e-2
       assert abs(result_mean.item() - 0.1_6688_3260_0116_7297					) < 1e-3
    else:
       assert abs(result_sum.item() - 1_19.8_48_75_48_82_81_25					) < 1e-2
       assert abs(result_mean.item() - 0.1560_5306_6253_6621					) < 1e-3







 def  lowercase_       (   self					):

    '''simple docstring'''


    A__											=					self.scheduler_classes[0]
    A__											=					self.get_scheduler_config()
    A__											=					scheduler_class(**UpperCAmelCase__					)

    scheduler.set_timesteps(self.num_inference_steps		,	device=UpperCAmelCase__					)

    A__											=					self.dummy_model()
    A__											=					self.dummy_sample_deter.to(UpperCAmelCase__					) * scheduler.init_noise_sigma

    for t in scheduler.timesteps:
       A__											=					scheduler.scale_model_input(UpperCAmelCase__		,	UpperCAmelCase__					)

       A__											=					model(UpperCAmelCase__		,	UpperCAmelCase__					)

       A__											=					scheduler.step(UpperCAmelCase__		,	UpperCAmelCase__		,	UpperCAmelCase__					)
       A__											=					output.prev_sample

    A__											=					torch.sum(torch.abs(UpperCAmelCase__					)					)
    A__											=					torch.mean(torch.abs(UpperCAmelCase__					)					)

    if torch_device in ["mps"]:
       assert abs(result_sum.item() - 1_67.46_95_73_97_46_09_38					) < 1e-2
       assert abs(result_mean.item() - 0.2_1805_9346_0798_2635					) < 1e-3
    elif torch_device in ["cuda"]:
       assert abs(result_sum.item() - 1_71.59_35_36_37_69_53_12					) < 1e-2
       assert abs(result_mean.item() - 0.2_2342_9083_8241_5771					) < 1e-3
    else:
       assert abs(result_sum.item() - 1_62.52_38_34_22_85_15_62					) < 1e-2
       assert abs(result_mean.item() - 0.211_6195_7085_1326					) < 1e-3







 def  lowercase_       (   self					):

    '''simple docstring'''


    A__											=					self.scheduler_classes[0]
    A__											=					self.get_scheduler_config()
    A__											=					scheduler_class(**UpperCAmelCase__		,	use_karras_sigmas=UpperCAmelCase__					)

    scheduler.set_timesteps(self.num_inference_steps		,	device=UpperCAmelCase__					)

    A__											=					self.dummy_model()
    A__											=					self.dummy_sample_deter.to(UpperCAmelCase__					) * scheduler.init_noise_sigma
    A__											=					sample.to(UpperCAmelCase__					)

    for t in scheduler.timesteps:
       A__											=					scheduler.scale_model_input(UpperCAmelCase__		,	UpperCAmelCase__					)

       A__											=					model(UpperCAmelCase__		,	UpperCAmelCase__					)

       A__											=					scheduler.step(UpperCAmelCase__		,	UpperCAmelCase__		,	UpperCAmelCase__					)
       A__											=					output.prev_sample

    A__											=					torch.sum(torch.abs(UpperCAmelCase__					)					)
    A__											=					torch.mean(torch.abs(UpperCAmelCase__					)					)

    if torch_device in ["mps"]:
       assert abs(result_sum.item() - 1_76.66_97_41_35_74_21_88					) < 1e-2
       assert abs(result_mean.item() - 0.2_3003_8727_3098_1811					) < 1e-2
    elif torch_device in ["cuda"]:
       assert abs(result_sum.item() - 1_77.63_65_35_64_45_31_25					) < 1e-2
       assert abs(result_mean.item() - 0.2_3003_8727_3098_1811					) < 1e-2
    else:
       assert abs(result_sum.item() - 1_70.3_13_52_23_38_86_72					) < 1e-2
       assert abs(result_mean.item() - 0.2_3003_8727_3098_1811					) < 1e-2


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





