



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			:




     def __init__(				self ,						a__ ,						a__=1_3 ,						a__=7 ,						a__=False ,						a__=True ,						a__=False ,						a__=False ,						a__=1_9 ,						a__=3_2 ,						a__=5 ,						a__=4 ,						a__=3_7 ,						a__="gelu" ,						a__=0.1 ,						a__=0.1 ,						a__=5_1_2 ,						a__=1_6 ,						a__=2 ,						a__=0.0_2 ,						a__=3 ,						a__=4 ,						a__=None ,						):
         A__		 =					parent
         A__		 =					batch_size
         A__		 =					seq_length
         A__		 =					is_training
         A__		 =					use_input_mask
         A__		 =					use_token_type_ids
         A__		 =					use_labels
         A__		 =					vocab_size
         A__		 =					hidden_size
         A__		 =					num_hidden_layers
         A__		 =					num_attention_heads
         A__		 =					intermediate_size
         A__		 =					hidden_act
         A__		 =					hidden_dropout_prob
         A__		 =					attention_probs_dropout_prob
         A__		 =					max_position_embeddings
         A__		 =					type_vocab_size
         A__		 =					type_sequence_label_size
         A__		 =					initializer_range
         A__		 =					num_labels
         A__		 =					num_choices
         A__		 =					scope




     def 		snake_case_		(				self):
         A__		 =					ids_tensor([self.batch_size, self.seq_length] ,						self.vocab_size)

         A__		 =					None
         if self.use_input_mask:
             A__		 =					random_attention_mask([self.batch_size, self.seq_length])

         A__		 =					None
         A__		 =					None
         A__		 =					None
         if self.use_labels:
             A__		 =					ids_tensor([self.batch_size] ,						self.type_sequence_label_size)
             A__		 =					ids_tensor([self.batch_size, self.seq_length] ,						self.num_labels)
             A__		 =					ids_tensor([self.batch_size] ,						self.num_choices)

         A__		 =					self.get_config()

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




     def 		snake_case_		(				self):
         A__		 =					EsmConfig(
             vocab_size=3_3 ,						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 		snake_case_		(				self ,						a__ ,						a__ ,						a__ ,						a__ ,						a__ ,						a__):
         A__		 =					EsmForProteinFolding(config=a__).float()
         model.to(a__)
         model.eval()
         A__		 =					model(a__ ,						attention_mask=a__)
         A__		 =					model(a__)
         A__		 =					model(a__)

         self.parent.assertEqual(result.positions.shape ,						(8, self.batch_size, self.seq_length, 1_4, 3))
         self.parent.assertEqual(result.angles.shape ,						(8, self.batch_size, self.seq_length, 7, 2))




     def 		snake_case_		(				self):
         A__		 =					self.prepare_config_and_inputs()
         (
             (
             A__
         ) ,	(
             A__
         ) ,	(
             A__
         ) ,	(
             A__
         ) ,	(
             A__
         ) ,	(
             A__
         ) ,	
         )		 =					config_and_inputs
         A__		 =					{'''input_ids''': input_ids, '''attention_mask''': input_mask}
         return config, inputs_dict







@require_torch
class     _UpperCAmelCase			(       A__			,      A__			,      unittest.TestCase     ):
     UpperCamelCase__		    =				False

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




     def 		snake_case_		(				self):
         A__		 =					EsmFoldModelTester(self)
         A__		 =					ConfigTester(self ,						config_class=a__ ,						hidden_size=3_7)




     def 		snake_case_		(				self):
         self.config_tester.run_common_tests()




     def 		snake_case_		(				self):
         A__		 =					self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_model(*a__)




     @unittest.skip('''Does not support attention outputs''')
     def 		snake_case_		(				self):
         pass




     @unittest.skip
     def 		snake_case_		(				self):
         pass




     @unittest.skip('''Esm does not support embedding resizing''')
     def 		snake_case_		(				self):
         pass




     @unittest.skip('''Esm does not support embedding resizing''')
     def 		snake_case_		(				self):
         pass




     @unittest.skip('''ESMFold does not support passing input embeds!''')
     def 		snake_case_		(				self):
         pass




     @unittest.skip('''ESMFold does not support head pruning.''')
     def 		snake_case_		(				self):
         pass




     @unittest.skip('''ESMFold does not support head pruning.''')
     def 		snake_case_		(				self):
         pass




     @unittest.skip('''ESMFold does not support head pruning.''')
     def 		snake_case_		(				self):
         pass




     @unittest.skip('''ESMFold does not support head pruning.''')
     def 		snake_case_		(				self):
         pass




     @unittest.skip('''ESMFold does not support head pruning.''')
     def 		snake_case_		(				self):
         pass




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




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




     @unittest.skip('''ESMFold only has one output format.''')
     def 		snake_case_		(				self):
         pass




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




     @unittest.skip('''ESMFold does not support input chunking.''')
     def 		snake_case_		(				self):
         pass




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




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




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




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




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




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







@require_torch
class     _UpperCAmelCase			(       A__     ):




     @slow
     def 		snake_case_		(				self):
         A__		 =					EsmForProteinFolding.from_pretrained('''facebook/esmfold_v1''').float()
         model.eval()
         A__		 =					torch.tensor([[0, 6, 4, 1_3, 5, 4, 1_6, 1_2, 1_1, 7, 2]])
         A__		 =					model(a__)['''positions''']
         A__		 =					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] ,						a__ ,						atol=1e-4))





from collections import OrderedDict
from typing import Mapping

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


_lowercase						  =		logging.get_logger(__name__)

_lowercase						  =		{
    "bert-base-uncased": "https://huggingface.co/bert-base-uncased/resolve/main/config.json",
    "bert-large-uncased": "https://huggingface.co/bert-large-uncased/resolve/main/config.json",
    "bert-base-cased": "https://huggingface.co/bert-base-cased/resolve/main/config.json",
    "bert-large-cased": "https://huggingface.co/bert-large-cased/resolve/main/config.json",
    "bert-base-multilingual-uncased": "https://huggingface.co/bert-base-multilingual-uncased/resolve/main/config.json",
    "bert-base-multilingual-cased": "https://huggingface.co/bert-base-multilingual-cased/resolve/main/config.json",
    "bert-base-chinese": "https://huggingface.co/bert-base-chinese/resolve/main/config.json",
    "bert-base-german-cased": "https://huggingface.co/bert-base-german-cased/resolve/main/config.json",
    "bert-large-uncased-whole-word-masking": (
        "https://huggingface.co/bert-large-uncased-whole-word-masking/resolve/main/config.json"
    ),
    "bert-large-cased-whole-word-masking": (
        "https://huggingface.co/bert-large-cased-whole-word-masking/resolve/main/config.json"
    ),
    "bert-large-uncased-whole-word-masking-finetuned-squad": (
        "https://huggingface.co/bert-large-uncased-whole-word-masking-finetuned-squad/resolve/main/config.json"
    ),
    "bert-large-cased-whole-word-masking-finetuned-squad": (
        "https://huggingface.co/bert-large-cased-whole-word-masking-finetuned-squad/resolve/main/config.json"
    ),
    "bert-base-cased-finetuned-mrpc": "https://huggingface.co/bert-base-cased-finetuned-mrpc/resolve/main/config.json",
    "bert-base-german-dbmdz-cased": "https://huggingface.co/bert-base-german-dbmdz-cased/resolve/main/config.json",
    "bert-base-german-dbmdz-uncased": "https://huggingface.co/bert-base-german-dbmdz-uncased/resolve/main/config.json",
    "cl-tohoku/bert-base-japanese": "https://huggingface.co/cl-tohoku/bert-base-japanese/resolve/main/config.json",
    "cl-tohoku/bert-base-japanese-whole-word-masking": (
        "https://huggingface.co/cl-tohoku/bert-base-japanese-whole-word-masking/resolve/main/config.json"
    ),
    "cl-tohoku/bert-base-japanese-char": (
        "https://huggingface.co/cl-tohoku/bert-base-japanese-char/resolve/main/config.json"
    ),
    "cl-tohoku/bert-base-japanese-char-whole-word-masking": (
        "https://huggingface.co/cl-tohoku/bert-base-japanese-char-whole-word-masking/resolve/main/config.json"
    ),
    "TurkuNLP/bert-base-finnish-cased-v1": (
        "https://huggingface.co/TurkuNLP/bert-base-finnish-cased-v1/resolve/main/config.json"
    ),
    "TurkuNLP/bert-base-finnish-uncased-v1": (
        "https://huggingface.co/TurkuNLP/bert-base-finnish-uncased-v1/resolve/main/config.json"
    ),
    "wietsedv/bert-base-dutch-cased": "https://huggingface.co/wietsedv/bert-base-dutch-cased/resolve/main/config.json",
    # See all BERT models at https://huggingface.co/models?filter=bert
}



class     _UpperCAmelCase			(       A__     ):
     UpperCamelCase__		    =				'''bert'''




     def __init__(				self ,						a__=3_0_5_2_2 ,						a__=7_6_8 ,						a__=1_2 ,						a__=1_2 ,						a__=3_0_7_2 ,						a__="gelu" ,						a__=0.1 ,						a__=0.1 ,						a__=5_1_2 ,						a__=2 ,						a__=0.0_2 ,						a__=1e-12 ,						a__=0 ,						a__="absolute" ,						a__=True ,						a__=None ,						**a__ ,						):
         super().__init__(pad_token_id=a__ ,						**a__)

         A__		 =					vocab_size
         A__		 =					hidden_size
         A__		 =					num_hidden_layers
         A__		 =					num_attention_heads
         A__		 =					hidden_act
         A__		 =					intermediate_size
         A__		 =					hidden_dropout_prob
         A__		 =					attention_probs_dropout_prob
         A__		 =					max_position_embeddings
         A__		 =					type_vocab_size
         A__		 =					initializer_range
         A__		 =					layer_norm_eps
         A__		 =					position_embedding_type
         A__		 =					use_cache
         A__		 =					classifier_dropout







class     _UpperCAmelCase			(       A__     ):




     @property
     def 		snake_case_		(				self):
         if self.task == "multiple-choice":
             A__		 =					{0: '''batch''', 1: '''choice''', 2: '''sequence'''}
         else:
             A__		 =					{0: '''batch''', 1: '''sequence'''}
         return OrderedDict(
             [
                 ('''input_ids''', dynamic_axis),
                 ('''attention_mask''', dynamic_axis),
                 ('''token_type_ids''', dynamic_axis),
             ])
