






"""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 							lowerCAmelCase__       :






      '''simple docstring'''

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




          '''simple docstring'''


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




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_   :						Optional[int]							=	ids_tensor([self.batch_size, self.seq_length]		,			self.vocab_size)

          SCREAMING_SNAKE_CASE_   :						Optional[Any]							=	None
          if self.use_input_mask:
              SCREAMING_SNAKE_CASE_   :						Union[str, Any]							=	random_attention_mask([self.batch_size, self.seq_length])

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

          SCREAMING_SNAKE_CASE_   :						List[str]							=	self.get_config()

          return config, input_ids, input_mask, sequence_labels, token_labels, choice_labels
      def _SCREAMING_SNAKE_CASE				(		self  :						Any):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_   :						str							=	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=lowercase_		,			esmfold_config={'''trunk''': {'''num_blocks''': 2}, '''fp16_esm''': False}		,			)
          return config
      def _SCREAMING_SNAKE_CASE				(		self  :						List[Any]		,			lowercase_  :						List[Any]		,			lowercase_  :						str		,			lowercase_  :						Any		,			lowercase_  :						Dict		,			lowercase_  :						Dict		,			lowercase_  :						str):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_   :						Dict							=	EsmForProteinFolding(config=lowercase_).float()
          model.to(lowercase_)
          model.eval()
          SCREAMING_SNAKE_CASE_   :						int							=	model(lowercase_		,			attention_mask=lowercase_)
          SCREAMING_SNAKE_CASE_   :						Tuple							=	model(lowercase_)
          SCREAMING_SNAKE_CASE_   :						int							=	model(lowercase_)

          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  :						Tuple):




          '''simple docstring'''


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




@require_torch
class 							lowerCAmelCase__       (    UpperCAmelCase__       ,	UpperCAmelCase__       ,	unittest.TestCase				):






      '''simple docstring'''

      __UpperCamelCase							=						False

      __UpperCamelCase							=						(EsmForProteinFolding,) if is_torch_available() else ()
      __UpperCamelCase							=						()
      __UpperCamelCase							=						{} if is_torch_available() else {}
      __UpperCamelCase							=						False
      def _SCREAMING_SNAKE_CASE				(		self  :						Optional[int]):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_   :						List[str]							=	EsmFoldModelTester(self)
          SCREAMING_SNAKE_CASE_   :						Any							=	ConfigTester(self		,			config_class=lowercase_		,			hidden_size=37)
      def _SCREAMING_SNAKE_CASE				(		self  :						Tuple):




          '''simple docstring'''


          self.config_tester.run_common_tests()
      def _SCREAMING_SNAKE_CASE				(		self  :						List[str]):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_   :						Any							=	self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_model(*lowercase_)
      @unittest.skip('''Does not support attention outputs''')
      def _SCREAMING_SNAKE_CASE				(		self  :						List[str]):




          '''simple docstring'''


          pass
      @unittest.skip
      def _SCREAMING_SNAKE_CASE				(		self  :						int):




          '''simple docstring'''


          pass
      @unittest.skip('''Esm does not support embedding resizing''')
      def _SCREAMING_SNAKE_CASE				(		self  :						Union[str, Any]):




          '''simple docstring'''


          pass
      @unittest.skip('''Esm does not support embedding resizing''')
      def _SCREAMING_SNAKE_CASE				(		self  :						List[str]):




          '''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  :						Optional[int]):




          '''simple docstring'''


          pass
      @unittest.skip('''ESMFold does not support head pruning.''')
      def _SCREAMING_SNAKE_CASE				(		self  :						Optional[int]):




          '''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 support head pruning.''')
      def _SCREAMING_SNAKE_CASE				(		self  :						str):




          '''simple docstring'''


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




          '''simple docstring'''


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




          '''simple docstring'''


          pass
      @unittest.skip('''ESMFold only has one output format.''')
      def _SCREAMING_SNAKE_CASE				(		self  :						Union[str, Any]):




          '''simple docstring'''


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




          '''simple docstring'''


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




          '''simple docstring'''


          pass
      @unittest.skip('''ESMFold doesn\'t respect you and it certainly doesn\'t respect your initialization arguments.''')
      def _SCREAMING_SNAKE_CASE				(		self  :						List[str]):




          '''simple docstring'''


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




          '''simple docstring'''


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




          '''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 data parallel.''')
      def _SCREAMING_SNAKE_CASE				(		self  :						List[str]):




          '''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 							lowerCAmelCase__       (    UpperCAmelCase__				):






      '''simple docstring'''

      @slow
      def _SCREAMING_SNAKE_CASE				(		self  :						Any):




          '''simple docstring'''


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










"""simple docstring"""






def       _A    (__a = 1_00							)				-> int:
    """simple docstring"""






    SCREAMING_SNAKE_CASE_   :						int							=	(n * (n + 1) // 2) ** 2
    SCREAMING_SNAKE_CASE_   :						Optional[Any]							=	n * (n + 1) * (2 * n + 1) // 6
    return sum_cubes - sum_squares


if __name__ == "__main__":
  print(f'''{solution() = }''')


