


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






			"""simple docstring"""






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

						'''simple docstring'''


						UpperCamelCase__       =							parent
						UpperCamelCase__       =							batch_size
						UpperCamelCase__       =							seq_length
						UpperCamelCase__       =							is_training
						UpperCamelCase__       =							use_input_mask
						UpperCamelCase__       =							use_token_type_ids
						UpperCamelCase__       =							use_labels
						UpperCamelCase__       =							vocab_size
						UpperCamelCase__       =							hidden_size
						UpperCamelCase__       =							num_hidden_layers
						UpperCamelCase__       =							num_attention_heads
						UpperCamelCase__       =							intermediate_size
						UpperCamelCase__       =							hidden_act
						UpperCamelCase__       =							hidden_dropout_prob
						UpperCamelCase__       =							attention_probs_dropout_prob
						UpperCamelCase__       =							max_position_embeddings
						UpperCamelCase__       =							type_vocab_size
						UpperCamelCase__       =							type_sequence_label_size
						UpperCamelCase__       =							initializer_range
						UpperCamelCase__       =							num_labels
						UpperCamelCase__       =							num_choices
						UpperCamelCase__       =							scope




			def 					snake_case__					(				self  ):

						'''simple docstring'''


						UpperCamelCase__       =							ids_tensor([self.batch_size, self.seq_length]   ,    self.vocab_size  )

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

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

						UpperCamelCase__       =							self.get_config()

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




			def 					snake_case__					(				self  ):

						'''simple docstring'''


						UpperCamelCase__       =							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 					snake_case__					(				self   ,    snake_case   ,    snake_case   ,    snake_case   ,    snake_case   ,    snake_case   ,    snake_case  ):

						'''simple docstring'''


						UpperCamelCase__       =							EsmForProteinFolding(config=snake_case  ).float()
						model.to(snake_case  )
						model.eval()
						UpperCamelCase__       =							model(snake_case   ,    attention_mask=snake_case  )
						UpperCamelCase__       =							model(snake_case  )
						UpperCamelCase__       =							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 					snake_case__					(				self  ):

						'''simple docstring'''


						UpperCamelCase__       =							self.prepare_config_and_inputs()
						(
						    (
						    UpperCamelCase__
						),   (
						    UpperCamelCase__
						),   (
						    UpperCamelCase__
						),   (
						    UpperCamelCase__
						),   (
						    UpperCamelCase__
						),   (
						    UpperCamelCase__
						),   
						)       =							config_and_inputs
						UpperCamelCase__       =							{"input_ids": input_ids, "attention_mask": input_mask}
						return config, inputs_dict







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






			"""simple docstring"""






			_UpperCamelCase		:       Union[str, Any] 			=		False

			_UpperCamelCase		:       Union[str, Any] 			=		(EsmForProteinFolding,) if is_torch_available() else ()
			_UpperCamelCase		:       Dict 			=		()
			_UpperCamelCase		:       Union[str, Any] 			=		{} if is_torch_available() else {}
			_UpperCamelCase		:       Optional[int] 			=		False




			def 					snake_case__					(				self  ):

						'''simple docstring'''


						UpperCamelCase__       =							EsmFoldModelTester(self  )
						UpperCamelCase__       =							ConfigTester(self   ,    config_class=snake_case   ,    hidden_size=37  )




			def 					snake_case__					(				self  ):

						'''simple docstring'''


						self.config_tester.run_common_tests()




			def 					snake_case__					(				self  ):

						'''simple docstring'''


						UpperCamelCase__       =							self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_model(*snake_case  )




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

						'''simple docstring'''


						pass




			@unittest.skip
			def 					snake_case__					(				self  ):

						'''simple docstring'''


						pass




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

						'''simple docstring'''


						pass




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

						'''simple docstring'''


						pass




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

						'''simple docstring'''


						pass




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

						'''simple docstring'''


						pass




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

						'''simple docstring'''


						pass




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

						'''simple docstring'''


						pass




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

						'''simple docstring'''


						pass




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

						'''simple docstring'''


						pass




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

						'''simple docstring'''


						pass




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

						'''simple docstring'''


						pass




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

						'''simple docstring'''


						pass




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

						'''simple docstring'''


						pass




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

						'''simple docstring'''


						pass




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

						'''simple docstring'''


						pass




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

						'''simple docstring'''


						pass




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

						'''simple docstring'''


						pass




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

						'''simple docstring'''


						pass




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

						'''simple docstring'''


						pass







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

						'''simple docstring'''


						pass







@require_torch
class 							lowerCamelCase__						(				UpperCAmelCase   ):






			"""simple docstring"""






			@slow
			def 					snake_case__					(				self  ):

						'''simple docstring'''


						UpperCamelCase__       =							EsmForProteinFolding.from_pretrained("facebook/esmfold_v1"  ).float()
						model.eval()
						UpperCamelCase__       =							torch.tensor([[0, 6, 4, 13, 5, 4, 16, 12, 11, 7, 2]]  )
						UpperCamelCase__       =							model(snake_case  )["positions"]
						UpperCamelCase__       =							torch.tensor([2.5828, 0.7993, -10.9334]   ,    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


__UpperCamelCase     		=							logging.get_logger(__name__)

__UpperCamelCase     		=							{
    'facebook/xmod-base': 'https://huggingface.co/facebook/xmod-base/resolve/main/config.json',
    'facebook/xmod-large-prenorm': 'https://huggingface.co/facebook/xmod-large-prenorm/resolve/main/config.json',
    'facebook/xmod-base-13-125k': 'https://huggingface.co/facebook/xmod-base-13-125k/resolve/main/config.json',
    'facebook/xmod-base-30-125k': 'https://huggingface.co/facebook/xmod-base-30-125k/resolve/main/config.json',
    'facebook/xmod-base-30-195k': 'https://huggingface.co/facebook/xmod-base-30-195k/resolve/main/config.json',
    'facebook/xmod-base-60-125k': 'https://huggingface.co/facebook/xmod-base-60-125k/resolve/main/config.json',
    'facebook/xmod-base-60-265k': 'https://huggingface.co/facebook/xmod-base-60-265k/resolve/main/config.json',
    'facebook/xmod-base-75-125k': 'https://huggingface.co/facebook/xmod-base-75-125k/resolve/main/config.json',
    'facebook/xmod-base-75-269k': 'https://huggingface.co/facebook/xmod-base-75-269k/resolve/main/config.json',
}







class 							lowerCamelCase__						(				UpperCAmelCase   ):






			"""simple docstring"""






			_UpperCamelCase		:       List[str] 			=		'xmod'




			def __init__(				self   ,    snake_case=30522   ,    snake_case=768   ,    snake_case=12   ,    snake_case=12   ,    snake_case=3072   ,    snake_case="gelu"   ,    snake_case=0.1   ,    snake_case=0.1   ,    snake_case=512   ,    snake_case=2   ,    snake_case=0.02   ,    snake_case=1E-1_2   ,    snake_case=1   ,    snake_case=0   ,    snake_case=2   ,    snake_case="absolute"   ,    snake_case=True   ,    snake_case=None   ,    snake_case=False   ,    snake_case=2   ,    snake_case=False   ,    snake_case=True   ,    snake_case=True   ,    snake_case=("en_XX",)   ,    snake_case=None   ,    **snake_case   ,    ):

						'''simple docstring'''


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

						UpperCamelCase__       =							vocab_size
						UpperCamelCase__       =							hidden_size
						UpperCamelCase__       =							num_hidden_layers
						UpperCamelCase__       =							num_attention_heads
						UpperCamelCase__       =							hidden_act
						UpperCamelCase__       =							intermediate_size
						UpperCamelCase__       =							hidden_dropout_prob
						UpperCamelCase__       =							attention_probs_dropout_prob
						UpperCamelCase__       =							max_position_embeddings
						UpperCamelCase__       =							type_vocab_size
						UpperCamelCase__       =							initializer_range
						UpperCamelCase__       =							layer_norm_eps
						UpperCamelCase__       =							position_embedding_type
						UpperCamelCase__       =							use_cache
						UpperCamelCase__       =							classifier_dropout
						UpperCamelCase__       =							pre_norm
						UpperCamelCase__       =							adapter_reduction_factor
						UpperCamelCase__       =							adapter_layer_norm
						UpperCamelCase__       =							adapter_reuse_layer_norm
						UpperCamelCase__       =							ln_before_adapter
						UpperCamelCase__       =							list(snake_case  )
						UpperCamelCase__       =							default_language







class 							lowerCamelCase__						(				UpperCAmelCase   ):






			"""simple docstring"""






			@property
			def 					snake_case__					(				self  ):

						'''simple docstring'''


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


