





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	:







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



        '''simple docstring'''




        a__     : int  =parent
        a__     : Optional[int]  =1_3
        a__     : Tuple  =7
        a__     : List[str]  =True
        a__     : Tuple  =True
        a__     : Optional[Any]  =True
        a__     : Union[str, Any]  =9_9
        a__     : Any  =3_2
        a__     : Dict  =2
        a__     : Dict  =4
        a__     : List[str]  =3_7
        a__     : str  ='gelu'
        a__     : Dict  =0.1
        a__     : Optional[int]  =0.1
        a__     : Optional[Any]  =5_1_2
        a__     : List[Any]  =1_6
        a__     : Dict  =2
        a__     : Optional[int]  =0.02
        a__     : List[Any]  =3
        a__     : Dict  =4
        a__     : str  =None







      def        _lowercase  (   self							)       ->						int:



        '''simple docstring'''




        a__     : int  =ids_tensor([self.batch_size, self.seq_length]		,				self.vocab_size							)

        a__     : Optional[Any]  =None
        if self.use_input_mask:
          a__     : List[str]  =random_attention_mask([self.batch_size, self.seq_length]							)

        a__     : Dict  =None
        a__     : List[str]  =None
        a__     : Any  =None
        if self.use_labels:
          a__     : Dict  =ids_tensor([self.batch_size]		,				self.type_sequence_label_size							)
          a__     : int  =ids_tensor([self.batch_size, self.seq_length]		,				self.num_labels							)
          a__     : int  =ids_tensor([self.batch_size]		,				self.num_choices							)

        a__     : str  =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							)       ->						List[Any]:



        '''simple docstring'''




        (
            a__							
        )     : Tuple  =self.prepare_config_and_inputs()

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



        '''simple docstring'''




        a__     : Optional[int]  =TFEsmModel(config=lowerCAmelCase__							)
        a__     : Optional[Any]  ={'input_ids': input_ids, 'attention_mask': input_mask}
        a__     : str  =model(lowerCAmelCase__							)

        a__     : Optional[Any]  =[input_ids, input_mask]
        a__     : int  =model(lowerCAmelCase__							)

        a__     : Dict  =model(lowerCAmelCase__							)

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







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



        '''simple docstring'''




        a__     : Tuple  =True

        a__     : Union[str, Any]  =TFEsmModel(config=lowerCAmelCase__							)
        a__     : Tuple  ={
            'input_ids': input_ids,
            'attention_mask': input_mask,
            'encoder_hidden_states': encoder_hidden_states,
            'encoder_attention_mask': encoder_attention_mask,
        }
        a__     : Tuple  =model(lowerCAmelCase__							)

        a__     : Tuple  =[input_ids, input_mask]
        a__     : List[Any]  =model(lowerCAmelCase__		,				encoder_hidden_states=lowerCAmelCase__							)

        # Also check the case where encoder outputs are not passed
        a__     : Dict  =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		,				lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__							)       ->						Optional[int]:



        '''simple docstring'''




        a__     : Any  =TFEsmForMaskedLM(config=lowerCAmelCase__							)
        a__     : List[Any]  =model([input_ids, input_mask]							)
        self.parent.assertEqual(result.logits.shape		,				(self.batch_size, self.seq_length, self.vocab_size)							)







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



        '''simple docstring'''




        a__     : Tuple  =self.num_labels
        a__     : int  =TFEsmForTokenClassification(config=lowerCAmelCase__							)
        a__     : Dict  ={'input_ids': input_ids, 'attention_mask': input_mask}
        a__     : int  =model(lowerCAmelCase__							)
        self.parent.assertEqual(result.logits.shape		,				(self.batch_size, self.seq_length, self.num_labels)							)







      def        _lowercase  (   self							)       ->						Any:



        '''simple docstring'''




        a__     : Optional[int]  =self.prepare_config_and_inputs()
        (
            a__							
        )     : Optional[int]  =config_and_inputs
        a__     : List[Any]  ={'input_ids': input_ids, 'attention_mask': input_mask}
        return config, inputs_dict







@require_tf
class     __lowerCAmelCase	(   __UpperCamelCase						,    __UpperCamelCase						,    unittest.TestCase):
      _lowercase    : List[Any]            =      (
          (
              TFEsmModel,
              TFEsmForMaskedLM,
              TFEsmForSequenceClassification,
              TFEsmForTokenClassification,
          )
          if is_tf_available()
          else ()
      )
      _lowercase    : Dict            =      (
          {
              """feature-extraction""": TFEsmModel,
              """fill-mask""": TFEsmForMaskedLM,
              """text-classification""": TFEsmForSequenceClassification,
              """token-classification""": TFEsmForTokenClassification,
              """zero-shot""": TFEsmForSequenceClassification,
          }
          if is_tf_available()
          else {}
      )
      _lowercase    : Tuple            =      False
      _lowercase    : Optional[int]            =      False







      def        _lowercase  (   self							)       ->						int:



        '''simple docstring'''




        a__     : Dict  =TFEsmModelTester(self							)
        a__     : Any  =ConfigTester(self		,				config_class=lowerCAmelCase__		,				hidden_size=3_7							)







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



        '''simple docstring'''




        self.config_tester.run_common_tests()







      def        _lowercase  (   self							)       ->						int:



        '''simple docstring'''




        a__     : Any  =self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_model(*lowerCAmelCase__							)







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



        '''simple docstring'''




        a__     : List[str]  =self.model_tester.prepare_config_and_inputs_for_decoder()
        self.model_tester.create_and_check_model_as_decoder(*lowerCAmelCase__							)







      def        _lowercase  (   self							)       ->						int:



        '''simple docstring'''




        a__     : str  =self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_masked_lm(*lowerCAmelCase__							)







      def        _lowercase  (   self							)       ->						str:



        '''simple docstring'''




        a__     : List[Any]  =self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_token_classification(*lowerCAmelCase__							)







      @slow
      def        _lowercase  (   self							)       ->						Dict:



        '''simple docstring'''




        for model_name in TF_ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
          a__     : str  =TFEsmModel.from_pretrained(lowerCAmelCase__							)
          self.assertIsNotNone(lowerCAmelCase__							)







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



        '''simple docstring'''




        pass







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



        '''simple docstring'''




        pass







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



        '''simple docstring'''




        a__     : List[str]  =self.model_tester.prepare_config_and_inputs_for_common()

        for model_class in self.all_model_classes:
          a__     : Union[str, Any]  =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
            a__     : Optional[int]  =model.get_bias()
            assert isinstance(lowerCAmelCase__		,				lowerCAmelCase__							)
            for k, v in name.items():
              assert isinstance(lowerCAmelCase__		,				tf.Variable							)
          else:
            a__     : Optional[Any]  =model.get_output_embeddings()
            assert x is None
            a__     : Dict  =model.get_bias()
            assert name is None







@require_tf
class     __lowerCAmelCase	(   unittest.TestCase):







      @slow
      def        _lowercase  (   self							)       ->						Dict:



        '''simple docstring'''




        a__     : Optional[Any]  =TFEsmForMaskedLM.from_pretrained("facebook/esm2_t6_8M_UR50D"							)

        a__     : List[Any]  =tf.constant([[0, 1, 2, 3, 4, 5]]							)
        a__     : Dict  =model(lowerCAmelCase__							)[0]
        a__     : List[str]  =[1, 6, 3_3]
        self.assertEqual(list(output.numpy().shape							)		,				lowerCAmelCase__							)
        # compare the actual values for a slice.
        a__     : Dict  =tf.constant(
            [
                [
                    [8.92_15_18, -1_0.5_8_9_8_1_4, -6.4_67_13_07],
                    [-6.3_96_71_56, -1_3.9_1_1_3_7_7, -1.1_21_19_15],
                    [-7.78_12_47, -1_3.9_5_1_5_5_7, -3.74_05_92],
                ]
            ]							)
        self.assertTrue(numpy.allclose(output[:, :3, :3].numpy()		,				expected_slice.numpy()		,				atol=1E-2							)							)







      @slow
      def        _lowercase  (   self							)       ->						Tuple:



        '''simple docstring'''




        a__     : str  =TFEsmModel.from_pretrained("facebook/esm2_t6_8M_UR50D"							)

        a__     : List[str]  =tf.constant([[0, 6, 4, 1_3, 5, 4, 1_6, 1_2, 1_1, 7, 2]]							)
        a__     : Optional[int]  =model(lowerCAmelCase__							)[0]
        # compare the actual values for a slice.
        a__     : str  =tf.constant(
            [
                [
                    [0.14_44_30_92, 0.54_12_53_27, 0.3_24_77_39],
                    [0.30_34_04_84, 0.00_52_66_76, 0.31_07_77_22],
                    [0.32_27_80_43, -0.24_98_70_96, 0.3_41_46_28],
                ]
            ]							)
        self.assertTrue(numpy.allclose(output[:, :3, :3].numpy()		,				expected_slice.numpy()		,				atol=1E-4							)							)







import random
from typing import Any
def 					_A	(		SCREAMING_SNAKE_CASE					:       list						):





  """simple docstring"""



  for _ in range(len(SCREAMING_SNAKE_CASE						)						):
    a__     : Tuple  =random.randint(0	,	len(SCREAMING_SNAKE_CASE						) - 1						)
    a__     : Any  =random.randint(0	,	len(SCREAMING_SNAKE_CASE						) - 1						)
    a__ ,							a__     : Tuple  =data[b], data[a]
  return data


if __name__ == "__main__":
     UpperCAmelCase     :     Dict							=	[0, 1, 2, 3, 4, 5, 6, 7]
     UpperCAmelCase     :     Optional[Any]							=	["""python""", """says""", """hello""", """!"""]
     print("""Fisher-Yates Shuffle:""")
     print("""List""", integers, strings)
     print("""FY Shuffle""", fisher_yates_shuffle(integers), fisher_yates_shuffle(strings))






