'''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 			__lowercase		:







					"""simple docstring"""






					def __init__(		self     ,			_lowerCamelCase     ,			_lowerCamelCase=1_3     ,			_lowerCamelCase=7     ,			_lowerCamelCase=False     ,			_lowerCamelCase=True     ,			_lowerCamelCase=False     ,			_lowerCamelCase=False     ,			_lowerCamelCase=1_9     ,			_lowerCamelCase=3_2     ,			_lowerCamelCase=5     ,			_lowerCamelCase=4     ,			_lowerCamelCase=3_7     ,			_lowerCamelCase="gelu"     ,			_lowerCamelCase=0.1     ,			_lowerCamelCase=0.1     ,			_lowerCamelCase=5_1_2     ,			_lowerCamelCase=1_6     ,			_lowerCamelCase=2     ,			_lowerCamelCase=0.0_2     ,			_lowerCamelCase=3     ,			_lowerCamelCase=4     ,			_lowerCamelCase=None     ,			):
									__UpperCamelCase	:	Union[str, Any]		    =  parent
									__UpperCamelCase	:	int		    =  batch_size
									__UpperCamelCase	:	Any		    =  seq_length
									__UpperCamelCase	:	List[Any]		    =  is_training
									__UpperCamelCase	:	List[Any]		    =  use_input_mask
									__UpperCamelCase	:	Optional[int]		    =  use_token_type_ids
									__UpperCamelCase	:	Union[str, Any]		    =  use_labels
									__UpperCamelCase	:	str		    =  vocab_size
									__UpperCamelCase	:	Optional[int]		    =  hidden_size
									__UpperCamelCase	:	Any		    =  num_hidden_layers
									__UpperCamelCase	:	Dict		    =  num_attention_heads
									__UpperCamelCase	:	Any		    =  intermediate_size
									__UpperCamelCase	:	List[Any]		    =  hidden_act
									__UpperCamelCase	:	Optional[Any]		    =  hidden_dropout_prob
									__UpperCamelCase	:	Any		    =  attention_probs_dropout_prob
									__UpperCamelCase	:	Optional[int]		    =  max_position_embeddings
									__UpperCamelCase	:	Any		    =  type_vocab_size
									__UpperCamelCase	:	Any		    =  type_sequence_label_size
									__UpperCamelCase	:	Tuple		    =  initializer_range
									__UpperCamelCase	:	Any		    =  num_labels
									__UpperCamelCase	:	Any		    =  num_choices
									__UpperCamelCase	:	Optional[Any]		    =  scope







					def lowerCAmelCase     (		self    ):
									__UpperCamelCase	:	List[Any]		    =  ids_tensor([self.batch_size, self.seq_length]     ,			self.vocab_size    )

									__UpperCamelCase	:	Dict		    =  None
									if self.use_input_mask:
													__UpperCamelCase	:	str		    =  random_attention_mask([self.batch_size, self.seq_length]    )

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

									__UpperCamelCase	:	str		    =  self.get_config()

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







					def lowerCAmelCase     (		self    ):
									__UpperCamelCase	:	Dict		    =  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=snake_case_     ,			esmfold_config={'trunk': {'num_blocks': 2}, 'fp16_esm': False}     ,			)
									return config







					def lowerCAmelCase     (		self     ,			_lowerCamelCase     ,			_lowerCamelCase     ,			_lowerCamelCase     ,			_lowerCamelCase     ,			_lowerCamelCase     ,			_lowerCamelCase    ):
									__UpperCamelCase	:	Tuple		    =  EsmForProteinFolding(config=snake_case_    ).float()
									model.to(snake_case_    )
									model.eval()
									__UpperCamelCase	:	Optional[Any]		    =  model(snake_case_     ,			attention_mask=snake_case_    )
									__UpperCamelCase	:	str		    =  model(snake_case_    )
									__UpperCamelCase	:	Optional[int]		    =  model(snake_case_    )

									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 lowerCAmelCase     (		self    ):
									__UpperCamelCase	:	str		    =  self.prepare_config_and_inputs()
									(
									    (
									    __UpperCamelCase						
									)    ,					(
									    __UpperCamelCase						
									)    ,					(
									    __UpperCamelCase						
									)    ,					(
									    __UpperCamelCase						
									)    ,					(
									    __UpperCamelCase						
									)    ,					(
									    __UpperCamelCase						
									)    ,											
									)	:	Optional[int]		    =  config_and_inputs
									__UpperCamelCase	:	Dict		    =  {'input_ids': input_ids, 'attention_mask': input_mask}
									return config, inputs_dict






@require_torch
class 			__lowercase		(						_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    ):
									__UpperCamelCase	:	List[Any]		    =  EsmFoldModelTester(self    )
									__UpperCamelCase	:	Optional[int]		    =  ConfigTester(self     ,			config_class=snake_case_     ,			hidden_size=3_7    )







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







					def lowerCAmelCase     (		self    ):
									__UpperCamelCase	:	int		    =  self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_model(*snake_case_    )







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







					@unittest.skip
					def lowerCAmelCase     (		self    ):
									pass







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







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







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







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







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







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







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







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







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







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







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







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







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







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







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







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







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







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







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






@require_torch
class 			__lowercase		(						_a						):







					"""simple docstring"""






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




'''simple docstring'''
a=  '''0.21.0'''

from .accelerator import Accelerator
from .big_modeling import (
    cpu_offload,
    cpu_offload_with_hook,
    disk_offload,
    dispatch_model,
    init_empty_weights,
    init_on_device,
    load_checkpoint_and_dispatch,
)
from .data_loader import skip_first_batches
from .launchers import debug_launcher, notebook_launcher
from .state import PartialState
from .utils import (
    DeepSpeedPlugin,
    DistributedDataParallelKwargs,
    DistributedType,
    FullyShardedDataParallelPlugin,
    GradScalerKwargs,
    InitProcessGroupKwargs,
    find_executable_batch_size,
    infer_auto_device_map,
    is_rich_available,
    load_checkpoint_in_model,
    synchronize_rng_states,
)


if is_rich_available():
  from .utils import rich



