





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







 '''simple docstring'''

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



   '''simple docstring'''



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

 def 							SCREAMING_SNAKE_CASE_					(self				:   List[str])       ->int:



   '''simple docstring'''



   lowerCamelCase__:      Tuple			  =ids_tensor([self.batch_size, self.seq_length]       , self.vocab_size)

   lowerCamelCase__:      Optional[Any]			  =None
   if self.use_input_mask:
     lowerCamelCase__:      Tuple			  =random_attention_mask([self.batch_size, self.seq_length])

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

   lowerCamelCase__:      Optional[Any]			  =self.get_config()

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

 def 							SCREAMING_SNAKE_CASE_					(self				:   Union[str, Any])       ->List[str]:



   '''simple docstring'''



   lowerCamelCase__:      Dict			  =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 							SCREAMING_SNAKE_CASE_					(self				:   Optional[Any]       , UpperCAmelCase_				:   int       , UpperCAmelCase_				:   List[str]       , UpperCAmelCase_				:   str       , UpperCAmelCase_				:   List[str]       , UpperCAmelCase_				:   List[str]       , UpperCAmelCase_				:   str)       ->List[Any]:



   '''simple docstring'''



   lowerCamelCase__:      str			  =EsmForProteinFolding(config=UpperCAmelCase_).float()
   model.to(UpperCAmelCase_)
   model.eval()
   lowerCamelCase__:      str			  =model(UpperCAmelCase_       , attention_mask=UpperCAmelCase_)
   lowerCamelCase__:      Optional[int]			  =model(UpperCAmelCase_)
   lowerCamelCase__:      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 							SCREAMING_SNAKE_CASE_					(self				:   List[Any])       ->List[Any]:



   '''simple docstring'''



   lowerCamelCase__:      Optional[Any]			  =self.prepare_config_and_inputs()
   (
       (
       lowerCamelCase__
   )  ,					(
       lowerCamelCase__
   )  ,					(
       lowerCamelCase__
   )  ,					(
       lowerCamelCase__
   )  ,					(
       lowerCamelCase__
   )  ,					(
       lowerCamelCase__
   )  ,					
   ):      Optional[Any]			  =config_and_inputs
   lowerCamelCase__:      Any			  ={"input_ids": input_ids, "attention_mask": input_mask}
   return config, inputs_dict



@require_torch
class    _SCREAMING_SNAKE_CASE						(	__SCREAMING_SNAKE_CASE    ,							__SCREAMING_SNAKE_CASE    ,							unittest.TestCase    ):







 '''simple docstring'''
 lowercase_		=			False

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

 def 							SCREAMING_SNAKE_CASE_					(self				:   Optional[Any])       ->Optional[int]:



   '''simple docstring'''



   lowerCamelCase__:      Optional[Any]			  =EsmFoldModelTester(self)
   lowerCamelCase__:      str			  =ConfigTester(self       , config_class=UpperCAmelCase_       , hidden_size=37)

 def 							SCREAMING_SNAKE_CASE_					(self				:   Union[str, Any])       ->List[Any]:



   '''simple docstring'''



   self.config_tester.run_common_tests()

 def 							SCREAMING_SNAKE_CASE_					(self				:   List[str])       ->Dict:



   '''simple docstring'''



   lowerCamelCase__:      Any			  =self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_model(*UpperCAmelCase_)

 @unittest.skip("Does not support attention outputs")
 def 							SCREAMING_SNAKE_CASE_					(self				:   List[Any])       ->Any:



   '''simple docstring'''



   pass

 @unittest.skip
 def 							SCREAMING_SNAKE_CASE_					(self				:   int)       ->Tuple:



   '''simple docstring'''



   pass

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



   '''simple docstring'''



   pass

 @unittest.skip("Esm does not support embedding resizing")
 def 							SCREAMING_SNAKE_CASE_					(self				:   Optional[int])       ->List[Any]:



   '''simple docstring'''



   pass

 @unittest.skip("ESMFold does not support passing input embeds!")
 def 							SCREAMING_SNAKE_CASE_					(self				:   Dict)       ->Dict:



   '''simple docstring'''



   pass

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



   '''simple docstring'''



   pass

 @unittest.skip("ESMFold does not support head pruning.")
 def 							SCREAMING_SNAKE_CASE_					(self				:   str)       ->Tuple:



   '''simple docstring'''



   pass

 @unittest.skip("ESMFold does not support head pruning.")
 def 							SCREAMING_SNAKE_CASE_					(self				:   Any)       ->Any:



   '''simple docstring'''



   pass

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



   '''simple docstring'''



   pass

 @unittest.skip("ESMFold does not support head pruning.")
 def 							SCREAMING_SNAKE_CASE_					(self				:   Union[str, Any])       ->Optional[int]:



   '''simple docstring'''



   pass

 @unittest.skip("ESMFold does not output hidden states in the normal way.")
 def 							SCREAMING_SNAKE_CASE_					(self				:   List[str])       ->Optional[Any]:



   '''simple docstring'''



   pass

 @unittest.skip("ESMfold does not output hidden states in the normal way.")
 def 							SCREAMING_SNAKE_CASE_					(self				:   List[str])       ->List[Any]:



   '''simple docstring'''



   pass

 @unittest.skip("ESMFold only has one output format.")
 def 							SCREAMING_SNAKE_CASE_					(self				:   Dict)       ->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[str])       ->Any:



   '''simple docstring'''



   pass

 @unittest.skip("ESMFold does not support input chunking.")
 def 							SCREAMING_SNAKE_CASE_					(self				:   Dict)       ->Optional[Any]:



   '''simple docstring'''



   pass

 @unittest.skip("ESMFold doesn't respect you and it certainly doesn't respect your initialization arguments.")
 def 							SCREAMING_SNAKE_CASE_					(self				:   int)       ->Union[str, Any]:



   '''simple docstring'''



   pass

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



   '''simple docstring'''



   pass

 @unittest.skip("ESMFold doesn't support torchscript compilation.")
 def 							SCREAMING_SNAKE_CASE_					(self				:   Any)       ->int:



   '''simple docstring'''



   pass

 @unittest.skip("ESMFold doesn't support torchscript compilation.")
 def 							SCREAMING_SNAKE_CASE_					(self				:   List[Any])       ->Optional[int]:



   '''simple docstring'''



   pass

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



   '''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				:   Union[str, Any])       ->int:



   '''simple docstring'''



   pass



@require_torch
class    _SCREAMING_SNAKE_CASE						(	__SCREAMING_SNAKE_CASE    ):







 '''simple docstring'''

 @slow
 def 							SCREAMING_SNAKE_CASE_					(self				:   Any)       ->Tuple:



   '''simple docstring'''



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








from collections import OrderedDict
from typing import Mapping

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig
from ...utils import logging


__lowerCamelCase       :     List[Any]   =     logging.get_logger(__name__)

__lowerCamelCase       :     str   =     {
    """roberta-base""": """https://huggingface.co/roberta-base/resolve/main/config.json""",
    """roberta-large""": """https://huggingface.co/roberta-large/resolve/main/config.json""",
    """roberta-large-mnli""": """https://huggingface.co/roberta-large-mnli/resolve/main/config.json""",
    """distilroberta-base""": """https://huggingface.co/distilroberta-base/resolve/main/config.json""",
    """roberta-base-openai-detector""": """https://huggingface.co/roberta-base-openai-detector/resolve/main/config.json""",
    """roberta-large-openai-detector""": """https://huggingface.co/roberta-large-openai-detector/resolve/main/config.json""",
}



class 					A__   (       __snake_case			):
      _UpperCAmelCase	:Union[str, Any]						    =   'roberta'





      def __init__( self       ,							A_=5_0265       ,							A_=768       ,							A_=12       ,							A_=12       ,							A_=3072       ,							A_="gelu"       ,							A_=0.1       ,							A_=0.1       ,							A_=512       ,							A_=2       ,							A_=0.02       ,							A_=1e-12       ,							A_=1       ,							A_=0       ,							A_=2       ,							A_="absolute"       ,							A_=True       ,							A_=None       ,							**A_       ,							):
       '''simple docstring'''




       super().__init__(pad_token_id=A_       ,							bos_token_id=A_       ,							eos_token_id=A_       ,							**A_   )

       UpperCamelCase	:					Optional[int]       	=   vocab_size
       UpperCamelCase	:					Dict       	=   hidden_size
       UpperCamelCase	:					str       	=   num_hidden_layers
       UpperCamelCase	:					Any       	=   num_attention_heads
       UpperCamelCase	:					List[str]       	=   hidden_act
       UpperCamelCase	:					Optional[Any]       	=   intermediate_size
       UpperCamelCase	:					Tuple       	=   hidden_dropout_prob
       UpperCamelCase	:					Tuple       	=   attention_probs_dropout_prob
       UpperCamelCase	:					Tuple       	=   max_position_embeddings
       UpperCamelCase	:					Any       	=   type_vocab_size
       UpperCamelCase	:					int       	=   initializer_range
       UpperCamelCase	:					str       	=   layer_norm_eps
       UpperCamelCase	:					Dict       	=   position_embedding_type
       UpperCamelCase	:					Any       	=   use_cache
       UpperCamelCase	:					Union[str, Any]       	=   classifier_dropout



class 					A__   (       __snake_case			):





      @property
      def        __UpperCamelCase( self   ):
       '''simple docstring'''




       if self.task == "multiple-choice":
        UpperCamelCase	:					Optional[int]       	=   {0: "batch", 1: "choice", 2: "sequence"}
       else:
        UpperCamelCase	:					Optional[int]       	=   {0: "batch", 1: "sequence"}
       return OrderedDict(
           [
               ("input_ids", dynamic_axis),
               ("attention_mask", dynamic_axis),
           ]   )
