
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 :

     """simple docstring"""

     def __init__(      self							,		lowerCAmelCase_							,		lowerCAmelCase_=13							,		lowerCAmelCase_=7							,		lowerCAmelCase_=False							,		lowerCAmelCase_=True							,		lowerCAmelCase_=False							,		lowerCAmelCase_=False							,		lowerCAmelCase_=19							,		lowerCAmelCase_=32							,		lowerCAmelCase_=5							,		lowerCAmelCase_=4							,		lowerCAmelCase_=37							,		lowerCAmelCase_="gelu"							,		lowerCAmelCase_=0.1							,		lowerCAmelCase_=0.1							,		lowerCAmelCase_=5_12							,		lowerCAmelCase_=16							,		lowerCAmelCase_=2							,		lowerCAmelCase_=0.02							,		lowerCAmelCase_=3							,		lowerCAmelCase_=4							,		lowerCAmelCase_=None							,		):

       '''simple docstring'''

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


     def 							_lowerCAmelCase  (      self ):

       '''simple docstring'''

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

       a_      :	Tuple      =							None
       if self.use_input_mask:
         a_      :	int      =							random_attention_mask([self.batch_size, self.seq_length] )

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

       a_      :	List[Any]      =							self.get_config()

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


     def 							_lowerCAmelCase  (      self ):

       '''simple docstring'''

       a_      :	Optional[int]      =							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 							_lowerCAmelCase  (      self							,		lowerCAmelCase_							,		lowerCAmelCase_							,		lowerCAmelCase_							,		lowerCAmelCase_							,		lowerCAmelCase_							,		lowerCAmelCase_ ):

       '''simple docstring'''

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

       '''simple docstring'''

       a_      :	Tuple      =							self.prepare_config_and_inputs()
       (
           a_							
       )      :	str      =							config_and_inputs
       a_      :	Dict      =							{'''input_ids''': input_ids, '''attention_mask''': input_mask}
       return config, inputs_dict







@require_torch
class  _UpperCAmelCase ( SCREAMING_SNAKE_CASE_	,SCREAMING_SNAKE_CASE_	,unittest.TestCase    ):

     """simple docstring"""

     a_    	=					False

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


     def 							_lowerCAmelCase  (      self ):

       '''simple docstring'''

       a_      :	Optional[Any]      =							EsmFoldModelTester(self )
       a_      :	List[str]      =							ConfigTester(self							,		config_class=UpperCamelCase__							,		hidden_size=37 )


     def 							_lowerCAmelCase  (      self ):

       '''simple docstring'''

       self.config_tester.run_common_tests()


     def 							_lowerCAmelCase  (      self ):

       '''simple docstring'''

       a_      :	List[str]      =							self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_model(*UpperCamelCase__ )


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

       '''simple docstring'''

       pass


     @unittest.skip
     def 							_lowerCAmelCase  (      self ):

       '''simple docstring'''

       pass


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

       '''simple docstring'''

       pass


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

       '''simple docstring'''

       pass


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

       '''simple docstring'''

       pass


     @unittest.skip("""ESMFold does not support head pruning.""" )
     def 							_lowerCAmelCase  (      self ):

       '''simple docstring'''

       pass


     @unittest.skip("""ESMFold does not support head pruning.""" )
     def 							_lowerCAmelCase  (      self ):

       '''simple docstring'''

       pass


     @unittest.skip("""ESMFold does not support head pruning.""" )
     def 							_lowerCAmelCase  (      self ):

       '''simple docstring'''

       pass


     @unittest.skip("""ESMFold does not support head pruning.""" )
     def 							_lowerCAmelCase  (      self ):

       '''simple docstring'''

       pass


     @unittest.skip("""ESMFold does not support head pruning.""" )
     def 							_lowerCAmelCase  (      self ):

       '''simple docstring'''

       pass


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

       '''simple docstring'''

       pass


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

       '''simple docstring'''

       pass


     @unittest.skip("""ESMFold only has one output format.""" )
     def 							_lowerCAmelCase  (      self ):

       '''simple docstring'''

       pass


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

       '''simple docstring'''

       pass


     @unittest.skip("""ESMFold does not support input chunking.""" )
     def 							_lowerCAmelCase  (      self ):

       '''simple docstring'''

       pass


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

       '''simple docstring'''

       pass


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

       '''simple docstring'''

       pass


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

       '''simple docstring'''

       pass


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

       '''simple docstring'''

       pass


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

       '''simple docstring'''

       pass






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

       '''simple docstring'''

       pass







@require_torch
class  _UpperCAmelCase ( SCREAMING_SNAKE_CASE_    ):

     """simple docstring"""

     @slow
     def 							_lowerCAmelCase  (      self ):

       '''simple docstring'''

       a_      :	Any      =							EsmForProteinFolding.from_pretrained("""facebook/esmfold_v1""" ).float()
       model.eval()
       a_      :	List[Any]      =							torch.tensor([[0, 6, 4, 13, 5, 4, 16, 12, 11, 7, 2]] )
       a_      :	List[Any]      =							model(UpperCamelCase__ )['''positions''']
       a_      :	Optional[int]      =							torch.tensor([2.5828, 0.7993, -10.93_34]							,		dtype=torch.floataa )
       self.assertTrue(torch.allclose(position_outputs[0, 0, 0, 0]							,		UpperCamelCase__							,		atol=1E-4 ) )




'''simple docstring'''

from torch import nn






def 	_snake_case						(      A_       :    Dict   ):
  """simple docstring"""







  if act_fn in ["swish", "silu"]:
    return nn.SiLU()
  elif act_fn == "mish":
    return nn.Mish()
  elif act_fn == "gelu":
    return nn.GELU()
  else:
    raise ValueError(f'''Unsupported activation function: {act_fn}'''   )


