






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  SCREAMING_SNAKE_CASE			:
     '''simple docstring'''


     def __init__(				self       :	Any					, UpperCAmelCase_       :	List[str]					, ):
       SCREAMING_SNAKE_CASE       :      str							=    parent
       SCREAMING_SNAKE_CASE       :      Optional[Any]							=    13
       SCREAMING_SNAKE_CASE       :      Dict							=    7
       SCREAMING_SNAKE_CASE       :      Optional[int]							=    True
       SCREAMING_SNAKE_CASE       :      List[str]							=    True
       SCREAMING_SNAKE_CASE       :      str							=    True
       SCREAMING_SNAKE_CASE       :      Optional[Any]							=    99
       SCREAMING_SNAKE_CASE       :      int							=    32
       SCREAMING_SNAKE_CASE       :      Optional[int]							=    2
       SCREAMING_SNAKE_CASE       :      int							=    4
       SCREAMING_SNAKE_CASE       :      Tuple							=    37
       SCREAMING_SNAKE_CASE       :      List[str]							=    "gelu"
       SCREAMING_SNAKE_CASE       :      Union[str, Any]							=    0.1
       SCREAMING_SNAKE_CASE       :      Optional[Any]							=    0.1
       SCREAMING_SNAKE_CASE       :      List[str]							=    512
       SCREAMING_SNAKE_CASE       :      Tuple							=    16
       SCREAMING_SNAKE_CASE       :      Optional[int]							=    2
       SCREAMING_SNAKE_CASE       :      str							=    0.02
       SCREAMING_SNAKE_CASE       :      Tuple							=    3
       SCREAMING_SNAKE_CASE       :      Optional[Any]							=    4
       SCREAMING_SNAKE_CASE       :      Tuple							=    None


     def 				_A       (				self       :	int      ):
       SCREAMING_SNAKE_CASE       :      Tuple							=    ids_tensor([self.batch_size, self.seq_length]					, self.vocab_size      )

       SCREAMING_SNAKE_CASE       :      Optional[int]							=    None
       if self.use_input_mask:
         SCREAMING_SNAKE_CASE       :      str							=    random_attention_mask([self.batch_size, self.seq_length]      )

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

       SCREAMING_SNAKE_CASE       :      List[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 				_A       (				self       :	List[Any]      ):
       (
           (
           SCREAMING_SNAKE_CASE      
       )						,  (
           SCREAMING_SNAKE_CASE      
       )						,  (
           SCREAMING_SNAKE_CASE      
       )						,  (
           SCREAMING_SNAKE_CASE      
       )						,  (
           SCREAMING_SNAKE_CASE      
       )						,  (
           SCREAMING_SNAKE_CASE      
       )						,        
       )       :      List[Any]							=    self.prepare_config_and_inputs()

       SCREAMING_SNAKE_CASE       :      List[str]							=    True
       SCREAMING_SNAKE_CASE       :      List[str]							=    floats_tensor([self.batch_size, self.seq_length, self.hidden_size]      )
       SCREAMING_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 				_A       (				self       :	Union[str, Any]					, UpperCAmelCase_       :	Union[str, Any]					, UpperCAmelCase_       :	Any					, UpperCAmelCase_       :	int					, UpperCAmelCase_       :	Union[str, Any]					, UpperCAmelCase_       :	Union[str, Any]					, UpperCAmelCase_       :	Any      ):
       SCREAMING_SNAKE_CASE       :      Union[str, Any]							=    TFEsmModel(config=_SCREAMING_SNAKE_CASE      )
       SCREAMING_SNAKE_CASE       :      Tuple							=    {"input_ids": input_ids, "attention_mask": input_mask}
       SCREAMING_SNAKE_CASE       :      Dict							=    model(_SCREAMING_SNAKE_CASE      )

       SCREAMING_SNAKE_CASE       :      Union[str, Any]							=    [input_ids, input_mask]
       SCREAMING_SNAKE_CASE       :      List[str]							=    model(_SCREAMING_SNAKE_CASE      )

       SCREAMING_SNAKE_CASE       :      Dict							=    model(_SCREAMING_SNAKE_CASE      )

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


     def 				_A       (				self       :	List[str]					, UpperCAmelCase_       :	Union[str, Any]					, UpperCAmelCase_       :	Tuple					, UpperCAmelCase_       :	Union[str, Any]					, UpperCAmelCase_       :	Tuple					, UpperCAmelCase_       :	str					, UpperCAmelCase_       :	Optional[Any]					, UpperCAmelCase_       :	Optional[int]					, UpperCAmelCase_       :	Dict					, ):
       SCREAMING_SNAKE_CASE       :      Union[str, Any]							=    True

       SCREAMING_SNAKE_CASE       :      Dict							=    TFEsmModel(config=_SCREAMING_SNAKE_CASE      )
       SCREAMING_SNAKE_CASE       :      int							=    {
           "input_ids": input_ids,
           "attention_mask": input_mask,
           "encoder_hidden_states": encoder_hidden_states,
           "encoder_attention_mask": encoder_attention_mask,
       }
       SCREAMING_SNAKE_CASE       :      Dict							=    model(_SCREAMING_SNAKE_CASE      )

       SCREAMING_SNAKE_CASE       :      Dict							=    [input_ids, input_mask]
       SCREAMING_SNAKE_CASE       :      Optional[int]							=    model(_SCREAMING_SNAKE_CASE					, encoder_hidden_states=_SCREAMING_SNAKE_CASE      )

       # Also check the case where encoder outputs are not passed
       SCREAMING_SNAKE_CASE       :      List[Any]							=    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 				_A       (				self       :	Tuple					, UpperCAmelCase_       :	Union[str, Any]					, UpperCAmelCase_       :	Optional[int]					, UpperCAmelCase_       :	int					, UpperCAmelCase_       :	int					, UpperCAmelCase_       :	int					, UpperCAmelCase_       :	Any      ):
       SCREAMING_SNAKE_CASE       :      List[str]							=    TFEsmForMaskedLM(config=_SCREAMING_SNAKE_CASE      )
       SCREAMING_SNAKE_CASE       :      Dict							=    model([input_ids, input_mask]      )
       self.parent.assertEqual(result.logits.shape					, (self.batch_size, self.seq_length, self.vocab_size)      )


     def 				_A       (				self       :	Optional[Any]					, UpperCAmelCase_       :	Dict					, UpperCAmelCase_       :	Any					, UpperCAmelCase_       :	Optional[int]					, UpperCAmelCase_       :	Optional[int]					, UpperCAmelCase_       :	str					, UpperCAmelCase_       :	Union[str, Any]      ):
       SCREAMING_SNAKE_CASE       :      Tuple							=    self.num_labels
       SCREAMING_SNAKE_CASE       :      int							=    TFEsmForTokenClassification(config=_SCREAMING_SNAKE_CASE      )
       SCREAMING_SNAKE_CASE       :      Union[str, Any]							=    {"input_ids": input_ids, "attention_mask": input_mask}
       SCREAMING_SNAKE_CASE       :      Dict							=    model(_SCREAMING_SNAKE_CASE      )
       self.parent.assertEqual(result.logits.shape					, (self.batch_size, self.seq_length, self.num_labels)      )






     def 				_A       (				self       :	Tuple      ):
       SCREAMING_SNAKE_CASE       :      Optional[Any]							=    self.prepare_config_and_inputs()
       (
           (
           SCREAMING_SNAKE_CASE      
       )						,  (
           SCREAMING_SNAKE_CASE      
       )						,  (
           SCREAMING_SNAKE_CASE      
       )						,  (
           SCREAMING_SNAKE_CASE      
       )						,  (
           SCREAMING_SNAKE_CASE      
       )						,  (
           SCREAMING_SNAKE_CASE      
       )						,        
       )       :      Union[str, Any]							=    config_and_inputs
       SCREAMING_SNAKE_CASE       :      List[str]							=    {"input_ids": input_ids, "attention_mask": input_mask}
       return config, inputs_dict





@require_tf
class  SCREAMING_SNAKE_CASE			(			lowerCAmelCase__      ,		lowerCAmelCase__      ,		unittest.TestCase					):
     '''simple docstring'''


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


     def 				_A       (				self       :	Dict      ):
       SCREAMING_SNAKE_CASE       :      Tuple							=    TFEsmModelTester(self      )
       SCREAMING_SNAKE_CASE       :      List[Any]							=    ConfigTester(self					, config_class=_SCREAMING_SNAKE_CASE					, hidden_size=37      )


     def 				_A       (				self       :	Optional[Any]      ):
       self.config_tester.run_common_tests()


     def 				_A       (				self       :	Optional[int]      ):
       SCREAMING_SNAKE_CASE       :      Tuple							=    self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_model(*_SCREAMING_SNAKE_CASE      )


     def 				_A       (				self       :	Dict      ):
       SCREAMING_SNAKE_CASE       :      List[Any]							=    self.model_tester.prepare_config_and_inputs_for_decoder()
       self.model_tester.create_and_check_model_as_decoder(*_SCREAMING_SNAKE_CASE      )


     def 				_A       (				self       :	Any      ):
       SCREAMING_SNAKE_CASE       :      Union[str, Any]							=    self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_masked_lm(*_SCREAMING_SNAKE_CASE      )


     def 				_A       (				self       :	Dict      ):
       SCREAMING_SNAKE_CASE       :      Tuple							=    self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_token_classification(*_SCREAMING_SNAKE_CASE      )


     @slow
     def 				_A       (				self       :	Tuple      ):
       for model_name in TF_ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
         SCREAMING_SNAKE_CASE       :      Optional[Any]							=    TFEsmModel.from_pretrained(_SCREAMING_SNAKE_CASE      )
         self.assertIsNotNone(_SCREAMING_SNAKE_CASE      )


     @unittest.skip("Protein models do not support embedding resizing."      )
     def 				_A       (				self       :	Dict      ):
       pass


     @unittest.skip("Protein models do not support embedding resizing."      )
     def 				_A       (				self       :	str      ):
       pass






     def 				_A       (				self       :	Union[str, Any]      ):
       SCREAMING_SNAKE_CASE						,  SCREAMING_SNAKE_CASE       :      str							=    self.model_tester.prepare_config_and_inputs_for_common()

       for model_class in self.all_model_classes:
         SCREAMING_SNAKE_CASE       :      Any							=    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
           SCREAMING_SNAKE_CASE       :      List[str]							=    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:
           SCREAMING_SNAKE_CASE       :      List[str]							=    model.get_output_embeddings()
           assert x is None
           SCREAMING_SNAKE_CASE       :      str							=    model.get_bias()
           assert name is None





@require_tf
class  SCREAMING_SNAKE_CASE			(			unittest.TestCase					):
     '''simple docstring'''


     @slow
     def 				_A       (				self       :	Dict      ):
       SCREAMING_SNAKE_CASE       :      Union[str, Any]							=    TFEsmForMaskedLM.from_pretrained("facebook/esm2_t6_8M_UR50D"      )

       SCREAMING_SNAKE_CASE       :      Dict							=    tf.constant([[0, 1, 2, 3, 4, 5]]      )
       SCREAMING_SNAKE_CASE       :      Tuple							=    model(_SCREAMING_SNAKE_CASE      )[0]
       SCREAMING_SNAKE_CASE       :      List[str]							=    [1, 6, 33]
       self.assertEqual(list(output.numpy().shape      )					, _SCREAMING_SNAKE_CASE      )
       # compare the actual values for a slice.
       SCREAMING_SNAKE_CASE       :      Tuple							=    tf.constant(
           [
               [
                   [8.921_518, -10.58_9814, -6.4_671_307],
                   [-6.3_967_156, -13.91_1377, -1.1_211_915],
                   [-7.781_247, -13.95_1557, -3.740_592],
               ]
           ]      )
       self.assertTrue(numpy.allclose(output[:, :3, :3].numpy()					, expected_slice.numpy()					, atol=1E-2      )      )






     @slow
     def 				_A       (				self       :	Union[str, Any]      ):
       SCREAMING_SNAKE_CASE       :      Any							=    TFEsmModel.from_pretrained("facebook/esm2_t6_8M_UR50D"      )

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














from ...utils import (
    OptionalDependencyNotAvailable,
    is_flax_available,
    is_torch_available,
    is_transformers_available,
)


try:
   if not (is_transformers_available() and is_torch_available()):
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   from ...utils.dummy_torch_and_transformers_objects import *  # noqa F403
else:
   from .multicontrolnet import MultiControlNetModel
   from .pipeline_controlnet import StableDiffusionControlNetPipeline
   from .pipeline_controlnet_imgaimg import StableDiffusionControlNetImgaImgPipeline
   from .pipeline_controlnet_inpaint import StableDiffusionControlNetInpaintPipeline


if is_transformers_available() and is_flax_available():
   from .pipeline_flax_controlnet import FlaxStableDiffusionControlNetPipeline






