

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   __lowercase      :
 def __init__(							self      ,							A_      ,							A_=13      ,							A_=7      ,							A_=False      ,							A_=True      ,							A_=False      ,							A_=False      ,							A_=19      ,							A_=32      ,							A_=5      ,							A_=4      ,							A_=37      ,							A_="gelu"      ,							A_=0.1      ,							A_=0.1      ,							A_=512      ,							A_=16      ,							A_=2      ,							A_=0.02      ,							A_=3      ,							A_=4      ,							A_=None      ,							)					->Optional[int]:



       '''simple docstring'''





       __lowerCAmelCase					:						Optional[int]					=					parent
       __lowerCAmelCase					:						Optional[int]					=					batch_size
       __lowerCAmelCase					:						Any					=					seq_length
       __lowerCAmelCase					:						Tuple					=					is_training
       __lowerCAmelCase					:						int					=					use_input_mask
       __lowerCAmelCase					:						Optional[Any]					=					use_token_type_ids
       __lowerCAmelCase					:						List[str]					=					use_labels
       __lowerCAmelCase					:						Optional[int]					=					vocab_size
       __lowerCAmelCase					:						Optional[int]					=					hidden_size
       __lowerCAmelCase					:						Any					=					num_hidden_layers
       __lowerCAmelCase					:						Optional[Any]					=					num_attention_heads
       __lowerCAmelCase					:						Union[str, Any]					=					intermediate_size
       __lowerCAmelCase					:						Tuple					=					hidden_act
       __lowerCAmelCase					:						Optional[Any]					=					hidden_dropout_prob
       __lowerCAmelCase					:						Tuple					=					attention_probs_dropout_prob
       __lowerCAmelCase					:						List[Any]					=					max_position_embeddings
       __lowerCAmelCase					:						List[Any]					=					type_vocab_size
       __lowerCAmelCase					:						Tuple					=					type_sequence_label_size
       __lowerCAmelCase					:						int					=					initializer_range
       __lowerCAmelCase					:						Dict					=					num_labels
       __lowerCAmelCase					:						List[str]					=					num_choices
       __lowerCAmelCase					:						str					=					scope
 def 	UpperCamelCase__  (							self					)					->str:



       '''simple docstring'''





       __lowerCAmelCase					:						str					=					ids_tensor([self.batch_size, self.seq_length]      ,							self.vocab_size					)

       __lowerCAmelCase					:						str					=					None
       if self.use_input_mask:
             __lowerCAmelCase					:						List[str]					=					random_attention_mask([self.batch_size, self.seq_length]					)

       __lowerCAmelCase					:						Optional[int]					=					None
       __lowerCAmelCase					:						int					=					None
       __lowerCAmelCase					:						Tuple					=					None
       if self.use_labels:
             __lowerCAmelCase					:						str					=					ids_tensor([self.batch_size]      ,							self.type_sequence_label_size					)
             __lowerCAmelCase					:						Any					=					ids_tensor([self.batch_size, self.seq_length]      ,							self.num_labels					)
             __lowerCAmelCase					:						str					=					ids_tensor([self.batch_size]      ,							self.num_choices					)

       __lowerCAmelCase					:						Union[str, Any]					=					self.get_config()

       return config, input_ids, input_mask, sequence_labels, token_labels, choice_labels
 def 	UpperCamelCase__  (							self					)					->str:



       '''simple docstring'''





       __lowerCAmelCase					:						Optional[Any]					=					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 	UpperCamelCase__  (							self      ,							A_      ,							A_      ,							A_      ,							A_      ,							A_      ,							A_					)					->str:



       '''simple docstring'''





       __lowerCAmelCase					:						Dict					=					EsmForProteinFolding(config=A_					).float()
       model.to(A_					)
       model.eval()
       __lowerCAmelCase					:						Union[str, Any]					=					model(A_      ,							attention_mask=A_					)
       __lowerCAmelCase					:						Any					=					model(A_					)
       __lowerCAmelCase					:						Optional[Any]					=					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 	UpperCamelCase__  (							self					)					->str:



       '''simple docstring'''





       __lowerCAmelCase					:						List[Any]					=					self.prepare_config_and_inputs()
       (
           (
           __lowerCAmelCase
       ),							(
           __lowerCAmelCase
       ),							(
           __lowerCAmelCase
       ),							(
           __lowerCAmelCase
       ),							(
           __lowerCAmelCase
       ),							(
           __lowerCAmelCase
       ),							
       )					:						Union[str, Any]					=					config_and_inputs
       __lowerCAmelCase					:						List[Any]					=					{'''input_ids''': input_ids, '''attention_mask''': input_mask}
       return config, inputs_dict







@require_torch
class   __lowercase      (_UpperCAmelCase		,   _UpperCAmelCase		,   unittest.TestCase ):
 _UpperCamelCase			  =  False

 _UpperCamelCase			  =  (EsmForProteinFolding,) if is_torch_available() else ()
 _UpperCamelCase			  =  ()
 _UpperCamelCase			  =  {} if is_torch_available() else {}
 _UpperCamelCase			  =  False
 def 	UpperCamelCase__  (							self					)					->Optional[int]:



       '''simple docstring'''





       __lowerCAmelCase					:						List[str]					=					EsmFoldModelTester(self					)
       __lowerCAmelCase					:						int					=					ConfigTester(self      ,							config_class=A_      ,							hidden_size=37					)
 def 	UpperCamelCase__  (							self					)					->List[str]:



       '''simple docstring'''





       self.config_tester.run_common_tests()
 def 	UpperCamelCase__  (							self					)					->int:



       '''simple docstring'''





       __lowerCAmelCase					:						List[str]					=					self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_model(*A_					)
 @unittest.skip('''Does not support attention outputs'''					)
 def 	UpperCamelCase__  (							self					)					->Optional[Any]:



       '''simple docstring'''





       pass
 @unittest.skip
 def 	UpperCamelCase__  (							self					)					->Dict:



       '''simple docstring'''





       pass
 @unittest.skip('''Esm does not support embedding resizing'''					)
 def 	UpperCamelCase__  (							self					)					->int:



       '''simple docstring'''





       pass
 @unittest.skip('''Esm does not support embedding resizing'''					)
 def 	UpperCamelCase__  (							self					)					->Optional[Any]:



       '''simple docstring'''





       pass
 @unittest.skip('''ESMFold does not support passing input embeds!'''					)
 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 support head pruning.'''					)
 def 	UpperCamelCase__  (							self					)					->Optional[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					)					->List[Any]:



       '''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					)					->str:



       '''simple docstring'''





       pass
 @unittest.skip('''ESMFold only has one output format.'''					)
 def 	UpperCamelCase__  (							self					)					->Dict:



       '''simple docstring'''





       pass
 @unittest.skip('''This test doesn\'t work for ESMFold and doesn\'t test core functionality'''					)
 def 	UpperCamelCase__  (							self					)					->Optional[Any]:



       '''simple docstring'''





       pass
 @unittest.skip('''ESMFold does not support input chunking.'''					)
 def 	UpperCamelCase__  (							self					)					->List[str]:



       '''simple docstring'''





       pass
 @unittest.skip('''ESMFold doesn\'t respect you and it certainly doesn\'t respect your initialization arguments.'''					)
 def 	UpperCamelCase__  (							self					)					->Optional[int]:



       '''simple docstring'''





       pass
 @unittest.skip('''ESMFold doesn\'t support torchscript compilation.'''					)
 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					)					->Dict:



       '''simple docstring'''





       pass
 @unittest.skip('''ESMFold doesn\'t support data parallel.'''					)
 def 	UpperCamelCase__  (							self					)					->Tuple:



       '''simple docstring'''





       pass





 @unittest.skip('''Will be fixed soon by reducing the size of the model used for common tests.'''					)
 def 	UpperCamelCase__  (							self					)					->List[str]:



       '''simple docstring'''





       pass







@require_torch
class   __lowercase      (_UpperCAmelCase ):
 @slow
 def 	UpperCamelCase__  (							self					)					->List[str]:



       '''simple docstring'''





       __lowerCAmelCase					:						Dict					=					EsmForProteinFolding.from_pretrained('''facebook/esmfold_v1'''					).float()
       model.eval()
       __lowerCAmelCase					:						Optional[Any]					=					torch.tensor([[0, 6, 4, 13, 5, 4, 16, 12, 11, 7, 2]]					)
       __lowerCAmelCase					:						List[str]					=					model(A_					)['''positions''']
       __lowerCAmelCase					:						Any					=					torch.tensor([2.5_828, 0.7_993, -10.9_334]      ,							dtype=torch.floataa					)
       self.assertTrue(torch.allclose(position_outputs[0, 0, 0, 0]      ,							A_      ,							atol=1e-4					)					)





from dataclasses import dataclass
from typing import Optional, Tuple, Union

import numpy as np
import torch

from ..configuration_utils import ConfigMixin, register_to_config
from ..utils import BaseOutput, randn_tensor
from .scheduling_utils import SchedulerMixin







@dataclass
class   __lowercase      (_UpperCAmelCase ):
 _UpperCamelCase			  =  42
 _UpperCamelCase			  =  42
 _UpperCamelCase			  =  None







class   __lowercase      (_UpperCAmelCase		,   _UpperCAmelCase ):
 _UpperCamelCase			  =  2
 @register_to_config
 def __init__(							self      ,							A_ = 0.02      ,							A_ = 100      ,							A_ = 1.007      ,							A_ = 80      ,							A_ = 0.05      ,							A_ = 50      ,							)					->int:



       '''simple docstring'''





       __lowerCAmelCase					:						Optional[int]					=					sigma_max

       # setable values
       __lowerCAmelCase					:						int					=					None
       __lowerCAmelCase					:						np.IntTensor					=					None
       __lowerCAmelCase					:						torch.FloatTensor					=					None  # sigma(t_i)
 def 	UpperCamelCase__  (							self      ,							A_      ,							A_ = None					)					->torch.FloatTensor:



       '''simple docstring'''





       return sample
 def 	UpperCamelCase__  (							self      ,							A_      ,							A_ = None					)					->List[str]:



       '''simple docstring'''





       __lowerCAmelCase					:						str					=					num_inference_steps
       __lowerCAmelCase					:						Dict					=					np.arange(0      ,							self.num_inference_steps					)[::-1].copy()
       __lowerCAmelCase					:						Optional[Any]					=					torch.from_numpy(A_					).to(A_					)
       __lowerCAmelCase					:						Tuple					=					[
           (
               self.config.sigma_max**2
               * (self.config.sigma_min**2 / self.config.sigma_max**2) ** (i / (num_inference_steps - 1))
           )
           for i in self.timesteps
       ]
       __lowerCAmelCase					:						Optional[int]					=					torch.tensor(A_      ,							dtype=torch.floataa      ,							device=A_					)
 def 	UpperCamelCase__  (							self      ,							A_      ,							A_      ,							A_ = None					)					->Tuple[torch.FloatTensor, float]:



       '''simple docstring'''





       if self.config.s_min <= sigma <= self.config.s_max:
             __lowerCAmelCase					:						List[str]					=					min(self.config.s_churn / self.num_inference_steps      ,							2**0.5 - 1					)
       else:
             __lowerCAmelCase					:						List[str]					=					0

       # sample eps ~ N(0, S_noise^2 * I)
       __lowerCAmelCase					:						int					=					self.config.s_noise * randn_tensor(sample.shape      ,							generator=A_					).to(sample.device					)
       __lowerCAmelCase					:						str					=					sigma + gamma * sigma
       __lowerCAmelCase					:						Any					=					sample + ((sigma_hat**2 - sigma**2) ** 0.5 * eps)

       return sample_hat, sigma_hat
 def 	UpperCamelCase__  (							self      ,							A_      ,							A_      ,							A_      ,							A_      ,							A_ = True      ,							)					->Union[KarrasVeOutput, Tuple]:



       '''simple docstring'''





       __lowerCAmelCase					:						Union[str, Any]					=					sample_hat + sigma_hat * model_output
       __lowerCAmelCase					:						int					=					(sample_hat - pred_original_sample) / sigma_hat
       __lowerCAmelCase					:						Tuple					=					sample_hat + (sigma_prev - sigma_hat) * derivative

       if not return_dict:
             return (sample_prev, derivative)

       return KarrasVeOutput(
           prev_sample=A_      ,							derivative=A_      ,							pred_original_sample=A_					)
 def 	UpperCamelCase__  (							self      ,							A_      ,							A_      ,							A_      ,							A_      ,							A_      ,							A_      ,							A_ = True      ,							)					->Union[KarrasVeOutput, Tuple]:



       '''simple docstring'''





       __lowerCAmelCase					:						str					=					sample_prev + sigma_prev * model_output
       __lowerCAmelCase					:						List[Any]					=					(sample_prev - pred_original_sample) / sigma_prev
       __lowerCAmelCase					:						Any					=					sample_hat + (sigma_prev - sigma_hat) * (0.5 * derivative + 0.5 * derivative_corr)

       if not return_dict:
             return (sample_prev, derivative)

       return KarrasVeOutput(
           prev_sample=A_      ,							derivative=A_      ,							pred_original_sample=A_					)





 def 	UpperCamelCase__  (							self      ,							A_      ,							A_      ,							A_					)					->Any:



       '''simple docstring'''





       raise NotImplementedError()


