



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 ,						a__ ,						a__=1_3 ,						a__=7 ,						a__=False ,						a__=True ,						a__=False ,						a__=True ,						a__=3_3 ,						a__=3_2 ,						a__=5 ,						a__=4 ,						a__=3_7 ,						a__="gelu" ,						a__=0.1 ,						a__=0.1 ,						a__=5_1_2 ,						a__=1_6 ,						a__=2 ,						a__=0.0_2 ,						a__=3 ,						a__=4 ,						a__=None ,						):
         A__		 =					parent
         A__		 =					batch_size
         A__		 =					seq_length
         A__		 =					is_training
         A__		 =					use_input_mask
         A__		 =					use_token_type_ids
         A__		 =					use_labels
         A__		 =					vocab_size
         A__		 =					hidden_size
         A__		 =					num_hidden_layers
         A__		 =					num_attention_heads
         A__		 =					intermediate_size
         A__		 =					hidden_act
         A__		 =					hidden_dropout_prob
         A__		 =					attention_probs_dropout_prob
         A__		 =					max_position_embeddings
         A__		 =					type_vocab_size
         A__		 =					type_sequence_label_size
         A__		 =					initializer_range
         A__		 =					num_labels
         A__		 =					num_choices
         A__		 =					scope




     def 		snake_case_		(				self):
         A__		 =					ids_tensor([self.batch_size, self.seq_length] ,						self.vocab_size)

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

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

         A__		 =					self.get_config()

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




     def 		snake_case_		(				self):
         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 		snake_case_		(				self ,						a__ ,						a__ ,						a__ ,						a__ ,						a__ ,						a__):
         A__		 =					EsmModel(config=a__)
         model.to(a__)
         model.eval()
         A__		 =					model(a__ ,						attention_mask=a__)
         A__		 =					model(a__)
         A__		 =					model(a__)

         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 		snake_case_		(				self ,						a__ ,						a__ ,						a__ ,						a__ ,						a__ ,						a__):
         A__		 =					EsmForMaskedLM(config=a__)
         model.to(a__)
         model.eval()
         A__		 =					model(a__ ,						attention_mask=a__ ,						labels=a__)
         self.parent.assertEqual(result.logits.shape ,						(self.batch_size, self.seq_length, self.vocab_size))




     def 		snake_case_		(				self ,						a__ ,						a__ ,						a__ ,						a__ ,						a__ ,						a__):
         A__		 =					self.num_labels
         A__		 =					EsmForTokenClassification(config=a__)
         model.to(a__)
         model.eval()
         A__		 =					model(a__ ,						attention_mask=a__ ,						labels=a__)
         self.parent.assertEqual(result.logits.shape ,						(self.batch_size, self.seq_length, self.num_labels))




     def 		snake_case_		(				self):
         A__		 =					self.prepare_config_and_inputs()
         (
             (
             A__	
         ) ,	(
             A__	
         ) ,	(
             A__	
         ) ,	(
             A__	
         ) ,	(
             A__	
         ) ,	(
             A__	
         ) ,		
         )		 =					config_and_inputs
         A__		 =					{'''input_ids''': input_ids, '''attention_mask''': input_mask}
         return config, inputs_dict







@require_torch
class     _UpperCAmelCase			(       A__			,      A__			,      unittest.TestCase     ):
     UpperCamelCase__		    =				False

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




     def 		snake_case_		(				self):
         A__		 =					EsmModelTester(self)
         A__		 =					ConfigTester(self ,						config_class=a__ ,						hidden_size=3_7)




     def 		snake_case_		(				self):
         self.config_tester.run_common_tests()




     def 		snake_case_		(				self):
         A__		 =					self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_model(*a__)




     def 		snake_case_		(				self):
         A__		 =					self.model_tester.prepare_config_and_inputs()
         for type in ["absolute", "relative_key", "relative_key_query"]:
             A__		 =					type
             self.model_tester.create_and_check_model(*a__)




     def 		snake_case_		(				self):
         A__		 =					self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_masked_lm(*a__)




     def 		snake_case_		(				self):
         A__		 =					self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_token_classification(*a__)




     @slow
     def 		snake_case_		(				self):
         for model_name in ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
             A__		 =					EsmModel.from_pretrained(a__)
             self.assertIsNotNone(a__)




     def 		snake_case_		(				self):
         A__		 =					self.model_tester.prepare_config_and_inputs()[0]
         A__		 =					EsmEmbeddings(config=a__)

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




     def 		snake_case_		(				self):
         A__		 =					self.model_tester.prepare_config_and_inputs()[0]
         A__		 =					EsmEmbeddings(config=a__)

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




     @unittest.skip('''Esm does not support embedding resizing''')
     def 		snake_case_		(				self):
         pass




     @unittest.skip('''Esm does not support embedding resizing''')
     def 		snake_case_		(				self):
         pass




     @unittest.skip('''Will be fixed soon by reducing the size of the model used for common tests.''')
     def 		snake_case_		(				self):
         pass







@require_torch
class     _UpperCAmelCase			(       A__     ):




     @slow
     def 		snake_case_		(				self):
         with torch.no_grad():
             A__		 =					EsmForMaskedLM.from_pretrained('''facebook/esm2_t6_8M_UR50D''')
             model.eval()
             A__		 =					torch.tensor([[0, 1, 2, 3, 4, 5]])
             A__		 =					model(a__)[0]

             A__		 =					3_3

             A__		 =					torch.Size((1, 6, vocab_size))
             self.assertEqual(output.shape ,						a__)

             A__		 =					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] ,						a__ ,						atol=1e-4))




     @slow
     def 		snake_case_		(				self):
         with torch.no_grad():
             A__		 =					EsmModel.from_pretrained('''facebook/esm2_t6_8M_UR50D''')
             model.eval()

             A__		 =					torch.tensor([[0, 6, 4, 1_3, 5, 4, 1_6, 1_2, 1_1, 7, 2]])
             A__		 =					model(a__)[0]
             # compare the actual values for a slice.
             A__		 =					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] ,						a__ ,						atol=1e-4))








from typing import TYPE_CHECKING

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


_lowercase						  =		{
    "configuration_mobilevit": ["MOBILEVIT_PRETRAINED_CONFIG_ARCHIVE_MAP", "MobileViTConfig", "MobileViTOnnxConfig"],
}

try:
   if not is_vision_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   _lowercase						  =		["MobileViTFeatureExtractor"]
   _lowercase						  =		["MobileViTImageProcessor"]

try:
   if not is_torch_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   _lowercase						  =		[
       "MOBILEVIT_PRETRAINED_MODEL_ARCHIVE_LIST",
       "MobileViTForImageClassification",
       "MobileViTForSemanticSegmentation",
       "MobileViTModel",
       "MobileViTPreTrainedModel",
   ]

try:
   if not is_tf_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   _lowercase						  =		[
       "TF_MOBILEVIT_PRETRAINED_MODEL_ARCHIVE_LIST",
       "TFMobileViTForImageClassification",
       "TFMobileViTForSemanticSegmentation",
       "TFMobileViTModel",
       "TFMobileViTPreTrainedModel",
   ]

if TYPE_CHECKING:
   from .configuration_mobilevit import MOBILEVIT_PRETRAINED_CONFIG_ARCHIVE_MAP, MobileViTConfig, MobileViTOnnxConfig

   try:
      if not is_vision_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .feature_extraction_mobilevit import MobileViTFeatureExtractor
      from .image_processing_mobilevit import MobileViTImageProcessor

   try:
      if not is_torch_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .modeling_mobilevit import (
          MOBILEVIT_PRETRAINED_MODEL_ARCHIVE_LIST,
          MobileViTForImageClassification,
          MobileViTForSemanticSegmentation,
          MobileViTModel,
          MobileViTPreTrainedModel,
      )

   try:
      if not is_tf_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .modeling_tf_mobilevit import (
          TF_MOBILEVIT_PRETRAINED_MODEL_ARCHIVE_LIST,
          TFMobileViTForImageClassification,
          TFMobileViTForSemanticSegmentation,
          TFMobileViTModel,
          TFMobileViTPreTrainedModel,
      )


else:
   import sys

   _lowercase						  =		_LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)
