


'''simple docstring'''







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 						UpperCamelCase__       :



     """simple docstring"""






     def __init__(			self   ,						snake_case__   ,						):

            '''simple docstring'''

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







     def a						(			self ):

            '''simple docstring'''

            _lowerCAmelCase     :				Union[str, Any]       =							ids_tensor([self.batch_size, self.seq_length]   ,						self.vocab_size )

            _lowerCAmelCase     :				int       =							None
            if self.use_input_mask:
                   _lowerCAmelCase     :				Union[str, Any]       =							random_attention_mask([self.batch_size, self.seq_length] )

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

            _lowerCAmelCase     :				Optional[Any]       =							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 ):

            '''simple docstring'''

            (
                (
                _lowerCAmelCase
            )      ,  (
                _lowerCAmelCase
            )      ,  (
                _lowerCAmelCase
            )      ,  (
                _lowerCAmelCase
            )      ,  (
                _lowerCAmelCase
            )      ,  (
                _lowerCAmelCase
            )      ,  
            )     :				List[Any]       =							self.prepare_config_and_inputs()

            _lowerCAmelCase     :				Any       =							True
            _lowerCAmelCase     :				Any       =							floats_tensor([self.batch_size, self.seq_length, self.hidden_size] )
            _lowerCAmelCase     :				str       =							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   ,						snake_case__   ,						snake_case__   ,						snake_case__   ,						snake_case__   ,						snake_case__   ,						snake_case__ ):

            '''simple docstring'''

            _lowerCAmelCase     :				Dict       =							TFEsmModel(config=snake_case__ )
            _lowerCAmelCase     :				Optional[Any]       =							{'input_ids': input_ids, 'attention_mask': input_mask}
            _lowerCAmelCase     :				Tuple       =							model(snake_case__ )

            _lowerCAmelCase     :				Tuple       =							[input_ids, input_mask]
            _lowerCAmelCase     :				int       =							model(snake_case__ )

            _lowerCAmelCase     :				List[Any]       =							model(snake_case__ )

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







     def a						(			self   ,						snake_case__   ,						snake_case__   ,						snake_case__   ,						snake_case__   ,						snake_case__   ,						snake_case__   ,						snake_case__   ,						snake_case__   ,						):

            '''simple docstring'''

            _lowerCAmelCase     :				Optional[int]       =							True

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

            _lowerCAmelCase     :				Optional[Any]       =							[input_ids, input_mask]
            _lowerCAmelCase     :				Tuple       =							model(snake_case__   ,						encoder_hidden_states=snake_case__ )

            # Also check the case where encoder outputs are not passed
            _lowerCAmelCase     :				Tuple       =							model(snake_case__   ,						attention_mask=snake_case__ )

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







     def a						(			self   ,						snake_case__   ,						snake_case__   ,						snake_case__   ,						snake_case__   ,						snake_case__   ,						snake_case__ ):

            '''simple docstring'''

            _lowerCAmelCase     :				int       =							TFEsmForMaskedLM(config=snake_case__ )
            _lowerCAmelCase     :				Optional[int]       =							model([input_ids, input_mask] )
            self.parent.assertEqual(result.logits.shape   ,						(self.batch_size, self.seq_length, self.vocab_size) )







     def a						(			self   ,						snake_case__   ,						snake_case__   ,						snake_case__   ,						snake_case__   ,						snake_case__   ,						snake_case__ ):

            '''simple docstring'''

            _lowerCAmelCase     :				str       =							self.num_labels
            _lowerCAmelCase     :				Optional[int]       =							TFEsmForTokenClassification(config=snake_case__ )
            _lowerCAmelCase     :				Union[str, Any]       =							{'input_ids': input_ids, 'attention_mask': input_mask}
            _lowerCAmelCase     :				Dict       =							model(snake_case__ )
            self.parent.assertEqual(result.logits.shape   ,						(self.batch_size, self.seq_length, self.num_labels) )







     def a						(			self ):

            '''simple docstring'''

            _lowerCAmelCase     :				Tuple       =							self.prepare_config_and_inputs()
            (
                (
                _lowerCAmelCase
            )      ,  (
                _lowerCAmelCase
            )      ,  (
                _lowerCAmelCase
            )      ,  (
                _lowerCAmelCase
            )      ,  (
                _lowerCAmelCase
            )      ,  (
                _lowerCAmelCase
            )      ,  
            )     :				List[Any]       =							config_and_inputs
            _lowerCAmelCase     :				int       =							{'input_ids': input_ids, 'attention_mask': input_mask}
            return config, inputs_dict
@require_tf
class 						UpperCamelCase__       (			SCREAMING_SNAKE_CASE_     ,       SCREAMING_SNAKE_CASE_     ,       unittest.TestCase						):



     """simple docstring"""






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







     def a						(			self ):

            '''simple docstring'''

            _lowerCAmelCase     :				Tuple       =							TFEsmModelTester(self )
            _lowerCAmelCase     :				int       =							ConfigTester(self   ,						config_class=snake_case__   ,						hidden_size=37 )







     def a						(			self ):

            '''simple docstring'''

            self.config_tester.run_common_tests()







     def a						(			self ):

            '''simple docstring'''

            _lowerCAmelCase     :				Optional[Any]       =							self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_model(*snake_case__ )







     def a						(			self ):

            '''simple docstring'''

            _lowerCAmelCase     :				List[Any]       =							self.model_tester.prepare_config_and_inputs_for_decoder()
            self.model_tester.create_and_check_model_as_decoder(*snake_case__ )







     def a						(			self ):

            '''simple docstring'''

            _lowerCAmelCase     :				Dict       =							self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_for_masked_lm(*snake_case__ )







     def a						(			self ):

            '''simple docstring'''

            _lowerCAmelCase     :				str       =							self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_for_token_classification(*snake_case__ )







     @slow
     def a						(			self ):

            '''simple docstring'''

            for model_name in TF_ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                   _lowerCAmelCase     :				Any       =							TFEsmModel.from_pretrained(snake_case__ )
                   self.assertIsNotNone(snake_case__ )







     @unittest.skip('Protein models do not support embedding resizing.' )
     def a						(			self ):

            '''simple docstring'''

            pass







     @unittest.skip('Protein models do not support embedding resizing.' )
     def a						(			self ):

            '''simple docstring'''

            pass







     def a						(			self ):

            '''simple docstring'''

            _lowerCAmelCase      ,  _lowerCAmelCase     :				Optional[int]       =							self.model_tester.prepare_config_and_inputs_for_common()

            for model_class in self.all_model_classes:
                   _lowerCAmelCase     :				Optional[Any]       =							model_class(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
                          _lowerCAmelCase     :				Any       =							model.get_bias()
                          assert isinstance(snake_case__   ,						snake_case__ )
                          for k, v in name.items():
                                 assert isinstance(snake_case__   ,						tf.Variable )
                   else:
                          _lowerCAmelCase     :				Tuple       =							model.get_output_embeddings()
                          assert x is None
                          _lowerCAmelCase     :				str       =							model.get_bias()
                          assert name is None
@require_tf
class 						UpperCamelCase__       (			unittest.TestCase						):



     """simple docstring"""






     @slow
     def a						(			self ):

            '''simple docstring'''

            _lowerCAmelCase     :				Any       =							TFEsmForMaskedLM.from_pretrained('facebook/esm2_t6_8M_UR50D' )

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

            '''simple docstring'''

            _lowerCAmelCase     :				Tuple       =							TFEsmModel.from_pretrained('facebook/esm2_t6_8M_UR50D' )

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





'''simple docstring'''







from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_flax_available,
    is_tf_available,
    is_tokenizers_available,
    is_torch_available,
    is_vision_available,
)


lowerCAmelCase     :  int											=    {
    """configuration_owlvit""": [
        """OWLVIT_PRETRAINED_CONFIG_ARCHIVE_MAP""",
        """OwlViTConfig""",
        """OwlViTOnnxConfig""",
        """OwlViTTextConfig""",
        """OwlViTVisionConfig""",
    ],
    """processing_owlvit""": ["""OwlViTProcessor"""],
}


try:
       if not is_vision_available():
              raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
       pass
else:
       lowerCAmelCase     :  Dict											=    ["""OwlViTFeatureExtractor"""]
       lowerCAmelCase     :  int											=    ["""OwlViTImageProcessor"""]

try:
       if not is_torch_available():
              raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
       pass
else:
       lowerCAmelCase     :  Optional[Any]											=    [
           """OWLVIT_PRETRAINED_MODEL_ARCHIVE_LIST""",
           """OwlViTModel""",
           """OwlViTPreTrainedModel""",
           """OwlViTTextModel""",
           """OwlViTVisionModel""",
           """OwlViTForObjectDetection""",
       ]

if TYPE_CHECKING:
       from .configuration_owlvit import (
           OWLVIT_PRETRAINED_CONFIG_ARCHIVE_MAP,
           OwlViTConfig,
           OwlViTOnnxConfig,
           OwlViTTextConfig,
           OwlViTVisionConfig,
       )
       from .processing_owlvit import OwlViTProcessor

       try:
              if not is_vision_available():
                     raise OptionalDependencyNotAvailable()
       except OptionalDependencyNotAvailable:
              pass
       else:
              from .feature_extraction_owlvit import OwlViTFeatureExtractor
              from .image_processing_owlvit import OwlViTImageProcessor

       try:
              if not is_torch_available():
                     raise OptionalDependencyNotAvailable()
       except OptionalDependencyNotAvailable:
              pass
       else:
              from .modeling_owlvit import (
                  OWLVIT_PRETRAINED_MODEL_ARCHIVE_LIST,
                  OwlViTForObjectDetection,
                  OwlViTModel,
                  OwlViTPreTrainedModel,
                  OwlViTTextModel,
                  OwlViTVisionModel,
              )

else:
       import sys

       lowerCAmelCase     :  Union[str, Any]											=    _LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)

