



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



	def __init__(    self				:     List[Any]					,						snake_case_				:     str					,						snake_case_				:     Union[str, Any]=13					,						snake_case_				:     Optional[Any]=7					,						snake_case_				:     str=False					,						snake_case_				:     str=True					,						snake_case_				:     Optional[Any]=False					,						snake_case_				:     List[Any]=False					,						snake_case_				:     Union[str, Any]=19					,						snake_case_				:     Optional[int]=32					,						snake_case_				:     Any=5					,						snake_case_				:     Dict=4					,						snake_case_				:     int=37					,						snake_case_				:     Optional[Any]="gelu"					,						snake_case_				:     List[str]=0.1					,						snake_case_				:     int=0.1					,						snake_case_				:     str=5_12					,						snake_case_				:     Optional[Any]=16					,						snake_case_				:     Dict=2					,						snake_case_				:     str=0.0_2					,						snake_case_				:     str=3					,						snake_case_				:     Dict=4					,						snake_case_				:     Tuple=None					,						)->				Optional[int]:
						__lowerCAmelCase						 =parent
						__lowerCAmelCase						 =batch_size
						__lowerCAmelCase						 =seq_length
						__lowerCAmelCase						 =is_training
						__lowerCAmelCase						 =use_input_mask
						__lowerCAmelCase						 =use_token_type_ids
						__lowerCAmelCase						 =use_labels
						__lowerCAmelCase						 =vocab_size
						__lowerCAmelCase						 =hidden_size
						__lowerCAmelCase						 =num_hidden_layers
						__lowerCAmelCase						 =num_attention_heads
						__lowerCAmelCase						 =intermediate_size
						__lowerCAmelCase						 =hidden_act
						__lowerCAmelCase						 =hidden_dropout_prob
						__lowerCAmelCase						 =attention_probs_dropout_prob
						__lowerCAmelCase						 =max_position_embeddings
						__lowerCAmelCase						 =type_vocab_size
						__lowerCAmelCase						 =type_sequence_label_size
						__lowerCAmelCase						 =initializer_range
						__lowerCAmelCase						 =num_labels
						__lowerCAmelCase						 =num_choices
						__lowerCAmelCase						 =scope



	def 		UpperCamelCase				(    self				:     Union[str, Any])->				List[Any]:
						__lowerCAmelCase						 =ids_tensor([self.batch_size, self.seq_length]					,						self.vocab_size)

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

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

						__lowerCAmelCase						 =self.get_config()

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



	def 		UpperCamelCase				(    self				:     Tuple)->				Tuple:
						__lowerCAmelCase						 =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=snake_case_					,						esmfold_config={"""trunk""": {"""num_blocks""": 2}, """fp16_esm""": False}					,						)
						return config



	def 		UpperCamelCase				(    self				:     str					,						snake_case_				:     Tuple					,						snake_case_				:     List[str]					,						snake_case_				:     Optional[int]					,						snake_case_				:     Optional[int]					,						snake_case_				:     Union[str, Any]					,						snake_case_				:     Any)->				int:
						__lowerCAmelCase						 =EsmForProteinFolding(config=snake_case_).float()
						model.to(snake_case_)
						model.eval()
						__lowerCAmelCase						 =model(snake_case_					,						attention_mask=snake_case_)
						__lowerCAmelCase						 =model(snake_case_)
						__lowerCAmelCase						 =model(snake_case_)

						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 		UpperCamelCase				(    self				:     int)->				Any:
						__lowerCAmelCase						 =self.prepare_config_and_inputs()
						(
						    (
						    __lowerCAmelCase
						) ,						(
						    __lowerCAmelCase
						) ,						(
						    __lowerCAmelCase
						) ,						(
						    __lowerCAmelCase
						) ,						(
						    __lowerCAmelCase
						) ,						(
						    __lowerCAmelCase
						) ,						
						)						 =config_and_inputs
						__lowerCAmelCase						 ={"""input_ids""": input_ids, """attention_mask""": input_mask}
						return config, inputs_dict






@require_torch
class     __a			(  SCREAMING_SNAKE_CASE						,     SCREAMING_SNAKE_CASE						,     unittest.TestCase					):
	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 		UpperCamelCase				(    self				:     List[str])->				List[Any]:
						__lowerCAmelCase						 =EsmFoldModelTester(self)
						__lowerCAmelCase						 =ConfigTester(self					,						config_class=snake_case_					,						hidden_size=37)



	def 		UpperCamelCase				(    self				:     Union[str, Any])->				Optional[int]:
						self.config_tester.run_common_tests()



	def 		UpperCamelCase				(    self				:     Union[str, Any])->				Tuple:
						__lowerCAmelCase						 =self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_model(*snake_case_)



	@unittest.skip("""Does not support attention outputs""")
	def 		UpperCamelCase				(    self				:     Union[str, Any])->				Dict:
						pass



	@unittest.skip
	def 		UpperCamelCase				(    self				:     Tuple)->				Any:
						pass



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



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



	@unittest.skip("""ESMFold does not support passing input embeds!""")
	def 		UpperCamelCase				(    self				:     Any)->				Tuple:
						pass



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



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



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



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



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



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



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



	@unittest.skip("""ESMFold only has one output format.""")
	def 		UpperCamelCase				(    self				:     str)->				Optional[Any]:
						pass



	@unittest.skip("""This test doesn't work for ESMFold and doesn't test core functionality""")
	def 		UpperCamelCase				(    self				:     Dict)->				Optional[int]:
						pass



	@unittest.skip("""ESMFold does not support input chunking.""")
	def 		UpperCamelCase				(    self				:     int)->				Optional[int]:
						pass



	@unittest.skip("""ESMFold doesn't respect you and it certainly doesn't respect your initialization arguments.""")
	def 		UpperCamelCase				(    self				:     Optional[Any])->				List[Any]:
						pass



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



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



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



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



	@unittest.skip("""Will be fixed soon by reducing the size of the model used for common tests.""")
	def 		UpperCamelCase				(    self				:     Union[str, Any])->				Union[str, Any]:
						pass






@require_torch
class     __a			(  SCREAMING_SNAKE_CASE					):



	@slow
	def 		UpperCamelCase				(    self				:     str)->				Optional[int]:
						__lowerCAmelCase						 =EsmForProteinFolding.from_pretrained("""facebook/esmfold_v1""").float()
						model.eval()
						__lowerCAmelCase						 =torch.tensor([[0, 6, 4, 13, 5, 4, 16, 12, 11, 7, 2]])
						__lowerCAmelCase						 =model(snake_case_)["""positions"""]
						__lowerCAmelCase						 =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]					,						snake_case_					,						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_													=      {
    '''distilbert-base-uncased''': '''https://huggingface.co/distilbert-base-uncased/resolve/main/config.json''',
    '''distilbert-base-uncased-distilled-squad''': (
        '''https://huggingface.co/distilbert-base-uncased-distilled-squad/resolve/main/config.json'''
    ),
    '''distilbert-base-cased''': '''https://huggingface.co/distilbert-base-cased/resolve/main/config.json''',
    '''distilbert-base-cased-distilled-squad''': (
        '''https://huggingface.co/distilbert-base-cased-distilled-squad/resolve/main/config.json'''
    ),
    '''distilbert-base-german-cased''': '''https://huggingface.co/distilbert-base-german-cased/resolve/main/config.json''',
    '''distilbert-base-multilingual-cased''': (
        '''https://huggingface.co/distilbert-base-multilingual-cased/resolve/main/config.json'''
    ),
    '''distilbert-base-uncased-finetuned-sst-2-english''': (
        '''https://huggingface.co/distilbert-base-uncased-finetuned-sst-2-english/resolve/main/config.json'''
    ),
}






class     __a			(  SCREAMING_SNAKE_CASE					):
	SCREAMING_SNAKE_CASE  	=							"distilbert"
	SCREAMING_SNAKE_CASE  	=							{
	    "hidden_size": "dim",
	    "num_attention_heads": "n_heads",
	    "num_hidden_layers": "n_layers",
	}



	def __init__(    self				:     Optional[Any]					,						snake_case_				:     Tuple=3_05_22					,						snake_case_				:     Any=5_12					,						snake_case_				:     Dict=False					,						snake_case_				:     Tuple=6					,						snake_case_				:     Union[str, Any]=12					,						snake_case_				:     List[str]=7_68					,						snake_case_				:     Optional[Any]=4 * 7_68					,						snake_case_				:     Optional[Any]=0.1					,						snake_case_				:     Dict=0.1					,						snake_case_				:     Union[str, Any]="gelu"					,						snake_case_				:     Any=0.0_2					,						snake_case_				:     Optional[Any]=0.1					,						snake_case_				:     List[Any]=0.2					,						snake_case_				:     Optional[Any]=0					,						**snake_case_				:     List[Any]					,						)->				Tuple:
						__lowerCAmelCase						 =vocab_size
						__lowerCAmelCase						 =max_position_embeddings
						__lowerCAmelCase						 =sinusoidal_pos_embds
						__lowerCAmelCase						 =n_layers
						__lowerCAmelCase						 =n_heads
						__lowerCAmelCase						 =dim
						__lowerCAmelCase						 =hidden_dim
						__lowerCAmelCase						 =dropout
						__lowerCAmelCase						 =attention_dropout
						__lowerCAmelCase						 =activation
						__lowerCAmelCase						 =initializer_range
						__lowerCAmelCase						 =qa_dropout
						__lowerCAmelCase						 =seq_classif_dropout
						super().__init__(**snake_case_					,						pad_token_id=snake_case_)






class     __a			(  SCREAMING_SNAKE_CASE					):



	@property
	def 		UpperCamelCase				(    self				:     Union[str, Any])->				Mapping[str, Mapping[int, str]]:
						if self.task == "multiple-choice":
											__lowerCAmelCase						 ={0: """batch""", 1: """choice""", 2: """sequence"""}
						else:
											__lowerCAmelCase						 ={0: """batch""", 1: """sequence"""}
						return OrderedDict(
						    [
						        ("""input_ids""", dynamic_axis),
						        ("""attention_mask""", dynamic_axis),
						    ])




