'''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       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=512     ,       __lowerCAmelCase=16     ,       __lowerCAmelCase=2     ,       __lowerCAmelCase=0.0_2     ,       __lowerCAmelCase=3     ,       __lowerCAmelCase=4     ,       __lowerCAmelCase=None     ,       )  ->  str:
        lowercase__    :   Optional[int]						= parent
        lowercase__    :   List[Any]						= batch_size
        lowercase__    :   int						= seq_length
        lowercase__    :   int						= is_training
        lowercase__    :   int						= use_input_mask
        lowercase__    :   Optional[Any]						= use_token_type_ids
        lowercase__    :   List[str]						= use_labels
        lowercase__    :   str						= vocab_size
        lowercase__    :   int						= hidden_size
        lowercase__    :   Tuple						= num_hidden_layers
        lowercase__    :   List[Any]						= num_attention_heads
        lowercase__    :   Any						= intermediate_size
        lowercase__    :   List[Any]						= hidden_act
        lowercase__    :   Union[str, Any]						= hidden_dropout_prob
        lowercase__    :   List[str]						= attention_probs_dropout_prob
        lowercase__    :   Any						= max_position_embeddings
        lowercase__    :   List[Any]						= type_vocab_size
        lowercase__    :   Tuple						= type_sequence_label_size
        lowercase__    :   List[str]						= initializer_range
        lowercase__    :   Optional[int]						= num_labels
        lowercase__    :   List[str]						= num_choices
        lowercase__    :   Tuple						= scope

       def _lowerCAmelCase(						self			)  ->  Union[str, Any]:
        lowercase__    :   Union[str, Any]						= ids_tensor([self.batch_size, self.seq_length]     ,       self.vocab_size			)

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

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

        lowercase__    :   str						= self.get_config()

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

       def _lowerCAmelCase(						self			)  ->  str:
        lowercase__    :   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=__lowerCAmelCase     ,       esmfold_config={'''trunk''': {'''num_blocks''': 2}, '''fp16_esm''': False}     ,       )
        return config

       def _lowerCAmelCase(						self     ,       __lowerCAmelCase     ,       __lowerCAmelCase     ,       __lowerCAmelCase     ,       __lowerCAmelCase     ,       __lowerCAmelCase     ,       __lowerCAmelCase			)  ->  List[Any]:
        lowercase__    :   Dict						= EsmForProteinFolding(config=__lowerCAmelCase			).float()
        model.to(__lowerCAmelCase			)
        model.eval()
        lowercase__    :   Optional[int]						= model(__lowerCAmelCase     ,       attention_mask=__lowerCAmelCase			)
        lowercase__    :   Tuple						= model(__lowerCAmelCase			)
        lowercase__    :   Optional[Any]						= model(__lowerCAmelCase			)

        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			)  ->  Union[str, Any]:
        lowercase__    :   List[str]						= self.prepare_config_and_inputs()
        (
            (
            lowercase__
        )			,							(
            lowercase__
        )			,							(
            lowercase__
        )			,							(
            lowercase__
        )			,							(
            lowercase__
        )			,							(
            lowercase__
        )			,							
        )    :   Any						= config_and_inputs
        lowercase__    :   str						= {'''input_ids''': input_ids, '''attention_mask''': input_mask}
        return config, inputs_dict



@require_torch
class       UpperCAmelCase  (						a__   ,						a__   ,						unittest.TestCase ):





       '''simple docstring'''



       SCREAMING_SNAKE_CASE 		=   False

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

       def _lowerCAmelCase(						self			)  ->  str:
        lowercase__    :   List[Any]						= EsmFoldModelTester(self			)
        lowercase__    :   Dict						= ConfigTester(self     ,       config_class=__lowerCAmelCase     ,       hidden_size=37			)

       def _lowerCAmelCase(						self			)  ->  Optional[int]:
        self.config_tester.run_common_tests()

       def _lowerCAmelCase(						self			)  ->  List[Any]:
        lowercase__    :   Optional[int]						= self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_model(*__lowerCAmelCase			)

       @unittest.skip('''Does not support attention outputs'''			)
       def _lowerCAmelCase(						self			)  ->  Optional[Any]:
        pass

       @unittest.skip
       def _lowerCAmelCase(						self			)  ->  Optional[Any]:
        pass

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

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

       @unittest.skip('''ESMFold does not support passing input embeds!'''			)
       def _lowerCAmelCase(						self			)  ->  int:
        pass

       @unittest.skip('''ESMFold does not support head pruning.'''			)
       def _lowerCAmelCase(						self			)  ->  Dict:
        pass

       @unittest.skip('''ESMFold does not support head pruning.'''			)
       def _lowerCAmelCase(						self			)  ->  Any:
        pass

       @unittest.skip('''ESMFold does not support head pruning.'''			)
       def _lowerCAmelCase(						self			)  ->  Any:
        pass

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

       @unittest.skip('''ESMFold does not support head pruning.'''			)
       def _lowerCAmelCase(						self			)  ->  List[str]:
        pass

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

       @unittest.skip('''ESMfold does not output hidden states in the normal way.'''			)
       def _lowerCAmelCase(						self			)  ->  Union[str, Any]:
        pass

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

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

       @unittest.skip('''ESMFold does not support input chunking.'''			)
       def _lowerCAmelCase(						self			)  ->  Tuple:
        pass

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

       @unittest.skip('''ESMFold doesn\'t support torchscript compilation.'''			)
       def _lowerCAmelCase(						self			)  ->  List[Any]:
        pass

       @unittest.skip('''ESMFold doesn\'t support torchscript compilation.'''			)
       def _lowerCAmelCase(						self			)  ->  Any:
        pass

       @unittest.skip('''ESMFold doesn\'t support torchscript compilation.'''			)
       def _lowerCAmelCase(						self			)  ->  Optional[int]:
        pass

       @unittest.skip('''ESMFold doesn\'t support data parallel.'''			)
       def _lowerCAmelCase(						self			)  ->  Optional[int]:
        pass


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



@require_torch
class       UpperCAmelCase  (						a__ ):





       '''simple docstring'''



       @slow
       def _lowerCAmelCase(						self			)  ->  Any:
        lowercase__    :   int						= EsmForProteinFolding.from_pretrained('''facebook/esmfold_v1'''			).float()
        model.eval()
        lowercase__    :   Dict						= torch.tensor([[0, 6, 4, 13, 5, 4, 16, 12, 11, 7, 2]]			)
        lowercase__    :   Any						= model(__lowerCAmelCase			)['''positions''']
        lowercase__    :   Union[str, Any]						= torch.tensor([2.5_8_2_8, 0.7_9_9_3, -1_0.9_3_3_4]     ,       dtype=torch.floataa			)
        self.assertTrue(torch.allclose(position_outputs[0, 0, 0, 0]     ,       __lowerCAmelCase     ,       atol=1E-4			)			)


'''simple docstring'''

import re

from filelock import FileLock


try:
  import nltk

  __a:			Tuple	       =       True
except (ImportError, ModuleNotFoundError):
  __a:			List[Any]	       =       False

if NLTK_AVAILABLE:
  with FileLock(""".lock""") as lock:
    nltk.download("""punkt""", quiet=True)
def       __UpperCamelCase	(       UpperCAmelCase					):
 re.sub('''<n>'''		,					''''''		,					UpperCAmelCase					)  # remove pegasus newline char
 assert NLTK_AVAILABLE, "nltk must be installed to separate newlines between sentences. (pip install nltk)"
 return "\n".join(nltk.sent_tokenize(UpperCAmelCase					)					)

