






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:  Optional[int]				,     UpperCAmelCase_:  Optional[Any]				,     UpperCAmelCase_:  Optional[int]=13				,     UpperCAmelCase_:  Any=7				,     UpperCAmelCase_:  Optional[Any]=False				,     UpperCAmelCase_:  Any=True				,     UpperCAmelCase_:  Dict=False				,     UpperCAmelCase_:  Tuple=False				,     UpperCAmelCase_:  Tuple=19				,     UpperCAmelCase_:  List[Any]=32				,     UpperCAmelCase_:  Optional[Any]=5				,     UpperCAmelCase_:  Dict=4				,     UpperCAmelCase_:  Tuple=37				,     UpperCAmelCase_:  str="gelu"				,     UpperCAmelCase_:  Tuple=0.1				,     UpperCAmelCase_:  Tuple=0.1				,     UpperCAmelCase_:  int=512				,     UpperCAmelCase_:  Dict=16				,     UpperCAmelCase_:  List[Any]=2				,     UpperCAmelCase_:  Optional[int]=0.02				,     UpperCAmelCase_:  Tuple=3				,     UpperCAmelCase_:  List[str]=4				,     UpperCAmelCase_:  Dict=None				,     ):







         '''simple docstring'''




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







   def 							UpperCamelCase				(     self:  Optional[Any]					):







         '''simple docstring'''




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

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

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

         _SCREAMING_SNAKE_CASE     				=  self.get_config()

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







   def 							UpperCamelCase				(     self:  Any					):







         '''simple docstring'''




         _SCREAMING_SNAKE_CASE     				=  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=UpperCAmelCase_				,     esmfold_config={"""trunk""": {"""num_blocks""": 2}, """fp16_esm""": False}				,     )
         return config







   def 							UpperCamelCase				(     self:  int				,     UpperCAmelCase_:  Optional[Any]				,     UpperCAmelCase_:  str				,     UpperCAmelCase_:  Union[str, Any]				,     UpperCAmelCase_:  Union[str, Any]				,     UpperCAmelCase_:  Optional[int]				,     UpperCAmelCase_:  Optional[int]					):







         '''simple docstring'''




         _SCREAMING_SNAKE_CASE     				=  EsmForProteinFolding(config=UpperCAmelCase_					).float()
         model.to(UpperCAmelCase_					)
         model.eval()
         _SCREAMING_SNAKE_CASE     				=  model(UpperCAmelCase_				,     attention_mask=UpperCAmelCase_					)
         _SCREAMING_SNAKE_CASE     				=  model(UpperCAmelCase_					)
         _SCREAMING_SNAKE_CASE     				=  model(UpperCAmelCase_					)

         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 							UpperCamelCase				(     self:  int					):







         '''simple docstring'''




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



@require_torch
class 				__UpperCAmelCase		(_UpperCAmelCase	,_UpperCAmelCase	,unittest.TestCase   ):
   __snake_case    :				List[str]									=							False

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







   def 							UpperCamelCase				(     self:  List[Any]					):







         '''simple docstring'''




         _SCREAMING_SNAKE_CASE     				=  EsmFoldModelTester(self					)
         _SCREAMING_SNAKE_CASE     				=  ConfigTester(self				,     config_class=UpperCAmelCase_				,     hidden_size=37					)







   def 							UpperCamelCase				(     self:  Tuple					):







         '''simple docstring'''




         self.config_tester.run_common_tests()







   def 							UpperCamelCase				(     self:  List[Any]					):







         '''simple docstring'''




         _SCREAMING_SNAKE_CASE     				=  self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_model(*UpperCAmelCase_					)







   @unittest.skip("""Does not support attention outputs"""					)
   def 							UpperCamelCase				(     self:  Tuple					):







         '''simple docstring'''




         pass







   @unittest.skip
   def 							UpperCamelCase				(     self:  List[Any]					):







         '''simple docstring'''




         pass







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







         '''simple docstring'''




         pass







   @unittest.skip("""Esm does not support embedding resizing"""					)
   def 							UpperCamelCase				(     self:  str					):







         '''simple docstring'''




         pass







   @unittest.skip("""ESMFold does not support passing input embeds!"""					)
   def 							UpperCamelCase				(     self:  int					):







         '''simple docstring'''




         pass







   @unittest.skip("""ESMFold does not support head pruning."""					)
   def 							UpperCamelCase				(     self:  Optional[int]					):







         '''simple docstring'''




         pass







   @unittest.skip("""ESMFold does not support head pruning."""					)
   def 							UpperCamelCase				(     self:  Union[str, Any]					):







         '''simple docstring'''




         pass







   @unittest.skip("""ESMFold does not support head pruning."""					)
   def 							UpperCamelCase				(     self:  Optional[Any]					):







         '''simple docstring'''




         pass







   @unittest.skip("""ESMFold does not support head pruning."""					)
   def 							UpperCamelCase				(     self:  Union[str, Any]					):







         '''simple docstring'''




         pass







   @unittest.skip("""ESMFold does not support head pruning."""					)
   def 							UpperCamelCase				(     self:  Optional[int]					):







         '''simple docstring'''




         pass







   @unittest.skip("""ESMFold does not output hidden states in the normal way."""					)
   def 							UpperCamelCase				(     self:  Any					):







         '''simple docstring'''




         pass







   @unittest.skip("""ESMfold does not output hidden states in the normal way."""					)
   def 							UpperCamelCase				(     self:  Dict					):







         '''simple docstring'''




         pass







   @unittest.skip("""ESMFold only has one output format."""					)
   def 							UpperCamelCase				(     self:  List[Any]					):







         '''simple docstring'''




         pass







   @unittest.skip("""This test doesn't work for ESMFold and doesn't test core functionality"""					)
   def 							UpperCamelCase				(     self:  Union[str, Any]					):







         '''simple docstring'''




         pass







   @unittest.skip("""ESMFold does not support input chunking."""					)
   def 							UpperCamelCase				(     self:  Optional[int]					):







         '''simple docstring'''




         pass







   @unittest.skip("""ESMFold doesn't respect you and it certainly doesn't respect your initialization arguments."""					)
   def 							UpperCamelCase				(     self:  Any					):







         '''simple docstring'''




         pass







   @unittest.skip("""ESMFold doesn't support torchscript compilation."""					)
   def 							UpperCamelCase				(     self:  Union[str, Any]					):







         '''simple docstring'''




         pass







   @unittest.skip("""ESMFold doesn't support torchscript compilation."""					)
   def 							UpperCamelCase				(     self:  List[str]					):







         '''simple docstring'''




         pass







   @unittest.skip("""ESMFold doesn't support torchscript compilation."""					)
   def 							UpperCamelCase				(     self:  List[str]					):







         '''simple docstring'''




         pass







   @unittest.skip("""ESMFold doesn't support data parallel."""					)
   def 							UpperCamelCase				(     self:  Union[str, Any]					):







         '''simple docstring'''




         pass







   @unittest.skip("""Will be fixed soon by reducing the size of the model used for common tests."""					)
   def 							UpperCamelCase				(     self:  int					):







         '''simple docstring'''




         pass



@require_torch
class 				__UpperCAmelCase		(_UpperCAmelCase   ):







   @slow
   def 							UpperCamelCase				(     self:  List[Any]					):







         '''simple docstring'''




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









import torch

from diffusers import UnCLIPScheduler

from .test_schedulers import SchedulerCommonTest



class 				__UpperCAmelCase		(_UpperCAmelCase   ):
			__snake_case    :				Optional[Any]									=							(UnCLIPScheduler,)







			def 							UpperCamelCase				(     self:  int				,     **UpperCAmelCase_:  Dict					):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  {
									    """num_train_timesteps""": 1_000,
									    """variance_type""": """fixed_small_log""",
									    """clip_sample""": True,
									    """clip_sample_range""": 1.0,
									    """prediction_type""": """epsilon""",
									}

									config.update(**UpperCAmelCase_					)
									return config







			def 							UpperCamelCase				(     self:  Optional[Any]					):







									'''simple docstring'''




									for timesteps in [1, 5, 100, 1_000]:
															self.check_over_configs(num_train_timesteps=UpperCAmelCase_					)







			def 							UpperCamelCase				(     self:  Any					):







									'''simple docstring'''




									for variance in ["fixed_small_log", "learned_range"]:
															self.check_over_configs(variance_type=UpperCAmelCase_					)







			def 							UpperCamelCase				(     self:  List[Any]					):







									'''simple docstring'''




									for clip_sample in [True, False]:
															self.check_over_configs(clip_sample=UpperCAmelCase_					)







			def 							UpperCamelCase				(     self:  Optional[int]					):







									'''simple docstring'''




									for clip_sample_range in [1, 5, 10, 20]:
															self.check_over_configs(clip_sample_range=UpperCAmelCase_					)







			def 							UpperCamelCase				(     self:  Union[str, Any]					):







									'''simple docstring'''




									for prediction_type in ["epsilon", "sample"]:
															self.check_over_configs(prediction_type=UpperCAmelCase_					)







			def 							UpperCamelCase				(     self:  Optional[Any]					):







									'''simple docstring'''




									for time_step in [0, 500, 999]:
															for prev_timestep in [None, 5, 100, 250, 500, 750]:
																					if prev_timestep is not None and prev_timestep >= time_step:
																											continue

																					self.check_over_forward(time_step=UpperCAmelCase_				,     prev_timestep=UpperCAmelCase_					)







			def 							UpperCamelCase				(     self:  Any					):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  self.scheduler_classes[0]
									_SCREAMING_SNAKE_CASE     				=  self.get_scheduler_config(variance_type="""fixed_small_log"""					)
									_SCREAMING_SNAKE_CASE     				=  scheduler_class(**UpperCAmelCase_					)

									assert torch.sum(torch.abs(scheduler._get_variance(0					) - 1.0_000E-10					)					) < 1E-5
									assert torch.sum(torch.abs(scheduler._get_variance(487					) - 0.0_54_96_25					)					) < 1E-5
									assert torch.sum(torch.abs(scheduler._get_variance(999					) - 0.9_99_49_87					)					) < 1E-5







			def 							UpperCamelCase				(     self:  int					):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  self.scheduler_classes[0]
									_SCREAMING_SNAKE_CASE     				=  self.get_scheduler_config(variance_type="""learned_range"""					)
									_SCREAMING_SNAKE_CASE     				=  scheduler_class(**UpperCAmelCase_					)

									_SCREAMING_SNAKE_CASE     				=  0.5

									assert scheduler._get_variance(1				,     predicted_variance=UpperCAmelCase_					) - -10.1_71_27_90 < 1E-5
									assert scheduler._get_variance(487				,     predicted_variance=UpperCAmelCase_					) - -5.7_99_80_52 < 1E-5
									assert scheduler._get_variance(999				,     predicted_variance=UpperCAmelCase_					) - -0.0_01_00_11 < 1E-5







			def 							UpperCamelCase				(     self:  List[Any]					):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  self.scheduler_classes[0]
									_SCREAMING_SNAKE_CASE     				=  self.get_scheduler_config()
									_SCREAMING_SNAKE_CASE     				=  scheduler_class(**UpperCAmelCase_					)

									_SCREAMING_SNAKE_CASE     				=  scheduler.timesteps

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

									for i, t in enumerate(UpperCAmelCase_					):
															# 1. predict noise residual
															_SCREAMING_SNAKE_CASE     				=  model(UpperCAmelCase_				,     UpperCAmelCase_					)

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

															_SCREAMING_SNAKE_CASE     				=  pred_prev_sample

									_SCREAMING_SNAKE_CASE     				=  torch.sum(torch.abs(UpperCAmelCase_					)					)
									_SCREAMING_SNAKE_CASE     				=  torch.mean(torch.abs(UpperCAmelCase_					)					)

									assert abs(result_sum.item() - 2_52.2_68_24_95					) < 1E-2
									assert abs(result_mean.item() - 0.3_28_47_43					) < 1E-3







			def 							UpperCamelCase				(     self:  List[str]					):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  self.scheduler_classes[0]
									_SCREAMING_SNAKE_CASE     				=  self.get_scheduler_config()
									_SCREAMING_SNAKE_CASE     				=  scheduler_class(**UpperCAmelCase_					)

									scheduler.set_timesteps(25					)

									_SCREAMING_SNAKE_CASE     				=  scheduler.timesteps

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

									for i, t in enumerate(UpperCAmelCase_					):
															# 1. predict noise residual
															_SCREAMING_SNAKE_CASE     				=  model(UpperCAmelCase_				,     UpperCAmelCase_					)

															if i + 1 == timesteps.shape[0]:
																					_SCREAMING_SNAKE_CASE     				=  None
															else:
																					_SCREAMING_SNAKE_CASE     				=  timesteps[i + 1]

															# 2. predict previous mean of sample x_t-1
															_SCREAMING_SNAKE_CASE     				=  scheduler.step(
															    UpperCAmelCase_				,     UpperCAmelCase_				,     UpperCAmelCase_				,     prev_timestep=UpperCAmelCase_				,     generator=UpperCAmelCase_					).prev_sample

															_SCREAMING_SNAKE_CASE     				=  pred_prev_sample

									_SCREAMING_SNAKE_CASE     				=  torch.sum(torch.abs(UpperCAmelCase_					)					)
									_SCREAMING_SNAKE_CASE     				=  torch.mean(torch.abs(UpperCAmelCase_					)					)

									assert abs(result_sum.item() - 2_58.2_04_49_83					) < 1E-2
									assert abs(result_mean.item() - 0.3_36_20_38					) < 1E-3







			def 							UpperCamelCase				(     self:  List[Any]					):







									'''simple docstring'''




									pass







			def 							UpperCamelCase				(     self:  str					):







									'''simple docstring'''




									pass

