





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					,    UpperCAmelCase					,    )							-> List[Any]:






        '''simple docstring'''






        __snake_case :					Optional[Any]          =							parent
        __snake_case :					List[str]          =							13
        __snake_case :					str          =							7
        __snake_case :					str          =							True
        __snake_case :					Dict          =							True
        __snake_case :					str          =							True
        __snake_case :					List[str]          =							99
        __snake_case :					List[Any]          =							32
        __snake_case :					List[str]          =							2
        __snake_case :					Union[str, Any]          =							4
        __snake_case :					List[str]          =							37
        __snake_case :					Optional[int]          =							"gelu"
        __snake_case :					int          =							0.1
        __snake_case :					List[str]          =							0.1
        __snake_case :					List[Any]          =							512
        __snake_case :					Optional[int]          =							16
        __snake_case :					Optional[Any]          =							2
        __snake_case :					Optional[Any]          =							0.02
        __snake_case :					Union[str, Any]          =							3
        __snake_case :					str          =							4
        __snake_case :					str          =							None







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






        '''simple docstring'''






        __snake_case :					Any          =							ids_tensor([self.batch_size, self.seq_length]					,    self.vocab_size )

        __snake_case :					int          =							None
        if self.use_input_mask:
          __snake_case :					Tuple          =							random_attention_mask([self.batch_size, self.seq_length] )

        __snake_case :					Optional[Any]          =							None
        __snake_case :					Union[str, Any]          =							None
        __snake_case :					Tuple          =							None
        if self.use_labels:
          __snake_case :					int          =							ids_tensor([self.batch_size]					,    self.type_sequence_label_size )
          __snake_case :					List[Any]          =							ids_tensor([self.batch_size, self.seq_length]					,    self.num_labels )
          __snake_case :					Optional[Any]          =							ids_tensor([self.batch_size]					,    self.num_choices )

        __snake_case :					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     UpperCAmelCase		(			self )							-> int:






        '''simple docstring'''






        (
            (
            __snake_case							
        )     , (
            __snake_case							
        )     , (
            __snake_case							
        )     , (
            __snake_case							
        )     , (
            __snake_case							
        )     , (
            __snake_case							
        )     , 							
        ) :					str          =							self.prepare_config_and_inputs()

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






        '''simple docstring'''






        __snake_case :					int          =							TFEsmModel(config=lowerCAmelCase__ )
        __snake_case :					Union[str, Any]          =							{"input_ids": input_ids, "attention_mask": input_mask}
        __snake_case :					int          =							model(lowerCAmelCase__ )

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

        __snake_case :					Tuple          =							model(lowerCAmelCase__ )

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







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






        '''simple docstring'''






        __snake_case :					List[Any]          =							True

        __snake_case :					Any          =							TFEsmModel(config=lowerCAmelCase__ )
        __snake_case :					Tuple          =							{
            "input_ids": input_ids,
            "attention_mask": input_mask,
            "encoder_hidden_states": encoder_hidden_states,
            "encoder_attention_mask": encoder_attention_mask,
        }
        __snake_case :					List[str]          =							model(lowerCAmelCase__ )

        __snake_case :					Union[str, Any]          =							[input_ids, input_mask]
        __snake_case :					List[str]          =							model(lowerCAmelCase__					,    encoder_hidden_states=lowerCAmelCase__ )

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

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







      def     UpperCAmelCase		(			self					,    UpperCAmelCase					,    UpperCAmelCase					,    UpperCAmelCase					,    UpperCAmelCase					,    UpperCAmelCase					,    UpperCAmelCase )							-> str:






        '''simple docstring'''






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







      def     UpperCAmelCase		(			self					,    UpperCAmelCase					,    UpperCAmelCase					,    UpperCAmelCase					,    UpperCAmelCase					,    UpperCAmelCase					,    UpperCAmelCase )							-> Dict:






        '''simple docstring'''






        __snake_case :					Optional[int]          =							self.num_labels
        __snake_case :					str          =							TFEsmForTokenClassification(config=lowerCAmelCase__ )
        __snake_case :					Tuple          =							{"input_ids": input_ids, "attention_mask": input_mask}
        __snake_case :					Dict          =							model(lowerCAmelCase__ )
        self.parent.assertEqual(result.logits.shape					,    (self.batch_size, self.seq_length, self.num_labels) )







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






        '''simple docstring'''






        __snake_case :					str          =							self.prepare_config_and_inputs()
        (
            (
            __snake_case							
        )     , (
            __snake_case							
        )     , (
            __snake_case							
        )     , (
            __snake_case							
        )     , (
            __snake_case							
        )     , (
            __snake_case							
        )     , 							
        ) :					str          =							config_and_inputs
        __snake_case :					str          =							{"input_ids": input_ids, "attention_mask": input_mask}
        return config, inputs_dict




@require_tf
class 					_lowerCamelCase			(     a			, a			, unittest.TestCase     ):



      """simple docstring"""






      UpperCAmelCase_    :   Optional[int]      =(
          (
              TFEsmModel,
              TFEsmForMaskedLM,
              TFEsmForSequenceClassification,
              TFEsmForTokenClassification,
          )
          if is_tf_available()
          else ()
      )
      UpperCAmelCase_    :   Optional[Any]      =(
          {
              'feature-extraction': TFEsmModel,
              'fill-mask': TFEsmForMaskedLM,
              'text-classification': TFEsmForSequenceClassification,
              'token-classification': TFEsmForTokenClassification,
              'zero-shot': TFEsmForSequenceClassification,
          }
          if is_tf_available()
          else {}
      )
      UpperCAmelCase_    :   Dict      =False
      UpperCAmelCase_    :   int      =False







      def     UpperCAmelCase		(			self )							-> Dict:






        '''simple docstring'''






        __snake_case :					Dict          =							TFEsmModelTester(self )
        __snake_case :					int          =							ConfigTester(self					,    config_class=lowerCAmelCase__					,    hidden_size=37 )







      def     UpperCAmelCase		(			self )							-> Dict:






        '''simple docstring'''






        self.config_tester.run_common_tests()







      def     UpperCAmelCase		(			self )							-> str:






        '''simple docstring'''






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







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






        '''simple docstring'''






        __snake_case :					Optional[int]          =							self.model_tester.prepare_config_and_inputs_for_decoder()
        self.model_tester.create_and_check_model_as_decoder(*lowerCAmelCase__ )







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






        '''simple docstring'''






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







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






        '''simple docstring'''






        __snake_case :					Union[str, Any]          =							self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_token_classification(*lowerCAmelCase__ )







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






        '''simple docstring'''






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







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






        '''simple docstring'''






        pass







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






        '''simple docstring'''






        pass







      def     UpperCAmelCase		(			self )							-> str:






        '''simple docstring'''






        __snake_case     , __snake_case :					Dict          =							self.model_tester.prepare_config_and_inputs_for_common()

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




@require_tf
class 					_lowerCamelCase			(     unittest.TestCase     ):



      """simple docstring"""






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






        '''simple docstring'''






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

        __snake_case :					Union[str, Any]          =							tf.constant([[0, 1, 2, 3, 4, 5]] )
        __snake_case :					str          =							model(lowerCAmelCase__ )[0]
        __snake_case :					Any          =							[1, 6, 33]
        self.assertEqual(list(output.numpy().shape )					,    lowerCAmelCase__ )
        # compare the actual values for a slice.
        __snake_case :					Any          =							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     UpperCAmelCase		(			self )							-> List[Any]:






        '''simple docstring'''






        __snake_case :					Optional[Any]          =							TFEsmModel.from_pretrained("facebook/esm2_t6_8M_UR50D" )

        __snake_case :					Any          =							tf.constant([[0, 6, 4, 13, 5, 4, 16, 12, 11, 7, 2]] )
        __snake_case :					List[str]          =							model(lowerCAmelCase__ )[0]
        # compare the actual values for a slice.
        __snake_case :					str          =							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 ) )








'''simple docstring'''

def    a__							(							lowercase   :	Dict,   lowercase   :	Optional[Any] )						->						int:





     """simple docstring"""

     _UpperCamelCase							       =			0
     _UpperCamelCase							       =			len(lowercase ) - 1

     while left <= right:
          # avoid divided by 0 during interpolation
          if sorted_collection[left] == sorted_collection[right]:
               if sorted_collection[left] == item:
                    return left
               else:
                    return None

          _UpperCamelCase							       =			left + ((item - sorted_collection[left]) * (right - left)) // (
              sorted_collection[right] - sorted_collection[left]
          )

          # out of range check
          if point < 0 or point >= len(lowercase ):
               return None

          _UpperCamelCase							       =			sorted_collection[point]
          if current_item == item:
               return point
          else:
               if point < left:
                    _UpperCamelCase							       =			left
                    _UpperCamelCase							       =			point
               elif point > right:
                    _UpperCamelCase							       =			right
                    _UpperCamelCase							       =			point
               else:
                    if item < current_item:
                         _UpperCamelCase							       =			point - 1
                    else:
                         _UpperCamelCase							       =			point + 1
     return None







def    a__							(							lowercase   :	int,   lowercase   :	Optional[int],   lowercase   :	List[Any],   lowercase   :	List[Any] )						->						Dict:





     """simple docstring"""

     if sorted_collection[left] == sorted_collection[right]:
          if sorted_collection[left] == item:
               return left
          else:
               return None

     _UpperCamelCase							       =			left + ((item - sorted_collection[left]) * (right - left)) // (
         sorted_collection[right] - sorted_collection[left]
     )

     # out of range check
     if point < 0 or point >= len(lowercase ):
          return None

     if sorted_collection[point] == item:
          return point
     elif point < left:
          return interpolation_search_by_recursion(lowercase,   lowercase,   lowercase,   lowercase )
     elif point > right:
          return interpolation_search_by_recursion(lowercase,   lowercase,   lowercase,   lowercase )
     else:
          if sorted_collection[point] > item:
               return interpolation_search_by_recursion(
                   lowercase,   lowercase,   lowercase,   point - 1 )
          else:
               return interpolation_search_by_recursion(
                   lowercase,   lowercase,   point + 1,   lowercase )







def    a__							(							lowercase   :	Any )						->						List[Any]:





     """simple docstring"""

     if collection != sorted(lowercase ):
          raise ValueError('''Collection must be ascending sorted''' )
     return True


if __name__ == "__main__":
   import sys
   lowercase__		:     int			=							0
   if debug == 1:
      lowercase__		:     Optional[Any]			=							[10, 30, 40, 45, 50, 66, 77, 93]
      try:
         __assert_sorted(collection)
      except ValueError:
         sys.exit('Sequence must be ascending sorted to apply interpolation search')
      lowercase__		:     List[str]			=							67

   lowercase__		:     Optional[int]			=							interpolation_search(collection, target)
   if result is not None:
      print(F"""{target} found at positions: {result}""")
   else:
      print('Not found')



