


from __future__ import annotations

import unittest

from transformers import EsmConfig, is_tf_available
from transformers.testing_utils import require_tf, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
 import numpy
 import tensorflow as tf

 from transformers.models.esm.modeling_tf_esm import (
     TF_ESM_PRETRAINED_MODEL_ARCHIVE_LIST,
     TFEsmForMaskedLM,
     TFEsmForSequenceClassification,
     TFEsmForTokenClassification,
     TFEsmModel,
 )






class 	__lowercase					:





      def __init__(		self    :       Any       ,							__lowerCamelCase    :       Dict       ,							)     -> int:






         """simple docstring"""







         UpperCAmelCase      					=  parent
         UpperCAmelCase      					=  1_3
         UpperCAmelCase      					=  7
         UpperCAmelCase      					=  True
         UpperCAmelCase      					=  True
         UpperCAmelCase      					=  True
         UpperCAmelCase      					=  9_9
         UpperCAmelCase      					=  3_2
         UpperCAmelCase      					=  2
         UpperCAmelCase      					=  4
         UpperCAmelCase      					=  3_7
         UpperCAmelCase      					=  """gelu"""
         UpperCAmelCase      					=  0.1
         UpperCAmelCase      					=  0.1
         UpperCAmelCase      					=  5_1_2
         UpperCAmelCase      					=  1_6
         UpperCAmelCase      					=  2
         UpperCAmelCase      					=  0.02
         UpperCAmelCase      					=  3
         UpperCAmelCase      					=  4
         UpperCAmelCase      					=  None





      def 							_lowercase					(		self    :       Dict							)     -> Optional[int]:






         """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      					=  EsmConfig(
             vocab_size=self.vocab_size       ,							hidden_size=self.hidden_size       ,							num_hidden_layers=self.num_hidden_layers       ,							pad_token_id=1       ,							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       ,							)

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





      def 							_lowercase					(		self    :       Any							)     -> List[str]:






         """simple docstring"""







         (
             (
             UpperCAmelCase
         )					,	(
             UpperCAmelCase
         )					,	(
             UpperCAmelCase
         )					,	(
             UpperCAmelCase
         )					,	(
             UpperCAmelCase
         )					,	(
             UpperCAmelCase
         )					,	
         )      					=  self.prepare_config_and_inputs()

         UpperCAmelCase      					=  True
         UpperCAmelCase      					=  floats_tensor([self.batch_size, self.seq_length, self.hidden_size]							)
         UpperCAmelCase      					=  ids_tensor([self.batch_size, self.seq_length]       ,							vocab_size=2							)

         return (
             config,
             input_ids,
             input_mask,
             sequence_labels,
             token_labels,
             choice_labels,
             encoder_hidden_states,
             encoder_attention_mask,
         )





      def 							_lowercase					(		self    :       Optional[int]       ,							__lowerCamelCase    :       Optional[int]       ,							__lowerCamelCase    :       int       ,							__lowerCamelCase    :       List[Any]       ,							__lowerCamelCase    :       int       ,							__lowerCamelCase    :       Optional[int]       ,							__lowerCamelCase    :       List[str]							)     -> Optional[int]:






         """simple docstring"""







         UpperCAmelCase      					=  TFEsmModel(config=__lowerCamelCase							)
         UpperCAmelCase      					=  {"""input_ids""": input_ids, """attention_mask""": input_mask}
         UpperCAmelCase      					=  model(__lowerCamelCase							)

         UpperCAmelCase      					=  [input_ids, input_mask]
         UpperCAmelCase      					=  model(__lowerCamelCase							)

         UpperCAmelCase      					=  model(__lowerCamelCase							)

         self.parent.assertEqual(result.last_hidden_state.shape       ,							(self.batch_size, self.seq_length, self.hidden_size)							)





      def 							_lowercase					(		self    :       List[str]       ,							__lowerCamelCase    :       Union[str, Any]       ,							__lowerCamelCase    :       Dict       ,							__lowerCamelCase    :       str       ,							__lowerCamelCase    :       Any       ,							__lowerCamelCase    :       Optional[Any]       ,							__lowerCamelCase    :       Optional[int]       ,							__lowerCamelCase    :       str       ,							__lowerCamelCase    :       Dict       ,							)     -> Dict:






         """simple docstring"""







         UpperCAmelCase      					=  True

         UpperCAmelCase      					=  TFEsmModel(config=__lowerCamelCase							)
         UpperCAmelCase      					=  {
             """input_ids""": input_ids,
             """attention_mask""": input_mask,
             """encoder_hidden_states""": encoder_hidden_states,
             """encoder_attention_mask""": encoder_attention_mask,
         }
         UpperCAmelCase      					=  model(__lowerCamelCase							)

         UpperCAmelCase      					=  [input_ids, input_mask]
         UpperCAmelCase      					=  model(__lowerCamelCase       ,							encoder_hidden_states=__lowerCamelCase							)

         # Also check the case where encoder outputs are not passed
         UpperCAmelCase      					=  model(__lowerCamelCase       ,							attention_mask=__lowerCamelCase							)

         self.parent.assertEqual(result.last_hidden_state.shape       ,							(self.batch_size, self.seq_length, self.hidden_size)							)





      def 							_lowercase					(		self    :       Dict       ,							__lowerCamelCase    :       Optional[int]       ,							__lowerCamelCase    :       int       ,							__lowerCamelCase    :       Dict       ,							__lowerCamelCase    :       Dict       ,							__lowerCamelCase    :       List[str]       ,							__lowerCamelCase    :       int							)     -> int:






         """simple docstring"""







         UpperCAmelCase      					=  TFEsmForMaskedLM(config=__lowerCamelCase							)
         UpperCAmelCase      					=  model([input_ids, input_mask]							)
         self.parent.assertEqual(result.logits.shape       ,							(self.batch_size, self.seq_length, self.vocab_size)							)





      def 							_lowercase					(		self    :       str       ,							__lowerCamelCase    :       Union[str, Any]       ,							__lowerCamelCase    :       List[str]       ,							__lowerCamelCase    :       Union[str, Any]       ,							__lowerCamelCase    :       Union[str, Any]       ,							__lowerCamelCase    :       List[Any]       ,							__lowerCamelCase    :       Optional[Any]							)     -> Tuple:






         """simple docstring"""







         UpperCAmelCase      					=  self.num_labels
         UpperCAmelCase      					=  TFEsmForTokenClassification(config=__lowerCamelCase							)
         UpperCAmelCase      					=  {"""input_ids""": input_ids, """attention_mask""": input_mask}
         UpperCAmelCase      					=  model(__lowerCamelCase							)
         self.parent.assertEqual(result.logits.shape       ,							(self.batch_size, self.seq_length, self.num_labels)							)





      def 							_lowercase					(		self    :       str							)     -> Union[str, Any]:






         """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_tf
class 	__lowercase					(	__snake_case			,       __snake_case			,       unittest.TestCase	):
      UpperCamelCase										=							(
          (
              TFEsmModel,
              TFEsmForMaskedLM,
              TFEsmForSequenceClassification,
              TFEsmForTokenClassification,
          )
          if is_tf_available()
          else ()
      )
      UpperCamelCase										=							(
          {
              '''feature-extraction''': TFEsmModel,
              '''fill-mask''': TFEsmForMaskedLM,
              '''text-classification''': TFEsmForSequenceClassification,
              '''token-classification''': TFEsmForTokenClassification,
              '''zero-shot''': TFEsmForSequenceClassification,
          }
          if is_tf_available()
          else {}
      )
      UpperCamelCase										=							False
      UpperCamelCase										=							False





      def 							_lowercase					(		self    :       Tuple							)     -> Any:






         """simple docstring"""







         UpperCAmelCase      					=  TFEsmModelTester(self							)
         UpperCAmelCase      					=  ConfigTester(self       ,							config_class=__lowerCamelCase       ,							hidden_size=3_7							)





      def 							_lowercase					(		self    :       Dict							)     -> Optional[int]:






         """simple docstring"""







         self.config_tester.run_common_tests()





      def 							_lowercase					(		self    :       Tuple							)     -> Optional[Any]:






         """simple docstring"""







         UpperCAmelCase      					=  self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_model(*__lowerCamelCase							)





      def 							_lowercase					(		self    :       List[Any]							)     -> Optional[Any]:






         """simple docstring"""







         UpperCAmelCase      					=  self.model_tester.prepare_config_and_inputs_for_decoder()
         self.model_tester.create_and_check_model_as_decoder(*__lowerCamelCase							)





      def 							_lowercase					(		self    :       Optional[Any]							)     -> Tuple:






         """simple docstring"""







         UpperCAmelCase      					=  self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_masked_lm(*__lowerCamelCase							)





      def 							_lowercase					(		self    :       List[Any]							)     -> List[str]:






         """simple docstring"""







         UpperCAmelCase      					=  self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_token_classification(*__lowerCamelCase							)





      @slow
      def 							_lowercase					(		self    :       Optional[Any]							)     -> Optional[Any]:






         """simple docstring"""







         for model_name in TF_ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
            UpperCAmelCase      					=  TFEsmModel.from_pretrained(__lowerCamelCase							)
            self.assertIsNotNone(__lowerCamelCase							)





      @unittest.skip("""Protein models do not support embedding resizing."""							)
      def 							_lowercase					(		self    :       List[Any]							)     -> Optional[int]:






         """simple docstring"""







         pass





      @unittest.skip("""Protein models do not support embedding resizing."""							)
      def 							_lowercase					(		self    :       Optional[int]							)     -> str:






         """simple docstring"""







         pass





      def 							_lowercase					(		self    :       Optional[int]							)     -> Union[str, Any]:






         """simple docstring"""







         UpperCAmelCase					,	UpperCAmelCase      					=  self.model_tester.prepare_config_and_inputs_for_common()

         for model_class in self.all_model_classes:
            UpperCAmelCase      					=  model_class(__lowerCamelCase							)
            assert isinstance(model.get_input_embeddings()       ,							tf.keras.layers.Layer							)
            if model_class is TFEsmForMaskedLM:
               # Output embedding test differs from the main test because they're a matrix, not a layer
               UpperCAmelCase      					=  model.get_bias()
               assert isinstance(__lowerCamelCase       ,							__lowerCamelCase							)
               for k, v in name.items():
                  assert isinstance(__lowerCamelCase       ,							tf.Variable							)
            else:
               UpperCAmelCase      					=  model.get_output_embeddings()
               assert x is None
               UpperCAmelCase      					=  model.get_bias()
               assert name is None






@require_tf
class 	__lowercase					(	unittest.TestCase	):





      @slow
      def 							_lowercase					(		self    :       Any							)     -> int:






         """simple docstring"""







         UpperCAmelCase      					=  TFEsmForMaskedLM.from_pretrained("""facebook/esm2_t6_8M_UR50D"""							)

         UpperCAmelCase      					=  tf.constant([[0, 1, 2, 3, 4, 5]]							)
         UpperCAmelCase      					=  model(__lowerCamelCase							)[0]
         UpperCAmelCase      					=  [1, 6, 3_3]
         self.assertEqual(list(output.numpy().shape							)       ,							__lowerCamelCase							)
         # compare the actual values for a slice.
         UpperCAmelCase      					=  tf.constant(
             [
                 [
                     [8.921_518, -10.589_814, -6.4_671_307],
                     [-6.3_967_156, -13.911_377, -1.1_211_915],
                     [-7.781_247, -13.951_557, -3.740_592],
                 ]
             ]							)
         self.assertTrue(numpy.allclose(output[:, :3, :3].numpy()       ,							expected_slice.numpy()       ,							atol=1e-2							)							)





      @slow
      def 							_lowercase					(		self    :       str							)     -> List[str]:






         """simple docstring"""







         UpperCAmelCase      					=  TFEsmModel.from_pretrained("""facebook/esm2_t6_8M_UR50D"""							)

         UpperCAmelCase      					=  tf.constant([[0, 6, 4, 1_3, 5, 4, 1_6, 1_2, 1_1, 7, 2]]							)
         UpperCAmelCase      					=  model(__lowerCamelCase							)[0]
         # compare the actual values for a slice.
         UpperCAmelCase      					=  tf.constant(
             [
                 [
                     [0.14_443_092, 0.54_125_327, 0.3_247_739],
                     [0.30_340_484, 0.00_526_676, 0.31_077_722],
                     [0.32_278_043, -0.24_987_096, 0.3_414_628],
                 ]
             ]							)
         self.assertTrue(numpy.allclose(output[:, :3, :3].numpy()       ,							expected_slice.numpy()       ,							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 	__lowercase					(	__snake_case	):





      def __init__(		self    :       Dict       ,							*__lowerCamelCase    :       Dict       ,							__lowerCamelCase    :       Union[str, Any]=None       ,							__lowerCamelCase    :       str=None       ,							**__lowerCamelCase    :       Optional[int]							)     -> Optional[int]:






         """simple docstring"""







         super().__init__(*__lowerCamelCase       ,							**__lowerCamelCase							)
         UpperCAmelCase      					=  eval_examples
         UpperCAmelCase      					=  post_process_function





      def 							_lowercase					(		self    :       Any       ,							__lowerCamelCase    :       int=None       ,							__lowerCamelCase    :       int=None       ,							__lowerCamelCase    :       Tuple=None       ,							__lowerCamelCase    :       str = "eval"							)     -> List[str]:






         """simple docstring"""







         UpperCAmelCase      					=  self.eval_dataset if eval_dataset is None else eval_dataset
         UpperCAmelCase      					=  self.get_eval_dataloader(__lowerCamelCase							)
         UpperCAmelCase      					=  self.eval_examples if eval_examples is None else eval_examples

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

         UpperCAmelCase      					=  self.callback_handler.on_evaluate(self.args       ,							self.state       ,							self.control       ,							__lowerCamelCase							)
         return metrics





      def 							_lowercase					(		self    :       Optional[Any]       ,							__lowerCamelCase    :       Union[str, Any]       ,							__lowerCamelCase    :       Any       ,							__lowerCamelCase    :       Dict=None       ,							__lowerCamelCase    :       str = "test"							)     -> Dict:






         """simple docstring"""







         UpperCAmelCase      					=  self.get_test_dataloader(__lowerCamelCase							)

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

         UpperCAmelCase      					=  self.post_process_function(__lowerCamelCase       ,							__lowerCamelCase       ,							output.predictions       ,							"""predict"""							)
         UpperCAmelCase      					=  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}_"""							):
               UpperCAmelCase      					=  metrics.pop(__lowerCamelCase							)
         metrics.update(output.metrics							)
         return PredictionOutput(predictions=predictions.predictions       ,							label_ids=predictions.label_ids       ,							metrics=__lowerCamelCase							)



