





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

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

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

     A__				=			self.get_config()

     return config, input_ids, input_mask, sequence_labels, token_labels, choice_labels
 def      UpperCAmelCase_							(				self      :Optional[Any]							)->							Tuple:
     A__				=			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=lowercase_						,       esmfold_config={"trunk": {"num_blocks": 2}, "fp16_esm": False}						,       )
     return config
 def      UpperCAmelCase_							(				self      :List[Any]						,       lowercase_      :List[Any]						,       lowercase_      :List[str]						,       lowercase_      :Optional[Any]						,       lowercase_      :int						,       lowercase_      :Dict						,       lowercase_      :Dict							)->							int:
     A__				=			EsmForProteinFolding(config=lowercase_							).float()
     model.to(lowercase_							)
     model.eval()
     A__				=			model(lowercase_						,       attention_mask=lowercase_							)
     A__				=			model(lowercase_							)
     A__				=			model(lowercase_							)

     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      :Tuple							)->							Union[str, Any]:
     A__				=			self.prepare_config_and_inputs()
     (
         (
         A__
     ),	(
         A__
     ),	(
         A__
     ),	(
         A__
     ),	(
         A__
     ),	(
         A__
     ),	
     )				=			config_and_inputs
     A__				=			{"input_ids": input_ids, "attention_mask": input_mask}
     return config, inputs_dict






@require_torch
class    UpperCAmelCase	(       UpperCamelCase__			,       UpperCamelCase__			,       unittest.TestCase				):
 __lowercase   		=				False

 __lowercase   		=				(EsmForProteinFolding,) if is_torch_available() else ()
 __lowercase   		=				()
 __lowercase   		=				{} if is_torch_available() else {}
 __lowercase   		=				False
 def      UpperCAmelCase_							(				self      :Union[str, Any]							)->							Union[str, Any]:
     A__				=			EsmFoldModelTester(self							)
     A__				=			ConfigTester(self						,       config_class=lowercase_						,       hidden_size=37							)
 def      UpperCAmelCase_							(				self      :Tuple							)->							List[str]:
     self.config_tester.run_common_tests()
 def      UpperCAmelCase_							(				self      :Tuple							)->							Optional[int]:
     A__				=			self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_model(*lowercase_							)
 @unittest.skip("Does not support attention outputs"							)
 def      UpperCAmelCase_							(				self      :Dict							)->							List[str]:
     pass
 @unittest.skip
 def      UpperCAmelCase_							(				self      :List[str]							)->							List[str]:
     pass
 @unittest.skip("Esm does not support embedding resizing"							)
 def      UpperCAmelCase_							(				self      :Union[str, Any]							)->							int:
     pass
 @unittest.skip("Esm does not support embedding resizing"							)
 def      UpperCAmelCase_							(				self      :Tuple							)->							Optional[Any]:
     pass
 @unittest.skip("ESMFold does not support passing input embeds!"							)
 def      UpperCAmelCase_							(				self      :Optional[int]							)->							Optional[Any]:
     pass
 @unittest.skip("ESMFold does not support head pruning."							)
 def      UpperCAmelCase_							(				self      :Tuple							)->							List[str]:
     pass
 @unittest.skip("ESMFold does not support head pruning."							)
 def      UpperCAmelCase_							(				self      :str							)->							Optional[int]:
     pass
 @unittest.skip("ESMFold does not support head pruning."							)
 def      UpperCAmelCase_							(				self      :Any							)->							str:
     pass
 @unittest.skip("ESMFold does not support head pruning."							)
 def      UpperCAmelCase_							(				self      :Tuple							)->							List[Any]:
     pass
 @unittest.skip("ESMFold does not support head pruning."							)
 def      UpperCAmelCase_							(				self      :Optional[Any]							)->							List[Any]:
     pass
 @unittest.skip("ESMFold does not output hidden states in the normal way."							)
 def      UpperCAmelCase_							(				self      :Dict							)->							Optional[int]:
     pass
 @unittest.skip("ESMfold does not output hidden states in the normal way."							)
 def      UpperCAmelCase_							(				self      :Tuple							)->							Union[str, Any]:
     pass
 @unittest.skip("ESMFold only has one output format."							)
 def      UpperCAmelCase_							(				self      :Optional[Any]							)->							Any:
     pass
 @unittest.skip("This test doesn't work for ESMFold and doesn't test core functionality"							)
 def      UpperCAmelCase_							(				self      :List[str]							)->							int:
     pass
 @unittest.skip("ESMFold does not support input chunking."							)
 def      UpperCAmelCase_							(				self      :Tuple							)->							List[str]:
     pass
 @unittest.skip("ESMFold doesn't respect you and it certainly doesn't respect your initialization arguments."							)
 def      UpperCAmelCase_							(				self      :List[str]							)->							List[str]:
     pass
 @unittest.skip("ESMFold doesn't support torchscript compilation."							)
 def      UpperCAmelCase_							(				self      :str							)->							Dict:
     pass
 @unittest.skip("ESMFold doesn't support torchscript compilation."							)
 def      UpperCAmelCase_							(				self      :Optional[int]							)->							List[Any]:
     pass
 @unittest.skip("ESMFold doesn't support torchscript compilation."							)
 def      UpperCAmelCase_							(				self      :Union[str, Any]							)->							Tuple:
     pass
 @unittest.skip("ESMFold doesn't support data parallel."							)
 def      UpperCAmelCase_							(				self      :Union[str, Any]							)->							Union[str, Any]:
     pass



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






@require_torch
class    UpperCAmelCase	(       UpperCamelCase__				):
 @slow
 def      UpperCAmelCase_							(				self      :Optional[Any]							)->							Optional[Any]:
     A__				=			EsmForProteinFolding.from_pretrained("facebook/esmfold_v1"							).float()
     model.eval()
     A__				=			torch.tensor([[0, 6, 4, 13, 5, 4, 16, 12, 11, 7, 2]]							)
     A__				=			model(lowercase_							)["positions"]
     A__				=			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]						,       lowercase_						,       atol=1E-4							)							)










'''simple docstring'''

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	(       UpperCamelCase__				):
 def __init__(				self      :Optional[Any]						,       *lowercase_      :int						,       lowercase_      :Any=None						,       lowercase_      :List[str]=None						,       **lowercase_      :Any							)->							Any:
     super().__init__(*lowercase_						,       **lowercase_							)
     A__				=			eval_examples
     A__				=			post_process_function
 def      UpperCAmelCase_							(				self      :str						,       lowercase_      :str=None						,       lowercase_      :Optional[int]=None						,       lowercase_      :Optional[int]=None						,       lowercase_      :str = "eval"							)->							Union[str, Any]:
     A__				=			self.eval_dataset if eval_dataset is None else eval_dataset
     A__				=			self.get_eval_dataloader(lowercase_							)
     A__				=			self.eval_examples if eval_examples is None else eval_examples

     # Temporarily disable metric computation, we will do it in the loop here.
     A__				=			self.compute_metrics
     A__				=			None
     A__				=			self.prediction_loop if self.args.use_legacy_prediction_loop else self.evaluation_loop
     A__				=			time.time()
     try:
         A__				=			eval_loop(
             lowercase_						,       description="Evaluation"						,       prediction_loss_only=True if compute_metrics is None else None						,       ignore_keys=lowercase_						,       metric_key_prefix=lowercase_						,       )
     finally:
         A__				=			compute_metrics
     A__				=			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(
             lowercase_						,       lowercase_						,       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
         A__				=			self.post_process_function(lowercase_						,       lowercase_						,       output.predictions							)
         A__				=			self.compute_metrics(lowercase_							)

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

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

     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()							)

     A__				=			self.callback_handler.on_evaluate(self.args						,       self.state						,       self.control						,       lowercase_							)
     return metrics



 def      UpperCAmelCase_							(				self      :List[str]						,       lowercase_      :List[Any]						,       lowercase_      :str						,       lowercase_      :Any=None						,       lowercase_      :str = "test"							)->							List[Any]:
     A__				=			self.get_test_dataloader(lowercase_							)

     # Temporarily disable metric computation, we will do it in the loop here.
     A__				=			self.compute_metrics
     A__				=			None
     A__				=			self.prediction_loop if self.args.use_legacy_prediction_loop else self.evaluation_loop
     A__				=			time.time()
     try:
         A__				=			eval_loop(
             lowercase_						,       description="Prediction"						,       prediction_loss_only=True if compute_metrics is None else None						,       ignore_keys=lowercase_						,       metric_key_prefix=lowercase_						,       )
     finally:
         A__				=			compute_metrics
     A__				=			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(
             lowercase_						,       lowercase_						,       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

     A__				=			self.post_process_function(lowercase_						,       lowercase_						,       output.predictions						,       "predict"							)
     A__				=			self.compute_metrics(lowercase_							)

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



