






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

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




				def 				lowercase__   (     self						:    int  ):

								'''simple docstring'''
								lowercase__    =			ids_tensor([self.batch_size, self.seq_length],  self.vocab_size  )

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

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

								lowercase__    =			self.get_config()

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




				def 				lowercase__   (     self						:    Optional[int]  ):

								'''simple docstring'''
								lowercase__    =			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 				lowercase__   (     self						:    Any,  lowerCamelCase						:    Union[str, Any],  lowerCamelCase						:    List[Any],  lowerCamelCase						:    Tuple,  lowerCamelCase						:    Optional[Any],  lowerCamelCase						:    Union[str, Any],  lowerCamelCase						:    Optional[int]  ):

								'''simple docstring'''
								lowercase__    =			EsmForProteinFolding(config=lowerCamelCase  ).float()
								model.to(lowerCamelCase  )
								model.eval()
								lowercase__    =			model(lowerCamelCase,  attention_mask=lowerCamelCase  )
								lowercase__    =			model(lowerCamelCase  )
								lowercase__    =			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 				lowercase__   (     self						:    List[Any]  ):

								'''simple docstring'''
								lowercase__    =			self.prepare_config_and_inputs()
								(
								    (
								    lowercase__
								) ,						(
								    lowercase__
								) ,						(
								    lowercase__
								) ,						(
								    lowercase__
								) ,						(
								    lowercase__
								) ,						(
								    lowercase__
								) ,						
								)    =			config_and_inputs
								lowercase__    =			{'''input_ids''': input_ids, '''attention_mask''': input_mask}
								return config, inputs_dict



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

				"""simple docstring"""




				lowercase__											=       False

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




				def 				lowercase__   (     self						:    Tuple  ):

								'''simple docstring'''
								lowercase__    =			EsmFoldModelTester(self  )
								lowercase__    =			ConfigTester(self,  config_class=lowerCamelCase,  hidden_size=37  )




				def 				lowercase__   (     self						:    List[str]  ):

								'''simple docstring'''
								self.config_tester.run_common_tests()




				def 				lowercase__   (     self						:    List[Any]  ):

								'''simple docstring'''
								lowercase__    =			self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_model(*lowerCamelCase  )




				@unittest.skip('''Does not support attention outputs'''  )
				def 				lowercase__   (     self						:    Optional[Any]  ):

								'''simple docstring'''
								pass




				@unittest.skip
				def 				lowercase__   (     self						:    Optional[int]  ):

								'''simple docstring'''
								pass




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

								'''simple docstring'''
								pass




				@unittest.skip('''Esm does not support embedding resizing'''  )
				def 				lowercase__   (     self						:    List[Any]  ):

								'''simple docstring'''
								pass




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

								'''simple docstring'''
								pass




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

								'''simple docstring'''
								pass




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

								'''simple docstring'''
								pass




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

								'''simple docstring'''
								pass




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

								'''simple docstring'''
								pass




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

								'''simple docstring'''
								pass




				@unittest.skip('''ESMFold does not output hidden states in the normal way.'''  )
				def 				lowercase__   (     self						:    Optional[Any]  ):

								'''simple docstring'''
								pass




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

								'''simple docstring'''
								pass




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

								'''simple docstring'''
								pass




				@unittest.skip('''This test doesn\'t work for ESMFold and doesn\'t test core functionality'''  )
				def 				lowercase__   (     self						:    Optional[Any]  ):

								'''simple docstring'''
								pass




				@unittest.skip('''ESMFold does not support input chunking.'''  )
				def 				lowercase__   (     self						:    Optional[int]  ):

								'''simple docstring'''
								pass




				@unittest.skip('''ESMFold doesn\'t respect you and it certainly doesn\'t respect your initialization arguments.'''  )
				def 				lowercase__   (     self						:    Union[str, Any]  ):

								'''simple docstring'''
								pass




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

								'''simple docstring'''
								pass




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

								'''simple docstring'''
								pass




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

								'''simple docstring'''
								pass




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

								'''simple docstring'''
								pass




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

								'''simple docstring'''
								pass



@require_torch
class       _UpperCAmelCase					(       A__	):

				"""simple docstring"""




				@slow
				def 				lowercase__   (     self						:    Optional[int]  ):

								'''simple docstring'''
								lowercase__    =			EsmForProteinFolding.from_pretrained('''facebook/esmfold_v1'''  ).float()
								model.eval()
								lowercase__    =			torch.tensor([[0, 6, 4, 13, 5, 4, 16, 12, 11, 7, 2]]  )
								lowercase__    =			model(lowerCamelCase  )['''positions''']
								lowercase__    =			torch.tensor([2.5828, 0.7993, -10.9334],  dtype=torch.floataa  )
								self.assertTrue(torch.allclose(position_outputs[0, 0, 0, 0],  lowerCamelCase,  atol=1E-4  )  )












import math
import time

from transformers import Trainer, is_torch_tpu_available
from transformers.trainer_utils import PredictionOutput, speed_metrics


if is_torch_tpu_available(check_device=False):
							import torch_xla.core.xla_model as xm
							import torch_xla.debug.metrics as met
class       _UpperCAmelCase					(       A__	):

				"""simple docstring"""




				def __init__(     self						:    Dict,  *lowerCamelCase						:    Union[str, Any],  lowerCamelCase						:    Union[str, Any]=None,  lowerCamelCase						:    Any=None,  **lowerCamelCase						:    str  ):

								'''simple docstring'''
								super().__init__(*lowerCamelCase,  **lowerCamelCase  )
								lowercase__    =			eval_examples
								lowercase__    =			post_process_function




				def 				lowercase__   (     self						:    int,  lowerCamelCase						:    str=None,  lowerCamelCase						:    Optional[Any]=None,  lowerCamelCase						:    Union[str, Any]=None,  lowerCamelCase						:    str = "eval"  ):

								'''simple docstring'''
								lowercase__    =			self.eval_dataset if eval_dataset is None else eval_dataset
								lowercase__    =			self.get_eval_dataloader(lowerCamelCase  )
								lowercase__    =			self.eval_examples if eval_examples is None else eval_examples

								# Temporarily disable metric computation, we will do it in the loop here.
								lowercase__    =			self.compute_metrics
								lowercase__    =			None
								lowercase__    =			self.prediction_loop if self.args.use_legacy_prediction_loop else self.evaluation_loop
								lowercase__    =			time.time()
								try:
												lowercase__    =			eval_loop(
												    lowerCamelCase,  description='''Evaluation''',  prediction_loss_only=True if compute_metrics is None else None,  ignore_keys=lowerCamelCase,  metric_key_prefix=lowerCamelCase,  )
								finally:
												lowercase__    =			compute_metrics
								lowercase__    =			self.args.eval_batch_size * self.args.world_size
								if F"""{metric_key_prefix}_jit_compilation_time""" in output.metrics:
												start_time += output.metrics[F"""{metric_key_prefix}_jit_compilation_time"""]
								output.metrics.update(
								    speed_metrics(
								        lowerCamelCase,  lowerCamelCase,  num_samples=output.num_samples,  num_steps=math.ceil(output.num_samples / total_batch_size  ),  )  )
								if self.post_process_function is not None and self.compute_metrics is not None and self.args.should_save:
												# Only the main node write the results by default
												lowercase__    =			self.post_process_function(lowerCamelCase,  lowerCamelCase,  output.predictions  )
												lowercase__    =			self.compute_metrics(lowerCamelCase  )

												# Prefix all keys with metric_key_prefix + '_'
												for key in list(metrics.keys()  ):
																if not key.startswith(F"""{metric_key_prefix}_"""  ):
																				lowercase__    =			metrics.pop(lowerCamelCase  )
												metrics.update(output.metrics  )
								else:
												lowercase__    =			output.metrics

								if self.args.should_log:
												# Only the main node log the results by default
												self.log(lowerCamelCase  )

								if self.args.tpu_metrics_debug or self.args.debug:
												# tpu-comment: Logging debug metrics for PyTorch/XLA (compile, execute times, ops, etc.)
												xm.master_print(met.metrics_report()  )

								lowercase__    =			self.callback_handler.on_evaluate(self.args,  self.state,  self.control,  lowerCamelCase  )
								return metrics




				def 				lowercase__   (     self						:    List[Any],  lowerCamelCase						:    Any,  lowerCamelCase						:    Dict,  lowerCamelCase						:    int=None,  lowerCamelCase						:    str = "test"  ):

								'''simple docstring'''
								lowercase__    =			self.get_test_dataloader(lowerCamelCase  )

								# Temporarily disable metric computation, we will do it in the loop here.
								lowercase__    =			self.compute_metrics
								lowercase__    =			None
								lowercase__    =			self.prediction_loop if self.args.use_legacy_prediction_loop else self.evaluation_loop
								lowercase__    =			time.time()
								try:
												lowercase__    =			eval_loop(
												    lowerCamelCase,  description='''Prediction''',  prediction_loss_only=True if compute_metrics is None else None,  ignore_keys=lowerCamelCase,  metric_key_prefix=lowerCamelCase,  )
								finally:
												lowercase__    =			compute_metrics
								lowercase__    =			self.args.eval_batch_size * self.args.world_size
								if F"""{metric_key_prefix}_jit_compilation_time""" in output.metrics:
												start_time += output.metrics[F"""{metric_key_prefix}_jit_compilation_time"""]
								output.metrics.update(
								    speed_metrics(
								        lowerCamelCase,  lowerCamelCase,  num_samples=output.num_samples,  num_steps=math.ceil(output.num_samples / total_batch_size  ),  )  )

								if self.post_process_function is None or self.compute_metrics is None:
												return output

								lowercase__    =			self.post_process_function(lowerCamelCase,  lowerCamelCase,  output.predictions,  '''predict'''  )
								lowercase__    =			self.compute_metrics(lowerCamelCase  )

								# Prefix all keys with metric_key_prefix + '_'
								for key in list(metrics.keys()  ):
												if not key.startswith(F"""{metric_key_prefix}_"""  ):
																lowercase__    =			metrics.pop(lowerCamelCase  )
								metrics.update(output.metrics  )
								return PredictionOutput(predictions=predictions.predictions,  label_ids=predictions.label_ids,  metrics=lowerCamelCase  )




