

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     __snake_case    :
       """simple docstring"""







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






              """simple docstring"""


              __snake_case       	=							parent
              __snake_case       	=							13
              __snake_case       	=							7
              __snake_case       	=							True
              __snake_case       	=							True
              __snake_case       	=							True
              __snake_case       	=							99
              __snake_case       	=							32
              __snake_case       	=							2
              __snake_case       	=							4
              __snake_case       	=							37
              __snake_case       	=							"""gelu"""
              __snake_case       	=							0.1
              __snake_case       	=							0.1
              __snake_case       	=							5_12
              __snake_case       	=							16
              __snake_case       	=							2
              __snake_case       	=							0.02
              __snake_case       	=							3
              __snake_case       	=							4
              __snake_case       	=							None


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






              """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 			a		(		self   )      ->							Optional[Any]:






              """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 			a		(		self      , _UpperCamelCase      , _UpperCamelCase      , _UpperCamelCase      , _UpperCamelCase      , _UpperCamelCase      , _UpperCamelCase   )      ->							str:






              """simple docstring"""


              __snake_case       	=							TFEsmModel(config=_UpperCamelCase   )
              __snake_case       	=							{"""input_ids""": input_ids, """attention_mask""": input_mask}
              __snake_case       	=							model(_UpperCamelCase   )

              __snake_case       	=							[input_ids, input_mask]
              __snake_case       	=							model(_UpperCamelCase   )

              __snake_case       	=							model(_UpperCamelCase   )

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


       def 			a		(		self      , _UpperCamelCase      , _UpperCamelCase      , _UpperCamelCase      , _UpperCamelCase      , _UpperCamelCase      , _UpperCamelCase      , _UpperCamelCase      , _UpperCamelCase      , )      ->							Optional[int]:






              """simple docstring"""


              __snake_case       	=							True

              __snake_case       	=							TFEsmModel(config=_UpperCamelCase   )
              __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(_UpperCamelCase   )

              __snake_case       	=							[input_ids, input_mask]
              __snake_case       	=							model(_UpperCamelCase      , encoder_hidden_states=_UpperCamelCase   )

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

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


       def 			a		(		self      , _UpperCamelCase      , _UpperCamelCase      , _UpperCamelCase      , _UpperCamelCase      , _UpperCamelCase      , _UpperCamelCase   )      ->							Tuple:






              """simple docstring"""


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


       def 			a		(		self      , _UpperCamelCase      , _UpperCamelCase      , _UpperCamelCase      , _UpperCamelCase      , _UpperCamelCase      , _UpperCamelCase   )      ->							Tuple:






              """simple docstring"""


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






       def 			a		(		self   )      ->							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     __snake_case    (				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 			a		(		self   )      ->							str:






              """simple docstring"""


              __snake_case       	=							TFEsmModelTester(self   )
              __snake_case       	=							ConfigTester(self      , config_class=_UpperCamelCase      , hidden_size=37   )


       def 			a		(		self   )      ->							Any:






              """simple docstring"""


              self.config_tester.run_common_tests()


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






              """simple docstring"""


              __snake_case       	=							self.model_tester.prepare_config_and_inputs()
              self.model_tester.create_and_check_model(*_UpperCamelCase   )


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






              """simple docstring"""


              __snake_case       	=							self.model_tester.prepare_config_and_inputs_for_decoder()
              self.model_tester.create_and_check_model_as_decoder(*_UpperCamelCase   )


       def 			a		(		self   )      ->							str:






              """simple docstring"""


              __snake_case       	=							self.model_tester.prepare_config_and_inputs()
              self.model_tester.create_and_check_for_masked_lm(*_UpperCamelCase   )


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






              """simple docstring"""


              __snake_case       	=							self.model_tester.prepare_config_and_inputs()
              self.model_tester.create_and_check_for_token_classification(*_UpperCamelCase   )


       @slow
       def 			a		(		self   )      ->							Dict:






              """simple docstring"""


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


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






              """simple docstring"""


              pass


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






              """simple docstring"""


              pass






       def 			a		(		self   )      ->							List[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(_UpperCamelCase   )
                     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(_UpperCamelCase      , _UpperCamelCase   )
                            for k, v in name.items():
                                   assert isinstance(_UpperCamelCase      , 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     __snake_case    (				unittest.TestCase  ):
       """simple docstring"""







       @slow
       def 			a		(		self   )      ->							int:






              """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(_UpperCamelCase   )[0]
              __snake_case       	=							[1, 6, 33]
              self.assertEqual(list(output.numpy().shape   )      , _UpperCamelCase   )
              # compare the actual values for a slice.
              __snake_case       	=							tf.constant(
                  [
                      [
                          [8.92_1518, -10.58_9814, -6.467_1307],
                          [-6.396_7156, -13.91_1377, -1.121_1915],
                          [-7.78_1247, -13.95_1557, -3.74_0592],
                      ]
                  ]   )
              self.assertTrue(numpy.allclose(output[:, :3, :3].numpy()      , expected_slice.numpy()      , atol=1E-2   )   )






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






              """simple docstring"""


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

              __snake_case       	=							tf.constant([[0, 6, 4, 13, 5, 4, 16, 12, 11, 7, 2]]   )
              __snake_case       	=							model(_UpperCamelCase   )[0]
              # compare the actual values for a slice.
              __snake_case       	=							tf.constant(
                  [
                      [
                          [0.1444_3092, 0.5412_5327, 0.324_7739],
                          [0.3034_0484, 0.0052_6676, 0.3107_7722],
                          [0.3227_8043, -0.2498_7096, 0.341_4628],
                      ]
                  ]   )
              self.assertTrue(numpy.allclose(output[:, :3, :3].numpy()      , expected_slice.numpy()      , atol=1E-4   )   )




def 						lowerCamelCase__      ( __A    :int	,__A    :int		):


       """simple docstring"""




       if a < 0 or b < 0:
              raise ValueError("""the value of both inputs must be positive"""		)

       __snake_case       	=							str(bin(__A		)		)[2:]  # remove the leading "0b"
       __snake_case       	=							str(bin(__A		)		)[2:]  # remove the leading "0b"

       __snake_case       	=							max(len(__A		)	,len(__A		)		)

       return "0b" + "".join(
           str(int(char_a == """1""" and char_b == """1"""		)		)
           for char_a, char_b in zip(a_binary.zfill(__A		)	,b_binary.zfill(__A		)		)		)


if __name__ == "__main__":
    import doctest

    doctest.testmod()

