






import unittest

from transformers import EsmConfig, is_torch_available
from transformers.testing_utils import TestCasePlus, require_torch, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
 import torch

 from transformers import EsmForMaskedLM, EsmForSequenceClassification, EsmForTokenClassification, EsmModel
 from transformers.models.esm.modeling_esm import (
     ESM_PRETRAINED_MODEL_ARCHIVE_LIST,
     EsmEmbeddings,
     create_position_ids_from_input_ids,
 )



class 				_UpperCamelCase      :


      def __init__(						self       ,    __UpperCamelCase       ,    __UpperCamelCase=1_3       ,    __UpperCamelCase=7       ,    __UpperCamelCase=False       ,    __UpperCamelCase=True       ,    __UpperCamelCase=False       ,    __UpperCamelCase=True       ,    __UpperCamelCase=3_3       ,    __UpperCamelCase=3_2       ,    __UpperCamelCase=5       ,    __UpperCamelCase=4       ,    __UpperCamelCase=3_7       ,    __UpperCamelCase="gelu"       ,    __UpperCamelCase=0.1       ,    __UpperCamelCase=0.1       ,    __UpperCamelCase=5_1_2       ,    __UpperCamelCase=1_6       ,    __UpperCamelCase=2       ,    __UpperCamelCase=0.0_2       ,    __UpperCamelCase=3       ,    __UpperCamelCase=4       ,    __UpperCamelCase=None       ,    )->						List[str]:
           __lowerCAmelCase             =       parent
           __lowerCAmelCase             =       batch_size
           __lowerCAmelCase             =       seq_length
           __lowerCAmelCase             =       is_training
           __lowerCAmelCase             =       use_input_mask
           __lowerCAmelCase             =       use_token_type_ids
           __lowerCAmelCase             =       use_labels
           __lowerCAmelCase             =       vocab_size
           __lowerCAmelCase             =       hidden_size
           __lowerCAmelCase             =       num_hidden_layers
           __lowerCAmelCase             =       num_attention_heads
           __lowerCAmelCase             =       intermediate_size
           __lowerCAmelCase             =       hidden_act
           __lowerCAmelCase             =       hidden_dropout_prob
           __lowerCAmelCase             =       attention_probs_dropout_prob
           __lowerCAmelCase             =       max_position_embeddings
           __lowerCAmelCase             =       type_vocab_size
           __lowerCAmelCase             =       type_sequence_label_size
           __lowerCAmelCase             =       initializer_range
           __lowerCAmelCase             =       num_labels
           __lowerCAmelCase             =       num_choices
           __lowerCAmelCase             =       scope


      def 			__UpperCAmelCase		(						self						)->						Dict:
           __lowerCAmelCase             =       ids_tensor([self.batch_size, self.seq_length]       ,    self.vocab_size						)

           __lowerCAmelCase             =       None
           if self.use_input_mask:
                __lowerCAmelCase             =       random_attention_mask([self.batch_size, self.seq_length]						)

           __lowerCAmelCase             =       None
           __lowerCAmelCase             =       None
           __lowerCAmelCase             =       None
           if self.use_labels:
                __lowerCAmelCase             =       ids_tensor([self.batch_size]       ,    self.type_sequence_label_size						)
                __lowerCAmelCase             =       ids_tensor([self.batch_size, self.seq_length]       ,    self.num_labels						)
                __lowerCAmelCase             =       ids_tensor([self.batch_size]       ,    self.num_choices						)

           __lowerCAmelCase             =       self.get_config()

           return config, input_ids, input_mask, sequence_labels, token_labels, choice_labels


      def 			__UpperCAmelCase		(						self						)->						Optional[Any]:
           return EsmConfig(
               vocab_size=self.vocab_size       ,    hidden_size=self.hidden_size       ,    pad_token_id=1       ,    num_hidden_layers=self.num_hidden_layers       ,    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       ,    )


      def 			__UpperCAmelCase		(						self       ,    __UpperCamelCase       ,    __UpperCamelCase       ,    __UpperCamelCase       ,    __UpperCamelCase       ,    __UpperCamelCase       ,    __UpperCamelCase						)->						Optional[int]:
           __lowerCAmelCase             =       EsmModel(config=__UpperCamelCase						)
           model.to(__UpperCamelCase						)
           model.eval()
           __lowerCAmelCase             =       model(__UpperCamelCase       ,    attention_mask=__UpperCamelCase						)
           __lowerCAmelCase             =       model(__UpperCamelCase						)
           __lowerCAmelCase             =       model(__UpperCamelCase						)

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


      def 			__UpperCAmelCase		(						self       ,    __UpperCamelCase       ,    __UpperCamelCase       ,    __UpperCamelCase       ,    __UpperCamelCase       ,    __UpperCamelCase       ,    __UpperCamelCase						)->						int:
           __lowerCAmelCase             =       EsmForMaskedLM(config=__UpperCamelCase						)
           model.to(__UpperCamelCase						)
           model.eval()
           __lowerCAmelCase             =       model(__UpperCamelCase       ,    attention_mask=__UpperCamelCase       ,    labels=__UpperCamelCase						)
           self.parent.assertEqual(result.logits.shape       ,    (self.batch_size, self.seq_length, self.vocab_size)						)


      def 			__UpperCAmelCase		(						self       ,    __UpperCamelCase       ,    __UpperCamelCase       ,    __UpperCamelCase       ,    __UpperCamelCase       ,    __UpperCamelCase       ,    __UpperCamelCase						)->						List[Any]:
           __lowerCAmelCase             =       self.num_labels
           __lowerCAmelCase             =       EsmForTokenClassification(config=__UpperCamelCase						)
           model.to(__UpperCamelCase						)
           model.eval()
           __lowerCAmelCase             =       model(__UpperCamelCase       ,    attention_mask=__UpperCamelCase       ,    labels=__UpperCamelCase						)
           self.parent.assertEqual(result.logits.shape       ,    (self.batch_size, self.seq_length, self.num_labels)						)


      def 			__UpperCAmelCase		(						self						)->						Optional[int]:
           __lowerCAmelCase             =       self.prepare_config_and_inputs()
           (
               (
               __lowerCAmelCase
           )      ,			(
               __lowerCAmelCase
           )      ,			(
               __lowerCAmelCase
           )      ,			(
               __lowerCAmelCase
           )      ,			(
               __lowerCAmelCase
           )      ,			(
               __lowerCAmelCase
           )      ,			
           )             =       config_and_inputs
           __lowerCAmelCase             =       {"input_ids": input_ids, "attention_mask": input_mask}
           return config, inputs_dict



@require_torch
class 				_UpperCamelCase      (a_				,		a_				,		unittest.TestCase      ):
      snake_case_	=				False

      snake_case_	=				(
          (
              EsmForMaskedLM,
              EsmModel,
              EsmForSequenceClassification,
              EsmForTokenClassification,
          )
          if is_torch_available()
          else ()
      )
      snake_case_	=				()
      snake_case_	=				(
          {
              """feature-extraction""": EsmModel,
              """fill-mask""": EsmForMaskedLM,
              """text-classification""": EsmForSequenceClassification,
              """token-classification""": EsmForTokenClassification,
              """zero-shot""": EsmForSequenceClassification,
          }
          if is_torch_available()
          else {}
      )
      snake_case_	=				True


      def 			__UpperCAmelCase		(						self						)->						Union[str, Any]:
           __lowerCAmelCase             =       EsmModelTester(self						)
           __lowerCAmelCase             =       ConfigTester(self       ,    config_class=__UpperCamelCase       ,    hidden_size=3_7						)


      def 			__UpperCAmelCase		(						self						)->						Dict:
           self.config_tester.run_common_tests()


      def 			__UpperCAmelCase		(						self						)->						Union[str, Any]:
           __lowerCAmelCase             =       self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_model(*__UpperCamelCase						)


      def 			__UpperCAmelCase		(						self						)->						List[str]:
           __lowerCAmelCase             =       self.model_tester.prepare_config_and_inputs()
           for type in ["absolute", "relative_key", "relative_key_query"]:
                __lowerCAmelCase             =       type
                self.model_tester.create_and_check_model(*__UpperCamelCase						)


      def 			__UpperCAmelCase		(						self						)->						List[str]:
           __lowerCAmelCase             =       self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_masked_lm(*__UpperCamelCase						)


      def 			__UpperCAmelCase		(						self						)->						Dict:
           __lowerCAmelCase             =       self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_token_classification(*__UpperCamelCase						)


      @slow
      def 			__UpperCAmelCase		(						self						)->						List[str]:
           for model_name in ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                __lowerCAmelCase             =       EsmModel.from_pretrained(__UpperCamelCase						)
                self.assertIsNotNone(__UpperCamelCase						)


      def 			__UpperCAmelCase		(						self						)->						List[str]:
           __lowerCAmelCase             =       self.model_tester.prepare_config_and_inputs()[0]
           __lowerCAmelCase             =       EsmEmbeddings(config=__UpperCamelCase						)

           __lowerCAmelCase             =       torch.as_tensor([[1_2, 3_1, 1_3, model.padding_idx]]						)
           __lowerCAmelCase             =       torch.as_tensor(
               [
                   [
                       0 + model.padding_idx + 1,
                       1 + model.padding_idx + 1,
                       2 + model.padding_idx + 1,
                       model.padding_idx,
                   ]
               ]						)
           __lowerCAmelCase             =       create_position_ids_from_input_ids(__UpperCamelCase       ,    model.padding_idx						)
           self.assertEqual(position_ids.shape       ,    expected_positions.shape						)
           self.assertTrue(torch.all(torch.eq(__UpperCamelCase       ,    __UpperCamelCase						)						)						)


      def 			__UpperCAmelCase		(						self						)->						List[str]:
           __lowerCAmelCase             =       self.model_tester.prepare_config_and_inputs()[0]
           __lowerCAmelCase             =       EsmEmbeddings(config=__UpperCamelCase						)

           __lowerCAmelCase             =       torch.empty(2       ,    4       ,    3_0						)
           __lowerCAmelCase             =       [
               0 + embeddings.padding_idx + 1,
               1 + embeddings.padding_idx + 1,
               2 + embeddings.padding_idx + 1,
               3 + embeddings.padding_idx + 1,
           ]
           __lowerCAmelCase             =       torch.as_tensor([expected_single_positions, expected_single_positions]						)
           __lowerCAmelCase             =       embeddings.create_position_ids_from_inputs_embeds(__UpperCamelCase						)
           self.assertEqual(position_ids.shape       ,    expected_positions.shape						)
           self.assertTrue(torch.all(torch.eq(__UpperCamelCase       ,    __UpperCamelCase						)						)						)


      @unittest.skip("Esm does not support embedding resizing"						)
      def 			__UpperCAmelCase		(						self						)->						Optional[int]:
           pass


      @unittest.skip("Esm does not support embedding resizing"						)
      def 			__UpperCAmelCase		(						self						)->						Any:
           pass


      @unittest.skip("Will be fixed soon by reducing the size of the model used for common tests."						)
      def 			__UpperCAmelCase		(						self						)->						Any:
           pass



@require_torch
class 				_UpperCamelCase      (a_      ):


      @slow
      def 			__UpperCAmelCase		(						self						)->						List[Any]:
           with torch.no_grad():
                __lowerCAmelCase             =       EsmForMaskedLM.from_pretrained("facebook/esm2_t6_8M_UR50D"						)
                model.eval()
                __lowerCAmelCase             =       torch.tensor([[0, 1, 2, 3, 4, 5]]						)
                __lowerCAmelCase             =       model(__UpperCamelCase						)[0]

                __lowerCAmelCase             =       3_3

                __lowerCAmelCase             =       torch.Size((1, 6, vocab_size)						)
                self.assertEqual(output.shape       ,    __UpperCamelCase						)

                __lowerCAmelCase             =       torch.tensor(
                    [[[8.9_2_1_5, -1_0.5_8_9_8, -6.4_6_7_1], [-6.3_9_6_7, -1_3.9_1_1_4, -1.1_2_1_2], [-7.7_8_1_2, -1_3.9_5_1_6, -3.7_4_0_6]]]						)
                self.assertTrue(torch.allclose(output[:, :3, :3]       ,    __UpperCamelCase       ,    atol=1e-4						)						)


      @slow
      def 			__UpperCAmelCase		(						self						)->						Tuple:
           with torch.no_grad():
                __lowerCAmelCase             =       EsmModel.from_pretrained("facebook/esm2_t6_8M_UR50D"						)
                model.eval()

                __lowerCAmelCase             =       torch.tensor([[0, 6, 4, 1_3, 5, 4, 1_6, 1_2, 1_1, 7, 2]]						)
                __lowerCAmelCase             =       model(__UpperCamelCase						)[0]
                # compare the actual values for a slice.
                __lowerCAmelCase             =       torch.tensor(
                    [[[0.1_4_4_4, 0.5_4_1_3, 0.3_2_4_8], [0.3_0_3_4, 0.0_0_5_3, 0.3_1_0_8], [0.3_2_2_8, -0.2_4_9_9, 0.3_4_1_5]]]						)
                self.assertTrue(torch.allclose(output[:, :3, :3]       ,    __UpperCamelCase       ,    atol=1e-4						)						)














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      :       List[Any]    =						{
    '''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      :       List[Any]    =						['''OwlViTFeatureExtractor''']
 lowerCamelCase      :       str    =						['''OwlViTImageProcessor''']

try:
 if not is_torch_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 lowerCamelCase      :       List[str]    =						[
     '''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      :       List[str]    =						_LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)






