'''simple docstring'''




import unittest

from transformers import EsmConfig, is_torch_available
from transformers.testing_utils import TestCasePlus, require_torch, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
  import torch

  from transformers.models.esm.modeling_esmfold import EsmForProteinFolding




class    a__				:






     def __init__(			self		:		Optional[Any]  ,    a		:		Optional[Any]  ,    a		:		str=13  ,    a		:		Optional[Any]=7  ,    a		:		Dict=False  ,    a		:		Dict=True  ,    a		:		Dict=False  ,    a		:		Any=False  ,    a		:		Union[str, Any]=19  ,    a		:		List[str]=32  ,    a		:		Union[str, Any]=5  ,    a		:		Union[str, Any]=4  ,    a		:		Any=37  ,    a		:		int="gelu"  ,    a		:		Optional[int]=0.1  ,    a		:		Dict=0.1  ,    a		:		Any=5_12  ,    a		:		Optional[Any]=16  ,    a		:		List[Any]=2  ,    a		:		int=0.02  ,    a		:		Optional[int]=3  ,    a		:		Optional[int]=4  ,    a		:		Union[str, Any]=None  ,    ):





         """simple docstring"""



         __lowerCamelCase          =  parent
         __lowerCamelCase          =  batch_size
         __lowerCamelCase          =  seq_length
         __lowerCamelCase          =  is_training
         __lowerCamelCase          =  use_input_mask
         __lowerCamelCase          =  use_token_type_ids
         __lowerCamelCase          =  use_labels
         __lowerCamelCase          =  vocab_size
         __lowerCamelCase          =  hidden_size
         __lowerCamelCase          =  num_hidden_layers
         __lowerCamelCase          =  num_attention_heads
         __lowerCamelCase          =  intermediate_size
         __lowerCamelCase          =  hidden_act
         __lowerCamelCase          =  hidden_dropout_prob
         __lowerCamelCase          =  attention_probs_dropout_prob
         __lowerCamelCase          =  max_position_embeddings
         __lowerCamelCase          =  type_vocab_size
         __lowerCamelCase          =  type_sequence_label_size
         __lowerCamelCase          =  initializer_range
         __lowerCamelCase          =  num_labels
         __lowerCamelCase          =  num_choices
         __lowerCamelCase          =  scope






     def        SCREAMING_SNAKE_CASE__       (			self		:		Tuple   ):





         """simple docstring"""



         __lowerCamelCase          =  ids_tensor([self.batch_size, self.seq_length]  ,    self.vocab_size   )

         __lowerCamelCase          =  None
         if self.use_input_mask:
             __lowerCamelCase          =  random_attention_mask([self.batch_size, self.seq_length]   )

         __lowerCamelCase          =  None
         __lowerCamelCase          =  None
         __lowerCamelCase          =  None
         if self.use_labels:
             __lowerCamelCase          =  ids_tensor([self.batch_size]  ,    self.type_sequence_label_size   )
             __lowerCamelCase          =  ids_tensor([self.batch_size, self.seq_length]  ,    self.num_labels   )
             __lowerCamelCase          =  ids_tensor([self.batch_size]  ,    self.num_choices   )

         __lowerCamelCase          =  self.get_config()

         return config, input_ids, input_mask, sequence_labels, token_labels, choice_labels






     def        SCREAMING_SNAKE_CASE__       (			self		:		str   ):





         """simple docstring"""



         __lowerCamelCase          =  EsmConfig(
             vocab_size=33  ,    hidden_size=self.hidden_size  ,    pad_token_id=1  ,    num_hidden_layers=self.num_hidden_layers  ,    num_attention_heads=self.num_attention_heads  ,    intermediate_size=self.intermediate_size  ,    hidden_act=self.hidden_act  ,    hidden_dropout_prob=self.hidden_dropout_prob  ,    attention_probs_dropout_prob=self.attention_probs_dropout_prob  ,    max_position_embeddings=self.max_position_embeddings  ,    type_vocab_size=self.type_vocab_size  ,    initializer_range=self.initializer_range  ,    is_folding_model=a  ,    esmfold_config={'''trunk''': {'''num_blocks''': 2}, '''fp16_esm''': False}  ,    )
         return config






     def        SCREAMING_SNAKE_CASE__       (			self		:		Union[str, Any]  ,    a		:		Tuple  ,    a		:		List[str]  ,    a		:		Tuple  ,    a		:		Tuple  ,    a		:		Tuple  ,    a		:		Union[str, Any]   ):





         """simple docstring"""



         __lowerCamelCase          =  EsmForProteinFolding(config=a   ).float()
         model.to(a   )
         model.eval()
         __lowerCamelCase          =  model(a  ,    attention_mask=a   )
         __lowerCamelCase          =  model(a   )
         __lowerCamelCase          =  model(a   )

         self.parent.assertEqual(result.positions.shape  ,    (8, self.batch_size, self.seq_length, 14, 3)   )
         self.parent.assertEqual(result.angles.shape  ,    (8, self.batch_size, self.seq_length, 7, 2)   )






     def        SCREAMING_SNAKE_CASE__       (			self		:		List[Any]   ):





         """simple docstring"""



         __lowerCamelCase          =  self.prepare_config_and_inputs()
         (
             (
             __lowerCamelCase
         )     ,  (
             __lowerCamelCase
         )     ,  (
             __lowerCamelCase
         )     ,  (
             __lowerCamelCase
         )     ,  (
             __lowerCamelCase
         )     ,  (
             __lowerCamelCase
         )     ,  
         )          =  config_and_inputs
         __lowerCamelCase          =  {'''input_ids''': input_ids, '''attention_mask''': input_mask}
         return config, inputs_dict




@require_torch
class    a__				(       UpperCAmelCase__							,       UpperCAmelCase__							,       unittest.TestCase			):
     lowerCamelCase  :					Dict							=False

     lowerCamelCase  :					Optional[int]							=(EsmForProteinFolding,) if is_torch_available() else ()
     lowerCamelCase  :					Optional[int]							=()
     lowerCamelCase  :					Optional[Any]							={} if is_torch_available() else {}
     lowerCamelCase  :					List[str]							=False






     def        SCREAMING_SNAKE_CASE__       (			self		:		List[str]   ):





         """simple docstring"""



         __lowerCamelCase          =  EsmFoldModelTester(self   )
         __lowerCamelCase          =  ConfigTester(self  ,    config_class=a  ,    hidden_size=37   )






     def        SCREAMING_SNAKE_CASE__       (			self		:		List[Any]   ):





         """simple docstring"""



         self.config_tester.run_common_tests()






     def        SCREAMING_SNAKE_CASE__       (			self		:		str   ):





         """simple docstring"""



         __lowerCamelCase          =  self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_model(*a   )






     @unittest.skip('''Does not support attention outputs'''   )
     def        SCREAMING_SNAKE_CASE__       (			self		:		List[str]   ):





         """simple docstring"""



         pass






     @unittest.skip
     def        SCREAMING_SNAKE_CASE__       (			self		:		Tuple   ):





         """simple docstring"""



         pass






     @unittest.skip('''Esm does not support embedding resizing'''   )
     def        SCREAMING_SNAKE_CASE__       (			self		:		Tuple   ):





         """simple docstring"""



         pass






     @unittest.skip('''Esm does not support embedding resizing'''   )
     def        SCREAMING_SNAKE_CASE__       (			self		:		int   ):





         """simple docstring"""



         pass






     @unittest.skip('''ESMFold does not support passing input embeds!'''   )
     def        SCREAMING_SNAKE_CASE__       (			self		:		List[str]   ):





         """simple docstring"""



         pass






     @unittest.skip('''ESMFold does not support head pruning.'''   )
     def        SCREAMING_SNAKE_CASE__       (			self		:		Any   ):





         """simple docstring"""



         pass






     @unittest.skip('''ESMFold does not support head pruning.'''   )
     def        SCREAMING_SNAKE_CASE__       (			self		:		Dict   ):





         """simple docstring"""



         pass






     @unittest.skip('''ESMFold does not support head pruning.'''   )
     def        SCREAMING_SNAKE_CASE__       (			self		:		Tuple   ):





         """simple docstring"""



         pass






     @unittest.skip('''ESMFold does not support head pruning.'''   )
     def        SCREAMING_SNAKE_CASE__       (			self		:		List[str]   ):





         """simple docstring"""



         pass






     @unittest.skip('''ESMFold does not support head pruning.'''   )
     def        SCREAMING_SNAKE_CASE__       (			self		:		int   ):





         """simple docstring"""



         pass






     @unittest.skip('''ESMFold does not output hidden states in the normal way.'''   )
     def        SCREAMING_SNAKE_CASE__       (			self		:		Optional[Any]   ):





         """simple docstring"""



         pass






     @unittest.skip('''ESMfold does not output hidden states in the normal way.'''   )
     def        SCREAMING_SNAKE_CASE__       (			self		:		Optional[int]   ):





         """simple docstring"""



         pass






     @unittest.skip('''ESMFold only has one output format.'''   )
     def        SCREAMING_SNAKE_CASE__       (			self		:		Tuple   ):





         """simple docstring"""



         pass






     @unittest.skip('''This test doesn\'t work for ESMFold and doesn\'t test core functionality'''   )
     def        SCREAMING_SNAKE_CASE__       (			self		:		Any   ):





         """simple docstring"""



         pass






     @unittest.skip('''ESMFold does not support input chunking.'''   )
     def        SCREAMING_SNAKE_CASE__       (			self		:		List[Any]   ):





         """simple docstring"""



         pass






     @unittest.skip('''ESMFold doesn\'t respect you and it certainly doesn\'t respect your initialization arguments.'''   )
     def        SCREAMING_SNAKE_CASE__       (			self		:		str   ):





         """simple docstring"""



         pass






     @unittest.skip('''ESMFold doesn\'t support torchscript compilation.'''   )
     def        SCREAMING_SNAKE_CASE__       (			self		:		Any   ):





         """simple docstring"""



         pass






     @unittest.skip('''ESMFold doesn\'t support torchscript compilation.'''   )
     def        SCREAMING_SNAKE_CASE__       (			self		:		Optional[int]   ):





         """simple docstring"""



         pass






     @unittest.skip('''ESMFold doesn\'t support torchscript compilation.'''   )
     def        SCREAMING_SNAKE_CASE__       (			self		:		int   ):





         """simple docstring"""



         pass






     @unittest.skip('''ESMFold doesn\'t support data parallel.'''   )
     def        SCREAMING_SNAKE_CASE__       (			self		:		Optional[int]   ):





         """simple docstring"""



         pass






     @unittest.skip('''Will be fixed soon by reducing the size of the model used for common tests.'''   )
     def        SCREAMING_SNAKE_CASE__       (			self		:		Dict   ):





         """simple docstring"""



         pass




@require_torch
class    a__				(       UpperCAmelCase__			):






     @slow
     def        SCREAMING_SNAKE_CASE__       (			self		:		str   ):





         """simple docstring"""



         __lowerCamelCase          =  EsmForProteinFolding.from_pretrained('''facebook/esmfold_v1'''   ).float()
         model.eval()
         __lowerCamelCase          =  torch.tensor([[0, 6, 4, 13, 5, 4, 16, 12, 11, 7, 2]]   )
         __lowerCamelCase          =  model(a   )['''positions''']
         __lowerCamelCase          =  torch.tensor([2.58_28, 0.79_93, -10.93_34]  ,    dtype=torch.floataa   )
         self.assertTrue(torch.allclose(position_outputs[0, 0, 0, 0]  ,    a  ,    atol=1e-4   )   )

'''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    a__				(       UpperCAmelCase__			):
     lowerCamelCase  :					Dict							=(DPMSolverSDEScheduler,)
     lowerCamelCase  :					List[str]							=1_0






     def        SCREAMING_SNAKE_CASE__       (			self		:		Optional[Any]  ,    **a		:		Optional[int]   ):





         """simple docstring"""



         __lowerCamelCase          =  {
             '''num_train_timesteps''': 11_00,
             '''beta_start''': 0.00_01,
             '''beta_end''': 0.02,
             '''beta_schedule''': '''linear''',
             '''noise_sampler_seed''': 0,
         }

         config.update(**a   )
         return config






     def        SCREAMING_SNAKE_CASE__       (			self		:		Optional[Any]   ):





         """simple docstring"""



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






     def        SCREAMING_SNAKE_CASE__       (			self		:		str   ):





         """simple docstring"""



         for beta_start, beta_end in zip([0.0_00_01, 0.00_01, 0.0_01]  ,    [0.00_02, 0.0_02, 0.02]   ):
             self.check_over_configs(beta_start=a  ,    beta_end=a   )






     def        SCREAMING_SNAKE_CASE__       (			self		:		List[Any]   ):





         """simple docstring"""



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






     def        SCREAMING_SNAKE_CASE__       (			self		:		str   ):





         """simple docstring"""



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






     def        SCREAMING_SNAKE_CASE__       (			self		:		Tuple   ):





         """simple docstring"""



         __lowerCamelCase          =  self.scheduler_classes[0]
         __lowerCamelCase          =  self.get_scheduler_config()
         __lowerCamelCase          =  scheduler_class(**a   )

         scheduler.set_timesteps(self.num_inference_steps   )

         __lowerCamelCase          =  self.dummy_model()
         __lowerCamelCase          =  self.dummy_sample_deter * scheduler.init_noise_sigma
         __lowerCamelCase          =  sample.to(a   )

         for i, t in enumerate(scheduler.timesteps   ):
             __lowerCamelCase          =  scheduler.scale_model_input(a  ,    a   )

             __lowerCamelCase          =  model(a  ,    a   )

             __lowerCamelCase          =  scheduler.step(a  ,    a  ,    a   )
             __lowerCamelCase          =  output.prev_sample

         __lowerCamelCase          =  torch.sum(torch.abs(a   )   )
         __lowerCamelCase          =  torch.mean(torch.abs(a   )   )

         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.21_78_70_59_64_56_52_77   ) < 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_23_42_90_68_92_29_96_52   ) < 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.2_11_61_95_70_85_13_26   ) < 1e-3






     def        SCREAMING_SNAKE_CASE__       (			self		:		List[str]   ):





         """simple docstring"""



         __lowerCamelCase          =  self.scheduler_classes[0]
         __lowerCamelCase          =  self.get_scheduler_config(prediction_type='''v_prediction'''   )
         __lowerCamelCase          =  scheduler_class(**a   )

         scheduler.set_timesteps(self.num_inference_steps   )

         __lowerCamelCase          =  self.dummy_model()
         __lowerCamelCase          =  self.dummy_sample_deter * scheduler.init_noise_sigma
         __lowerCamelCase          =  sample.to(a   )

         for i, t in enumerate(scheduler.timesteps   ):
             __lowerCamelCase          =  scheduler.scale_model_input(a  ,    a   )

             __lowerCamelCase          =  model(a  ,    a   )

             __lowerCamelCase          =  scheduler.step(a  ,    a  ,    a   )
             __lowerCamelCase          =  output.prev_sample

         __lowerCamelCase          =  torch.sum(torch.abs(a   )   )
         __lowerCamelCase          =  torch.mean(torch.abs(a   )   )

         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_62_26_28_90_14_81_62_84   ) < 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_66_88_32_60_01_16_72_97   ) < 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.15_60_53_06_62_53_66_21   ) < 1e-3






     def        SCREAMING_SNAKE_CASE__       (			self		:		Dict   ):





         """simple docstring"""



         __lowerCamelCase          =  self.scheduler_classes[0]
         __lowerCamelCase          =  self.get_scheduler_config()
         __lowerCamelCase          =  scheduler_class(**a   )

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

         __lowerCamelCase          =  self.dummy_model()
         __lowerCamelCase          =  self.dummy_sample_deter.to(a   ) * scheduler.init_noise_sigma

         for t in scheduler.timesteps:
             __lowerCamelCase          =  scheduler.scale_model_input(a  ,    a   )

             __lowerCamelCase          =  model(a  ,    a   )

             __lowerCamelCase          =  scheduler.step(a  ,    a  ,    a   )
             __lowerCamelCase          =  output.prev_sample

         __lowerCamelCase          =  torch.sum(torch.abs(a   )   )
         __lowerCamelCase          =  torch.mean(torch.abs(a   )   )

         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_18_05_93_46_07_98_26_35   ) < 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_23_42_90_83_82_41_57_71   ) < 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.2_11_61_95_70_85_13_26   ) < 1e-3






     def        SCREAMING_SNAKE_CASE__       (			self		:		Optional[int]   ):





         """simple docstring"""



         __lowerCamelCase          =  self.scheduler_classes[0]
         __lowerCamelCase          =  self.get_scheduler_config()
         __lowerCamelCase          =  scheduler_class(**a  ,    use_karras_sigmas=a   )

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

         __lowerCamelCase          =  self.dummy_model()
         __lowerCamelCase          =  self.dummy_sample_deter.to(a   ) * scheduler.init_noise_sigma
         __lowerCamelCase          =  sample.to(a   )

         for t in scheduler.timesteps:
             __lowerCamelCase          =  scheduler.scale_model_input(a  ,    a   )

             __lowerCamelCase          =  model(a  ,    a   )

             __lowerCamelCase          =  scheduler.step(a  ,    a  ,    a   )
             __lowerCamelCase          =  output.prev_sample

         __lowerCamelCase          =  torch.sum(torch.abs(a   )   )
         __lowerCamelCase          =  torch.mean(torch.abs(a   )   )

         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_30_03_87_27_30_98_18_11   ) < 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_30_03_87_27_30_98_18_11   ) < 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_30_03_87_27_30_98_18_11   ) < 1e-2
