



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     _UpperCAmelCase			:




     def __init__(				self ,						a__ ,						a__=1_3 ,						a__=7 ,						a__=False ,						a__=True ,						a__=False ,						a__=False ,						a__=1_9 ,						a__=3_2 ,						a__=5 ,						a__=4 ,						a__=3_7 ,						a__="gelu" ,						a__=0.1 ,						a__=0.1 ,						a__=5_1_2 ,						a__=1_6 ,						a__=2 ,						a__=0.0_2 ,						a__=3 ,						a__=4 ,						a__=None ,						):
         A__		 =					parent
         A__		 =					batch_size
         A__		 =					seq_length
         A__		 =					is_training
         A__		 =					use_input_mask
         A__		 =					use_token_type_ids
         A__		 =					use_labels
         A__		 =					vocab_size
         A__		 =					hidden_size
         A__		 =					num_hidden_layers
         A__		 =					num_attention_heads
         A__		 =					intermediate_size
         A__		 =					hidden_act
         A__		 =					hidden_dropout_prob
         A__		 =					attention_probs_dropout_prob
         A__		 =					max_position_embeddings
         A__		 =					type_vocab_size
         A__		 =					type_sequence_label_size
         A__		 =					initializer_range
         A__		 =					num_labels
         A__		 =					num_choices
         A__		 =					scope




     def 		snake_case_		(				self):
         A__		 =					ids_tensor([self.batch_size, self.seq_length] ,						self.vocab_size)

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

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

         A__		 =					self.get_config()

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




     def 		snake_case_		(				self):
         A__		 =					EsmConfig(
             vocab_size=3_3 ,						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 		snake_case_		(				self ,						a__ ,						a__ ,						a__ ,						a__ ,						a__ ,						a__):
         A__		 =					EsmForProteinFolding(config=a__).float()
         model.to(a__)
         model.eval()
         A__		 =					model(a__ ,						attention_mask=a__)
         A__		 =					model(a__)
         A__		 =					model(a__)

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




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







@require_torch
class     _UpperCAmelCase			(       A__			,      A__			,      unittest.TestCase     ):
     UpperCamelCase__		    =				False

     UpperCamelCase__		    =				(EsmForProteinFolding,) if is_torch_available() else ()
     UpperCamelCase__		    =				()
     UpperCamelCase__		    =				{} if is_torch_available() else {}
     UpperCamelCase__		    =				False




     def 		snake_case_		(				self):
         A__		 =					EsmFoldModelTester(self)
         A__		 =					ConfigTester(self ,						config_class=a__ ,						hidden_size=3_7)




     def 		snake_case_		(				self):
         self.config_tester.run_common_tests()




     def 		snake_case_		(				self):
         A__		 =					self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_model(*a__)




     @unittest.skip('''Does not support attention outputs''')
     def 		snake_case_		(				self):
         pass




     @unittest.skip
     def 		snake_case_		(				self):
         pass




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




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




     @unittest.skip('''ESMFold does not support passing input embeds!''')
     def 		snake_case_		(				self):
         pass




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




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




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




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




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




     @unittest.skip('''ESMFold does not output hidden states in the normal way.''')
     def 		snake_case_		(				self):
         pass




     @unittest.skip('''ESMfold does not output hidden states in the normal way.''')
     def 		snake_case_		(				self):
         pass




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




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




     @unittest.skip('''ESMFold does not support input chunking.''')
     def 		snake_case_		(				self):
         pass




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




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




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




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




     @unittest.skip('''ESMFold doesn\'t support data parallel.''')
     def 		snake_case_		(				self):
         pass




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







@require_torch
class     _UpperCAmelCase			(       A__     ):




     @slow
     def 		snake_case_		(				self):
         A__		 =					EsmForProteinFolding.from_pretrained('''facebook/esmfold_v1''').float()
         model.eval()
         A__		 =					torch.tensor([[0, 6, 4, 1_3, 5, 4, 1_6, 1_2, 1_1, 7, 2]])
         A__		 =					model(a__)['''positions''']
         A__		 =					torch.tensor([2.5_8_2_8, 0.7_9_9_3, -1_0.9_3_3_4] ,						dtype=torch.floataa)
         self.assertTrue(torch.allclose(position_outputs[0, 0, 0, 0] ,						a__ ,						atol=1e-4))





import torch

from diffusers import DDPMScheduler

from .test_schedulers import SchedulerCommonTest



class     _UpperCAmelCase			(       A__     ):
     UpperCamelCase__		    =				(DDPMScheduler,)




     def 		snake_case_		(				self ,						**a__):
         A__		 =					{
             '''num_train_timesteps''': 1_0_0_0,
             '''beta_start''': 0.0_0_0_1,
             '''beta_end''': 0.0_2,
             '''beta_schedule''': '''linear''',
             '''variance_type''': '''fixed_small''',
             '''clip_sample''': True,
         }

         config.update(**a__)
         return config




     def 		snake_case_		(				self):
         for timesteps in [1, 5, 1_0_0, 1_0_0_0]:
             self.check_over_configs(num_train_timesteps=a__)




     def 		snake_case_		(				self):
         for beta_start, beta_end in zip([0.0_0_0_1, 0.0_0_1, 0.0_1, 0.1] ,						[0.0_0_2, 0.0_2, 0.2, 2]):
             self.check_over_configs(beta_start=a__ ,						beta_end=a__)




     def 		snake_case_		(				self):
         for schedule in ["linear", "squaredcos_cap_v2"]:
             self.check_over_configs(beta_schedule=a__)




     def 		snake_case_		(				self):
         for variance in ["fixed_small", "fixed_large", "other"]:
             self.check_over_configs(variance_type=a__)




     def 		snake_case_		(				self):
         for clip_sample in [True, False]:
             self.check_over_configs(clip_sample=a__)




     def 		snake_case_		(				self):
         self.check_over_configs(thresholding=a__)
         for threshold in [0.5, 1.0, 2.0]:
             for prediction_type in ["epsilon", "sample", "v_prediction"]:
                 self.check_over_configs(
                     thresholding=a__ ,						prediction_type=a__ ,						sample_max_value=a__ ,						)




     def 		snake_case_		(				self):
         for prediction_type in ["epsilon", "sample", "v_prediction"]:
             self.check_over_configs(prediction_type=a__)




     def 		snake_case_		(				self):
         for t in [0, 5_0_0, 9_9_9]:
             self.check_over_forward(time_step=a__)




     def 		snake_case_		(				self):
         A__		 =					self.scheduler_classes[0]
         A__		 =					self.get_scheduler_config()
         A__		 =					scheduler_class(**a__)

         assert torch.sum(torch.abs(scheduler._get_variance(0) - 0.0)) < 1e-5
         assert torch.sum(torch.abs(scheduler._get_variance(4_8_7) - 0.0_0_9_7_9)) < 1e-5
         assert torch.sum(torch.abs(scheduler._get_variance(9_9_9) - 0.0_2)) < 1e-5




     def 		snake_case_		(				self):
         A__		 =					self.scheduler_classes[0]
         A__		 =					self.get_scheduler_config()
         A__		 =					scheduler_class(**a__)

         A__		 =					len(a__)

         A__		 =					self.dummy_model()
         A__		 =					self.dummy_sample_deter
         A__		 =					torch.manual_seed(0)

         for t in reversed(range(a__)):
             # 1. predict noise residual
             A__		 =					model(a__ ,						a__)

             # 2. predict previous mean of sample x_t-1
             A__		 =					scheduler.step(a__ ,						a__ ,						a__ ,						generator=a__).prev_sample

             # if t > 0:
             #     noise = self.dummy_sample_deter
             #     variance = scheduler.get_variance(t) ** (0.5) * noise
             #
             # sample = pred_prev_sample + variance
             A__		 =					pred_prev_sample

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

         assert abs(result_sum.item() - 2_5_8.9_6_0_6) < 1e-2
         assert abs(result_mean.item() - 0.3_3_7_2) < 1e-3




     def 		snake_case_		(				self):
         A__		 =					self.scheduler_classes[0]
         A__		 =					self.get_scheduler_config(prediction_type='''v_prediction''')
         A__		 =					scheduler_class(**a__)

         A__		 =					len(a__)

         A__		 =					self.dummy_model()
         A__		 =					self.dummy_sample_deter
         A__		 =					torch.manual_seed(0)

         for t in reversed(range(a__)):
             # 1. predict noise residual
             A__		 =					model(a__ ,						a__)

             # 2. predict previous mean of sample x_t-1
             A__		 =					scheduler.step(a__ ,						a__ ,						a__ ,						generator=a__).prev_sample

             # if t > 0:
             #     noise = self.dummy_sample_deter
             #     variance = scheduler.get_variance(t) ** (0.5) * noise
             #
             # sample = pred_prev_sample + variance
             A__		 =					pred_prev_sample

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

         assert abs(result_sum.item() - 2_0_2.0_2_9_6) < 1e-2
         assert abs(result_mean.item() - 0.2_6_3_1) < 1e-3




     def 		snake_case_		(				self):
         A__		 =					self.scheduler_classes[0]
         A__		 =					self.get_scheduler_config()
         A__		 =					scheduler_class(**a__)

         A__		 =					[1_0_0, 8_7, 5_0, 1, 0]

         scheduler.set_timesteps(timesteps=a__)

         A__		 =					scheduler.timesteps

         for i, timestep in enumerate(a__):
             if i == len(a__) - 1:
                 A__		 =					-1
             else:
                 A__		 =					timesteps[i + 1]

             A__		 =					scheduler.previous_timestep(a__)
             A__		 =					prev_t.item()

             self.assertEqual(a__ ,						a__)




     def 		snake_case_		(				self):
         A__		 =					self.scheduler_classes[0]
         A__		 =					self.get_scheduler_config()
         A__		 =					scheduler_class(**a__)

         A__		 =					[1_0_0, 8_7, 5_0, 5_1, 0]

         with self.assertRaises(a__ ,						msg='''`custom_timesteps` must be in descending order.'''):
             scheduler.set_timesteps(timesteps=a__)




     def 		snake_case_		(				self):
         A__		 =					self.scheduler_classes[0]
         A__		 =					self.get_scheduler_config()
         A__		 =					scheduler_class(**a__)

         A__		 =					[1_0_0, 8_7, 5_0, 1, 0]
         A__		 =					len(a__)

         with self.assertRaises(a__ ,						msg='''Can only pass one of `num_inference_steps` or `custom_timesteps`.'''):
             scheduler.set_timesteps(num_inference_steps=a__ ,						timesteps=a__)




     def 		snake_case_		(				self):
         A__		 =					self.scheduler_classes[0]
         A__		 =					self.get_scheduler_config()
         A__		 =					scheduler_class(**a__)

         A__		 =					[scheduler.config.num_train_timesteps]

         with self.assertRaises(
             a__ ,						msg='''`timesteps` must start before `self.config.train_timesteps`: {scheduler.config.num_train_timesteps}}''' ,						):
             scheduler.set_timesteps(timesteps=a__)
