





'''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				:			str      ,						lowerCamelCase_				:			Union[str, Any]      ,						lowerCamelCase_				:			Union[str, Any]=13      ,						lowerCamelCase_				:			Optional[int]=7      ,						lowerCamelCase_				:			int=False      ,						lowerCamelCase_				:			Optional[Any]=True      ,						lowerCamelCase_				:			Optional[int]=False      ,						lowerCamelCase_				:			Optional[int]=False      ,						lowerCamelCase_				:			List[Any]=19      ,						lowerCamelCase_				:			int=32      ,						lowerCamelCase_				:			List[Any]=5      ,						lowerCamelCase_				:			Tuple=4      ,						lowerCamelCase_				:			Optional[int]=37      ,						lowerCamelCase_				:			str="gelu"      ,						lowerCamelCase_				:			Optional[Any]=0.1      ,						lowerCamelCase_				:			Optional[Any]=0.1      ,						lowerCamelCase_				:			int=5_12      ,						lowerCamelCase_				:			Union[str, Any]=16      ,						lowerCamelCase_				:			List[str]=2      ,						lowerCamelCase_				:			Tuple=0.02      ,						lowerCamelCase_				:			Dict=3      ,						lowerCamelCase_				:			Optional[int]=4      ,						lowerCamelCase_				:			Tuple=None      ,						):

							'''simple docstring'''



							SCREAMING_SNAKE_CASE :				Optional[Any]		=			parent
							SCREAMING_SNAKE_CASE :				List[str]		=			batch_size
							SCREAMING_SNAKE_CASE :				List[Any]		=			seq_length
							SCREAMING_SNAKE_CASE :				Tuple		=			is_training
							SCREAMING_SNAKE_CASE :				Tuple		=			use_input_mask
							SCREAMING_SNAKE_CASE :				List[str]		=			use_token_type_ids
							SCREAMING_SNAKE_CASE :				Optional[int]		=			use_labels
							SCREAMING_SNAKE_CASE :				Union[str, Any]		=			vocab_size
							SCREAMING_SNAKE_CASE :				Dict		=			hidden_size
							SCREAMING_SNAKE_CASE :				Any		=			num_hidden_layers
							SCREAMING_SNAKE_CASE :				Union[str, Any]		=			num_attention_heads
							SCREAMING_SNAKE_CASE :				Any		=			intermediate_size
							SCREAMING_SNAKE_CASE :				Any		=			hidden_act
							SCREAMING_SNAKE_CASE :				Optional[Any]		=			hidden_dropout_prob
							SCREAMING_SNAKE_CASE :				Union[str, Any]		=			attention_probs_dropout_prob
							SCREAMING_SNAKE_CASE :				Union[str, Any]		=			max_position_embeddings
							SCREAMING_SNAKE_CASE :				int		=			type_vocab_size
							SCREAMING_SNAKE_CASE :				List[str]		=			type_sequence_label_size
							SCREAMING_SNAKE_CASE :				Optional[int]		=			initializer_range
							SCREAMING_SNAKE_CASE :				Dict		=			num_labels
							SCREAMING_SNAKE_CASE :				List[str]		=			num_choices
							SCREAMING_SNAKE_CASE :				List[Any]		=			scope


					def lowerCamelCase_			(						self				:			List[str]						):

							'''simple docstring'''



							SCREAMING_SNAKE_CASE :				Dict		=			ids_tensor([self.batch_size, self.seq_length]      ,						self.vocab_size						)

							SCREAMING_SNAKE_CASE :				Union[str, Any]		=			None
							if self.use_input_mask:
									SCREAMING_SNAKE_CASE :				Tuple		=			random_attention_mask([self.batch_size, self.seq_length]						)

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

							SCREAMING_SNAKE_CASE :				Union[str, Any]		=			self.get_config()

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


					def lowerCamelCase_			(						self				:			str						):

							'''simple docstring'''



							SCREAMING_SNAKE_CASE :				Union[str, Any]		=			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				:			List[Any]      ,						lowerCamelCase_				:			Tuple      ,						lowerCamelCase_				:			List[str]      ,						lowerCamelCase_				:			Dict      ,						lowerCamelCase_				:			int      ,						lowerCamelCase_				:			Optional[int]      ,						lowerCamelCase_				:			int						):

							'''simple docstring'''



							SCREAMING_SNAKE_CASE :				Union[str, Any]		=			EsmForProteinFolding(config=lowerCamelCase_						).float()
							model.to(lowerCamelCase_						)
							model.eval()
							SCREAMING_SNAKE_CASE :				Tuple		=			model(lowerCamelCase_      ,						attention_mask=lowerCamelCase_						)
							SCREAMING_SNAKE_CASE :				int		=			model(lowerCamelCase_						)
							SCREAMING_SNAKE_CASE :				str		=			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				:			Optional[int]						):

							'''simple docstring'''



							SCREAMING_SNAKE_CASE :				Dict		=			self.prepare_config_and_inputs()
							(
							    (
							    SCREAMING_SNAKE_CASE
							),							(
							    SCREAMING_SNAKE_CASE
							),							(
							    SCREAMING_SNAKE_CASE
							),							(
							    SCREAMING_SNAKE_CASE
							),							(
							    SCREAMING_SNAKE_CASE
							),							(
							    SCREAMING_SNAKE_CASE
							),							
							) :				Optional[Any]		=			config_and_inputs
							SCREAMING_SNAKE_CASE :				Union[str, Any]		=			{"""input_ids""": input_ids, """attention_mask""": input_mask}
							return config, inputs_dict







@require_torch
class      UpperCamelCase__					(					lowercase_						,					lowercase_						,					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				:			Union[str, Any]						):

							'''simple docstring'''



							SCREAMING_SNAKE_CASE :				int		=			EsmFoldModelTester(self						)
							SCREAMING_SNAKE_CASE :				str		=			ConfigTester(self      ,						config_class=lowerCamelCase_      ,						hidden_size=37						)


					def lowerCamelCase_			(						self				:			Tuple						):

							'''simple docstring'''



							self.config_tester.run_common_tests()


					def lowerCamelCase_			(						self				:			List[Any]						):

							'''simple docstring'''



							SCREAMING_SNAKE_CASE :				Any		=			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				:			Union[str, Any]						):

							'''simple docstring'''



							pass


					@unittest.skip
					def lowerCamelCase_			(						self				:			Tuple						):

							'''simple docstring'''



							pass


					@unittest.skip("""Esm does not support embedding resizing"""						)
					def lowerCamelCase_			(						self				:			Optional[int]						):

							'''simple docstring'''



							pass


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

							'''simple docstring'''



							pass


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

							'''simple docstring'''



							pass


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

							'''simple docstring'''



							pass


					@unittest.skip("""ESMFold does not support head pruning."""						)
					def lowerCamelCase_			(						self				:			Optional[int]						):

							'''simple docstring'''



							pass


					@unittest.skip("""ESMFold does not support head pruning."""						)
					def lowerCamelCase_			(						self				:			List[str]						):

							'''simple docstring'''



							pass


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

							'''simple docstring'''



							pass


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

							'''simple docstring'''



							pass


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

							'''simple docstring'''



							pass


					@unittest.skip("""ESMfold does not output hidden states in the normal way."""						)
					def lowerCamelCase_			(						self				:			Optional[int]						):

							'''simple docstring'''



							pass


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

							'''simple docstring'''



							pass


					@unittest.skip("""This test doesn't work for ESMFold and doesn't test core functionality"""						)
					def lowerCamelCase_			(						self				:			Union[str, Any]						):

							'''simple docstring'''



							pass


					@unittest.skip("""ESMFold does not support input chunking."""						)
					def lowerCamelCase_			(						self				:			Union[str, Any]						):

							'''simple docstring'''



							pass


					@unittest.skip("""ESMFold doesn't respect you and it certainly doesn't respect your initialization arguments."""						)
					def lowerCamelCase_			(						self				:			Optional[int]						):

							'''simple docstring'''



							pass


					@unittest.skip("""ESMFold doesn't support torchscript compilation."""						)
					def lowerCamelCase_			(						self				:			List[Any]						):

							'''simple docstring'''



							pass


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

							'''simple docstring'''



							pass


					@unittest.skip("""ESMFold doesn't support torchscript compilation."""						)
					def lowerCamelCase_			(						self				:			Optional[Any]						):

							'''simple docstring'''



							pass


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

							'''simple docstring'''



							pass



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

							'''simple docstring'''



							pass







@require_torch
class      UpperCamelCase__					(					lowercase_			):




					"""simple docstring"""


					@slow
					def lowerCamelCase_			(						self				:			List[Any]						):

							'''simple docstring'''



							SCREAMING_SNAKE_CASE :				Tuple		=			EsmForProteinFolding.from_pretrained("""facebook/esmfold_v1"""						).float()
							model.eval()
							SCREAMING_SNAKE_CASE :				int		=			torch.tensor([[0, 6, 4, 13, 5, 4, 16, 12, 11, 7, 2]]						)
							SCREAMING_SNAKE_CASE :				Union[str, Any]		=			model(lowerCamelCase_						)["""positions"""]
							SCREAMING_SNAKE_CASE :				int		=			torch.tensor([2.5_828, 0.7_993, -10.9_334]      ,						dtype=torch.floataa						)
							self.assertTrue(torch.allclose(position_outputs[0, 0, 0, 0]      ,						lowerCamelCase_      ,						atol=1e-4						)						)










'''simple docstring'''




import unittest
from typing import Tuple

import torch

from diffusers.utils import floats_tensor, randn_tensor, torch_all_close, torch_device
from diffusers.utils.testing_utils import require_torch




@require_torch
class      UpperCamelCase__					:




					"""simple docstring"""


					@property
					def lowerCamelCase_			(						self				:			Union[str, Any]						):

							'''simple docstring'''



							return self.get_dummy_input()


					@property
					def lowerCamelCase_			(						self				:			List[Any]						):

							'''simple docstring'''



							if self.block_type == "down":
									return (4, 32, 16, 16)
							elif self.block_type == "mid":
									return (4, 32, 32, 32)
							elif self.block_type == "up":
									return (4, 32, 64, 64)

							raise ValueError(f'''\'{self.block_type}\' is not a supported block_type. Set it to \'up\', \'mid\', or \'down\'.'''						)


					def lowerCamelCase_			(						self				:			Union[str, Any]      ,						lowerCamelCase_				:			Dict=True      ,						lowerCamelCase_				:			List[Any]=False      ,						lowerCamelCase_				:			str=False      ,						lowerCamelCase_				:			Optional[int]=False      ,						):

							'''simple docstring'''



							SCREAMING_SNAKE_CASE :				Dict		=			4
							SCREAMING_SNAKE_CASE :				Any		=			32
							SCREAMING_SNAKE_CASE :				Union[str, Any]		=			(32, 32)

							SCREAMING_SNAKE_CASE :				int		=			torch.manual_seed(0						)
							SCREAMING_SNAKE_CASE :				Optional[Any]		=			torch.device(lowerCamelCase_						)
							SCREAMING_SNAKE_CASE :				Tuple		=			(batch_size, num_channels) + sizes
							SCREAMING_SNAKE_CASE :				Union[str, Any]		=			randn_tensor(lowerCamelCase_      ,						generator=lowerCamelCase_      ,						device=lowerCamelCase_						)
							SCREAMING_SNAKE_CASE :				Union[str, Any]		=			{"""hidden_states""": hidden_states}

							if include_temb:
									SCREAMING_SNAKE_CASE :				Union[str, Any]		=			1_28
									SCREAMING_SNAKE_CASE :				Optional[Any]		=			randn_tensor((batch_size, temb_channels)      ,						generator=lowerCamelCase_      ,						device=lowerCamelCase_						)

							if include_res_hidden_states_tuple:
									SCREAMING_SNAKE_CASE :				str		=			torch.manual_seed(1						)
									SCREAMING_SNAKE_CASE :				int		=			(randn_tensor(lowerCamelCase_      ,						generator=lowerCamelCase_      ,						device=lowerCamelCase_						),)

							if include_encoder_hidden_states:
									SCREAMING_SNAKE_CASE :				Optional[int]		=			floats_tensor((batch_size, 32, 32)						).to(lowerCamelCase_						)

							if include_skip_sample:
									SCREAMING_SNAKE_CASE :				str		=			randn_tensor(((batch_size, 3) + sizes)      ,						generator=lowerCamelCase_      ,						device=lowerCamelCase_						)

							return dummy_input


					def lowerCamelCase_			(						self				:			Optional[Any]						):

							'''simple docstring'''



							SCREAMING_SNAKE_CASE :				Dict		=			{
							    """in_channels""": 32,
							    """out_channels""": 32,
							    """temb_channels""": 1_28,
							}
							if self.block_type == "up":
									SCREAMING_SNAKE_CASE :				List[str]		=			32

							if self.block_type == "mid":
									init_dict.pop("""out_channels"""						)

							SCREAMING_SNAKE_CASE :				Tuple		=			self.dummy_input
							return init_dict, inputs_dict


					def lowerCamelCase_			(						self				:			Any      ,						lowerCamelCase_				:			Any						):

							'''simple docstring'''



							SCREAMING_SNAKE_CASE,							SCREAMING_SNAKE_CASE :				Optional[Any]		=			self.prepare_init_args_and_inputs_for_common()
							SCREAMING_SNAKE_CASE :				Tuple		=			self.block_class(**lowerCamelCase_						)
							unet_block.to(lowerCamelCase_						)
							unet_block.eval()

							with torch.no_grad():
									SCREAMING_SNAKE_CASE :				List[str]		=			unet_block(**lowerCamelCase_						)

							if isinstance(lowerCamelCase_      ,						lowerCamelCase_						):
									SCREAMING_SNAKE_CASE :				Dict		=			output[0]

							self.assertEqual(output.shape      ,						self.output_shape						)

							SCREAMING_SNAKE_CASE :				Dict		=			output[0, -1, -3:, -3:]
							SCREAMING_SNAKE_CASE :				List[Any]		=			torch.tensor(lowerCamelCase_						).to(lowerCamelCase_						)
							assert torch_all_close(output_slice.flatten()      ,						lowerCamelCase_      ,						atol=5e-3						)



					@unittest.skipIf(torch_device == """mps"""      ,						"""Training is not supported in mps"""						)
					def lowerCamelCase_			(						self				:			int						):

							'''simple docstring'''



							SCREAMING_SNAKE_CASE,							SCREAMING_SNAKE_CASE :				Union[str, Any]		=			self.prepare_init_args_and_inputs_for_common()
							SCREAMING_SNAKE_CASE :				Union[str, Any]		=			self.block_class(**lowerCamelCase_						)
							model.to(lowerCamelCase_						)
							model.train()
							SCREAMING_SNAKE_CASE :				Tuple		=			model(**lowerCamelCase_						)

							if isinstance(lowerCamelCase_      ,						lowerCamelCase_						):
									SCREAMING_SNAKE_CASE :				Optional[Any]		=			output[0]

							SCREAMING_SNAKE_CASE :				Union[str, Any]		=			torch.device(lowerCamelCase_						)
							SCREAMING_SNAKE_CASE :				Optional[int]		=			randn_tensor(output.shape      ,						device=lowerCamelCase_						)
							SCREAMING_SNAKE_CASE :				str		=			torch.nn.functional.mse_loss(lowerCamelCase_      ,						lowerCamelCase_						)
							loss.backward()



