



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							:	Any ,						lowerCAmelCase							:	str ,						) -> Tuple:

    """simple docstring"""







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

   def    UpperCamelCase_     (    self							:	Optional[Any]) -> str:

    """simple docstring"""







    _snake_case			:     List[str] 		=    ids_tensor([self.batch_size, self.seq_length] ,						self.vocab_size)

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

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

    _snake_case			:     Dict 		=    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							:	Union[str, Any]) -> Union[str, Any]:

    """simple docstring"""







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

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

    """simple docstring"""







    _snake_case			:     List[Any] 		=    TFEsmModel(config=__SCREAMING_SNAKE_CASE)
    _snake_case			:     Tuple 		=    {"""input_ids""": input_ids, """attention_mask""": input_mask}
    _snake_case			:     str 		=    model(__SCREAMING_SNAKE_CASE)

    _snake_case			:     Optional[int] 		=    [input_ids, input_mask]
    _snake_case			:     str 		=    model(__SCREAMING_SNAKE_CASE)

    _snake_case			:     Optional[Any] 		=    model(__SCREAMING_SNAKE_CASE)

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

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

    """simple docstring"""







    _snake_case			:     List[Any] 		=    True

    _snake_case			:     int 		=    TFEsmModel(config=__SCREAMING_SNAKE_CASE)
    _snake_case			:     Optional[Any] 		=    {
        """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(__SCREAMING_SNAKE_CASE)

    _snake_case			:     Optional[Any] 		=    [input_ids, input_mask]
    _snake_case			:     str 		=    model(__SCREAMING_SNAKE_CASE ,						encoder_hidden_states=__SCREAMING_SNAKE_CASE)

    # Also check the case where encoder outputs are not passed
    _snake_case			:     int 		=    model(__SCREAMING_SNAKE_CASE ,						attention_mask=__SCREAMING_SNAKE_CASE)

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

   def    UpperCamelCase_     (    self							:	Optional[int] ,						lowerCAmelCase							:	str ,						lowerCAmelCase							:	Optional[Any] ,						lowerCAmelCase							:	int ,						lowerCAmelCase							:	Optional[Any] ,						lowerCAmelCase							:	int ,						lowerCAmelCase							:	List[Any]) -> int:

    """simple docstring"""







    _snake_case			:     Optional[Any] 		=    TFEsmForMaskedLM(config=__SCREAMING_SNAKE_CASE)
    _snake_case			:     str 		=    model([input_ids, input_mask])
    self.parent.assertEqual(result.logits.shape ,						(self.batch_size, self.seq_length, self.vocab_size))

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

    """simple docstring"""







    _snake_case			:     List[Any] 		=    self.num_labels
    _snake_case			:     str 		=    TFEsmForTokenClassification(config=__SCREAMING_SNAKE_CASE)
    _snake_case			:     str 		=    {"""input_ids""": input_ids, """attention_mask""": input_mask}
    _snake_case			:     Any 		=    model(__SCREAMING_SNAKE_CASE)
    self.parent.assertEqual(result.logits.shape ,						(self.batch_size, self.seq_length, self.num_labels))


   def    UpperCamelCase_     (    self							:	List[Any]) -> int:

    """simple docstring"""







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







@require_tf
class 						snake_case  (		SCREAMING_SNAKE_CASE_    ,SCREAMING_SNAKE_CASE_    ,unittest.TestCase     ):



   '''simple docstring'''






   snake_case_    :       Tuple								=						(
       (
           TFEsmModel,
           TFEsmForMaskedLM,
           TFEsmForSequenceClassification,
           TFEsmForTokenClassification,
       )
       if is_tf_available()
       else ()
   )
   snake_case_    :       Any								=						(
       {
           """feature-extraction""": TFEsmModel,
           """fill-mask""": TFEsmForMaskedLM,
           """text-classification""": TFEsmForSequenceClassification,
           """token-classification""": TFEsmForTokenClassification,
           """zero-shot""": TFEsmForSequenceClassification,
       }
       if is_tf_available()
       else {}
   )
   snake_case_    :       List[Any]								=						False
   snake_case_    :       Tuple								=						False

   def    UpperCamelCase_     (    self							:	Dict) -> Any:

    """simple docstring"""







    _snake_case			:     Optional[Any] 		=    TFEsmModelTester(self)
    _snake_case			:     str 		=    ConfigTester(self ,						config_class=__SCREAMING_SNAKE_CASE ,						hidden_size=37)

   def    UpperCamelCase_     (    self							:	Optional[Any]) -> Any:

    """simple docstring"""







    self.config_tester.run_common_tests()

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

    """simple docstring"""







    _snake_case			:     Union[str, Any] 		=    self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_model(*__SCREAMING_SNAKE_CASE)

   def    UpperCamelCase_     (    self							:	Any) -> Optional[Any]:

    """simple docstring"""







    _snake_case			:     Union[str, Any] 		=    self.model_tester.prepare_config_and_inputs_for_decoder()
    self.model_tester.create_and_check_model_as_decoder(*__SCREAMING_SNAKE_CASE)

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

    """simple docstring"""







    _snake_case			:     Tuple 		=    self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_for_masked_lm(*__SCREAMING_SNAKE_CASE)

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

    """simple docstring"""







    _snake_case			:     Optional[Any] 		=    self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_for_token_classification(*__SCREAMING_SNAKE_CASE)

   @slow
   def    UpperCamelCase_     (    self							:	str) -> int:

    """simple docstring"""







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

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

    """simple docstring"""







    pass

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

    """simple docstring"""







    pass


   def    UpperCamelCase_     (    self							:	str) -> Tuple:

    """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			:     Optional[int] 		=    model_class(__SCREAMING_SNAKE_CASE)
     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(__SCREAMING_SNAKE_CASE ,						__SCREAMING_SNAKE_CASE)
      for k, v in name.items():
       assert isinstance(__SCREAMING_SNAKE_CASE ,						tf.Variable)
     else:
      _snake_case			:     str 		=    model.get_output_embeddings()
      assert x is None
      _snake_case			:     str 		=    model.get_bias()
      assert name is None







@require_tf
class 						snake_case  (		unittest.TestCase     ):



   '''simple docstring'''






   @slow
   def    UpperCamelCase_     (    self							:	Tuple) -> Any:

    """simple docstring"""







    _snake_case			:     str 		=    TFEsmForMaskedLM.from_pretrained("""facebook/esm2_t6_8M_UR50D""")

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

    """simple docstring"""







    _snake_case			:     int 		=    TFEsmModel.from_pretrained("""facebook/esm2_t6_8M_UR50D""")

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



from ...configuration_utils import PretrainedConfig
from ...utils import logging
from ...utils.backbone_utils import BackboneConfigMixin, get_aligned_output_features_output_indices


UpperCAmelCase :       List[str]			   =     logging.get_logger(__name__)







class 						lowerCAmelCase__   (      a ,					a					):



 """simple docstring"""



 lowerCAmelCase__  					=      "maskformer-swin"

 lowerCAmelCase__  					=      {
     "num_attention_heads": "num_heads",
     "num_hidden_layers": "num_layers",
 }
 def __init__(		self					:       str					,							__SCREAMING_SNAKE_CASE					:       Tuple=224					,							__SCREAMING_SNAKE_CASE					:       str=4					,							__SCREAMING_SNAKE_CASE					:       Union[str, Any]=3					,							__SCREAMING_SNAKE_CASE					:       Optional[Any]=96					,							__SCREAMING_SNAKE_CASE					:       Optional[Any]=[2, 2, 6, 2]					,							__SCREAMING_SNAKE_CASE					:       Any=[3, 6, 12, 24]					,							__SCREAMING_SNAKE_CASE					:       Dict=7					,							__SCREAMING_SNAKE_CASE					:       Dict=4.0					,							__SCREAMING_SNAKE_CASE					:       Dict=True					,							__SCREAMING_SNAKE_CASE					:       str=0.0					,							__SCREAMING_SNAKE_CASE					:       int=0.0					,							__SCREAMING_SNAKE_CASE					:       str=0.1					,							__SCREAMING_SNAKE_CASE					:       List[Any]="gelu"					,							__SCREAMING_SNAKE_CASE					:       str=False					,							__SCREAMING_SNAKE_CASE					:       Optional[int]=0.02					,							__SCREAMING_SNAKE_CASE					:       Optional[int]=1E-5					,							__SCREAMING_SNAKE_CASE					:       Dict=None					,							__SCREAMING_SNAKE_CASE					:       Dict=None					,							**__SCREAMING_SNAKE_CASE					:       Tuple					,							)  -> Tuple:


      """simple docstring"""




      super().__init__(**__SCREAMING_SNAKE_CASE			)

      __SCREAMING_SNAKE_CASE     =      image_size
      __SCREAMING_SNAKE_CASE     =      patch_size
      __SCREAMING_SNAKE_CASE     =      num_channels
      __SCREAMING_SNAKE_CASE     =      embed_dim
      __SCREAMING_SNAKE_CASE     =      depths
      __SCREAMING_SNAKE_CASE     =      len(__SCREAMING_SNAKE_CASE			)
      __SCREAMING_SNAKE_CASE     =      num_heads
      __SCREAMING_SNAKE_CASE     =      window_size
      __SCREAMING_SNAKE_CASE     =      mlp_ratio
      __SCREAMING_SNAKE_CASE     =      qkv_bias
      __SCREAMING_SNAKE_CASE     =      hidden_dropout_prob
      __SCREAMING_SNAKE_CASE     =      attention_probs_dropout_prob
      __SCREAMING_SNAKE_CASE     =      drop_path_rate
      __SCREAMING_SNAKE_CASE     =      hidden_act
      __SCREAMING_SNAKE_CASE     =      use_absolute_embeddings
      __SCREAMING_SNAKE_CASE     =      layer_norm_eps
      __SCREAMING_SNAKE_CASE     =      initializer_range
      # we set the hidden_size attribute in order to make Swin work with VisionEncoderDecoderModel
      # this indicates the channel dimension after the last stage of the model
      __SCREAMING_SNAKE_CASE     =      int(embed_dim * 2 ** (len(__SCREAMING_SNAKE_CASE			) - 1)			)
      __SCREAMING_SNAKE_CASE     =      ["""stem"""] + [f'stage{idx}' for idx in range(1					,							len(__SCREAMING_SNAKE_CASE			) + 1			)]
      __SCREAMING_SNAKE_CASE   ,			__SCREAMING_SNAKE_CASE     =      get_aligned_output_features_output_indices(
          out_features=__SCREAMING_SNAKE_CASE					,							out_indices=__SCREAMING_SNAKE_CASE					,							stage_names=self.stage_names			)






