


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 			__lowerCAmelCase   :






  """simple docstring"""





  def __init__(   self				:				List[str]			, _lowerCAmelCase				:				str			, )      ->       Optional[int]:




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



  def     lowerCAmelCase__     (   self				:				List[Any]     )      ->       Any:




     """simple docstring"""
     snake_case_    			=    ids_tensor([self.batch_size, self.seq_length]			, self.vocab_size     )

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

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

     snake_case_    			=    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     lowerCAmelCase__     (   self				:				Optional[int]     )      ->       List[str]:




     """simple docstring"""
     (
         (
         snake_case_    
     )			,			(
         snake_case_    
     )			,			(
         snake_case_    
     )			,			(
         snake_case_    
     )			,			(
         snake_case_    
     )			,			(
         snake_case_    
     )			,			    
     )    			=    self.prepare_config_and_inputs()

     snake_case_    			=    True
     snake_case_    			=    floats_tensor([self.batch_size, self.seq_length, self.hidden_size]     )
     snake_case_    			=    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     lowerCAmelCase__     (   self				:				Union[str, Any]			, _lowerCAmelCase				:				Optional[int]			, _lowerCAmelCase				:				List[Any]			, _lowerCAmelCase				:				Tuple			, _lowerCAmelCase				:				int			, _lowerCAmelCase				:				Tuple			, _lowerCAmelCase				:				str     )      ->       Any:




     """simple docstring"""
     snake_case_    			=    TFEsmModel(config=_lowerCAmelCase     )
     snake_case_    			=    {"input_ids": input_ids, "attention_mask": input_mask}
     snake_case_    			=    model(_lowerCAmelCase     )

     snake_case_    			=    [input_ids, input_mask]
     snake_case_    			=    model(_lowerCAmelCase     )

     snake_case_    			=    model(_lowerCAmelCase     )

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



  def     lowerCAmelCase__     (   self				:				Dict			, _lowerCAmelCase				:				List[str]			, _lowerCAmelCase				:				Union[str, Any]			, _lowerCAmelCase				:				Dict			, _lowerCAmelCase				:				List[Any]			, _lowerCAmelCase				:				int			, _lowerCAmelCase				:				List[str]			, _lowerCAmelCase				:				List[Any]			, _lowerCAmelCase				:				Tuple			, )      ->       Tuple:




     """simple docstring"""
     snake_case_    			=    True

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

     snake_case_    			=    [input_ids, input_mask]
     snake_case_    			=    model(_lowerCAmelCase			, encoder_hidden_states=_lowerCAmelCase     )

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

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



  def     lowerCAmelCase__     (   self				:				List[str]			, _lowerCAmelCase				:				int			, _lowerCAmelCase				:				Tuple			, _lowerCAmelCase				:				str			, _lowerCAmelCase				:				Optional[int]			, _lowerCAmelCase				:				Tuple			, _lowerCAmelCase				:				List[str]     )      ->       Optional[int]:




     """simple docstring"""
     snake_case_    			=    TFEsmForMaskedLM(config=_lowerCAmelCase     )
     snake_case_    			=    model([input_ids, input_mask]     )
     self.parent.assertEqual(result.logits.shape			, (self.batch_size, self.seq_length, self.vocab_size)     )



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




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



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




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






@require_tf
class 			__lowerCAmelCase   (       snake_case_       ,     snake_case_       ,     unittest.TestCase    ):






  """simple docstring"""





  _SCREAMING_SNAKE_CASE									=				(
      (
          TFEsmModel,
          TFEsmForMaskedLM,
          TFEsmForSequenceClassification,
          TFEsmForTokenClassification,
      )
      if is_tf_available()
      else ()
  )
  _SCREAMING_SNAKE_CASE									=				(
      {
          'feature-extraction': TFEsmModel,
          'fill-mask': TFEsmForMaskedLM,
          'text-classification': TFEsmForSequenceClassification,
          'token-classification': TFEsmForTokenClassification,
          'zero-shot': TFEsmForSequenceClassification,
      }
      if is_tf_available()
      else {}
  )
  _SCREAMING_SNAKE_CASE									=				False
  _SCREAMING_SNAKE_CASE									=				False



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




     """simple docstring"""
     snake_case_    			=    TFEsmModelTester(self     )
     snake_case_    			=    ConfigTester(self			, config_class=_lowerCAmelCase			, hidden_size=3_7     )



  def     lowerCAmelCase__     (   self				:				int     )      ->       Any:




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



  def     lowerCAmelCase__     (   self				:				str     )      ->       str:




     """simple docstring"""
     snake_case_    			=    self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_model(*_lowerCAmelCase     )



  def     lowerCAmelCase__     (   self				:				Optional[int]     )      ->       Tuple:




     """simple docstring"""
     snake_case_    			=    self.model_tester.prepare_config_and_inputs_for_decoder()
     self.model_tester.create_and_check_model_as_decoder(*_lowerCAmelCase     )



  def     lowerCAmelCase__     (   self				:				int     )      ->       Dict:




     """simple docstring"""
     snake_case_    			=    self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_masked_lm(*_lowerCAmelCase     )



  def     lowerCAmelCase__     (   self				:				Dict     )      ->       List[str]:




     """simple docstring"""
     snake_case_    			=    self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_token_classification(*_lowerCAmelCase     )



  @slow
  def     lowerCAmelCase__     (   self				:				Tuple     )      ->       Dict:




     """simple docstring"""
     for model_name in TF_ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
        snake_case_    			=    TFEsmModel.from_pretrained(_lowerCAmelCase     )
        self.assertIsNotNone(_lowerCAmelCase     )



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




     """simple docstring"""
     pass



  @unittest.skip("Protein models do not support embedding resizing."     )
  def     lowerCAmelCase__     (   self				:				int     )      ->       Any:




     """simple docstring"""
     pass



  def     lowerCAmelCase__     (   self				:				Optional[int]     )      ->       Optional[Any]:




     """simple docstring"""
     snake_case_			,			snake_case_    			=    self.model_tester.prepare_config_and_inputs_for_common()

     for model_class in self.all_model_classes:
        snake_case_    			=    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
           snake_case_    			=    model.get_bias()
           assert isinstance(_lowerCAmelCase			, _lowerCAmelCase     )
           for k, v in name.items():
              assert isinstance(_lowerCAmelCase			, tf.Variable     )
        else:
           snake_case_    			=    model.get_output_embeddings()
           assert x is None
           snake_case_    			=    model.get_bias()
           assert name is None






@require_tf
class 			__lowerCAmelCase   (       unittest.TestCase    ):






  """simple docstring"""





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




     """simple docstring"""
     snake_case_    			=    TFEsmForMaskedLM.from_pretrained("facebook/esm2_t6_8M_UR50D"     )

     snake_case_    			=    tf.constant([[0, 1, 2, 3, 4, 5]]     )
     snake_case_    			=    model(_lowerCAmelCase     )[0]
     snake_case_    			=    [1, 6, 3_3]
     self.assertEqual(list(output.numpy().shape     )			, _lowerCAmelCase     )
     # compare the actual values for a slice.
     snake_case_    			=    tf.constant(
         [
             [
                 [8.921_518, -1_0.5_8_9_8_1_4, -6.4_671_307],
                 [-6.3_967_156, -1_3.9_1_1_3_7_7, -1.1_211_915],
                 [-7.781_247, -1_3.9_5_1_5_5_7, -3.740_592],
             ]
         ]     )
     self.assertTrue(numpy.allclose(output[:, :3, :3].numpy()			, expected_slice.numpy()			, atol=1e-2     )     )



  @slow
  def     lowerCAmelCase__     (   self				:				Union[str, Any]     )      ->       Optional[Any]:




     """simple docstring"""
     snake_case_    			=    TFEsmModel.from_pretrained("facebook/esm2_t6_8M_UR50D"     )

     snake_case_    			=    tf.constant([[0, 6, 4, 1_3, 5, 4, 1_6, 1_2, 1_1, 7, 2]]     )
     snake_case_    			=    model(_lowerCAmelCase     )[0]
     # compare the actual values for a slice.
     snake_case_    			=    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 typing import Dict, List, Optional

from torch.utils.data import Dataset

from transformers import SeqaSeqTrainer, 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     _SCREAMING_SNAKE_CASE     (      snake_case_						):



	def __init__(    self						,  *lowercase						,  lowercase=None						,  lowercase=None						,  **lowercase )				->       List[str]:
					super().__init__(*lowercase						,  **lowercase )
					lowerCamelCase_						=			eval_examples
					lowerCamelCase_						=			post_process_function



	def     SCREAMING_SNAKE_CASE_(    self						,  lowercase = None						,  lowercase=None						,  lowercase = None						,  lowercase = "eval"						,  **lowercase						,  )				->       Dict[str, float]:
					lowerCamelCase_						=			gen_kwargs.copy()
					lowerCamelCase_						=			(
					    gen_kwargs["max_length"] if gen_kwargs.get("max_length" ) is not None else self.args.generation_max_length
					)
					lowerCamelCase_						=			(
					    gen_kwargs["num_beams"] if gen_kwargs.get("num_beams" ) is not None else self.args.generation_num_beams
					)
					lowerCamelCase_						=			gen_kwargs

					lowerCamelCase_						=			self.eval_dataset if eval_dataset is None else eval_dataset
					lowerCamelCase_						=			self.get_eval_dataloader(lowercase )
					lowerCamelCase_						=			self.eval_examples if eval_examples is None else eval_examples

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

									metrics.update(output.metrics )
					else:
									lowerCamelCase_						=			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() )

					lowerCamelCase_						=			self.callback_handler.on_evaluate(self.args						,  self.state						,  self.control						,  lowercase )
					return metrics



	def     SCREAMING_SNAKE_CASE_(    self						,  lowercase						,  lowercase						,  lowercase=None						,  lowercase = "test"						,  **lowercase )				->       Union[str, Any]:
					lowerCamelCase_						=			gen_kwargs.copy()

					lowerCamelCase_						=			self.get_test_dataloader(lowercase )

					# Temporarily disable metric computation, we will do it in the loop here.
					lowerCamelCase_						=			self.compute_metrics
					lowerCamelCase_						=			None
					lowerCamelCase_						=			time.time()
					lowerCamelCase_						=			self.prediction_loop if self.args.use_legacy_prediction_loop else self.evaluation_loop
					try:
									lowerCamelCase_						=			eval_loop(
									    lowercase						,  description="Prediction"						,  prediction_loss_only=True if compute_metrics is None else None						,  ignore_keys=lowercase						,  metric_key_prefix=lowercase						,  )
					finally:
									lowerCamelCase_						=			compute_metrics

					lowerCamelCase_						=			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

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


