
"""simple docstring"""






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






			def __init__(self  :	Tuple      ,SCREAMING_SNAKE_CASE_  :	int      ,SCREAMING_SNAKE_CASE_  :	Any=13      ,SCREAMING_SNAKE_CASE_  :	Dict=7      ,SCREAMING_SNAKE_CASE_  :	Optional[Any]=False      ,SCREAMING_SNAKE_CASE_  :	List[str]=True      ,SCREAMING_SNAKE_CASE_  :	List[Any]=False      ,SCREAMING_SNAKE_CASE_  :	str=True      ,SCREAMING_SNAKE_CASE_  :	Optional[Any]=33      ,SCREAMING_SNAKE_CASE_  :	int=32      ,SCREAMING_SNAKE_CASE_  :	List[str]=5      ,SCREAMING_SNAKE_CASE_  :	List[str]=4      ,SCREAMING_SNAKE_CASE_  :	int=37      ,SCREAMING_SNAKE_CASE_  :	List[str]="gelu"      ,SCREAMING_SNAKE_CASE_  :	Dict=0.1      ,SCREAMING_SNAKE_CASE_  :	Dict=0.1      ,SCREAMING_SNAKE_CASE_  :	Union[str, Any]=512      ,SCREAMING_SNAKE_CASE_  :	Any=16      ,SCREAMING_SNAKE_CASE_  :	Optional[int]=2      ,SCREAMING_SNAKE_CASE_  :	List[Any]=0.02      ,SCREAMING_SNAKE_CASE_  :	Optional[int]=3      ,SCREAMING_SNAKE_CASE_  :	int=4      ,SCREAMING_SNAKE_CASE_  :	Any=None      ,)       ->	Dict:





				"""simple docstring"""


				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  :	Union[str, Any]			)       ->	Optional[int]:





				"""simple docstring"""


				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]			)       ->	Optional[int]:





				"""simple docstring"""


				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  :	Optional[int]      ,SCREAMING_SNAKE_CASE_  :	Optional[Any]      ,SCREAMING_SNAKE_CASE_  :	List[str]      ,SCREAMING_SNAKE_CASE_  :	str      ,SCREAMING_SNAKE_CASE_  :	Tuple      ,SCREAMING_SNAKE_CASE_  :	Optional[int]      ,SCREAMING_SNAKE_CASE_  :	Union[str, Any]			)       ->	Tuple:





				"""simple docstring"""


				lowerCAmelCase												=  EsmModel(config=SCREAMING_SNAKE_CASE_			)
				model.to(SCREAMING_SNAKE_CASE_			)
				model.eval()
				lowerCAmelCase												=  model(SCREAMING_SNAKE_CASE_      ,attention_mask=SCREAMING_SNAKE_CASE_			)
				lowerCAmelCase												=  model(SCREAMING_SNAKE_CASE_			)
				lowerCAmelCase												=  model(SCREAMING_SNAKE_CASE_			)

				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  :	str      ,SCREAMING_SNAKE_CASE_  :	Union[str, Any]      ,SCREAMING_SNAKE_CASE_  :	Union[str, Any]      ,SCREAMING_SNAKE_CASE_  :	List[str]      ,SCREAMING_SNAKE_CASE_  :	Optional[int]      ,SCREAMING_SNAKE_CASE_  :	str      ,SCREAMING_SNAKE_CASE_  :	Dict			)       ->	Tuple:





				"""simple docstring"""


				lowerCAmelCase												=  EsmForMaskedLM(config=SCREAMING_SNAKE_CASE_			)
				model.to(SCREAMING_SNAKE_CASE_			)
				model.eval()
				lowerCAmelCase												=  model(SCREAMING_SNAKE_CASE_      ,attention_mask=SCREAMING_SNAKE_CASE_      ,labels=SCREAMING_SNAKE_CASE_			)
				self.parent.assertEqual(result.logits.shape      ,(self.batch_size, self.seq_length, self.vocab_size)			)






			def 					UpperCAmelCase							(self  :	List[Any]      ,SCREAMING_SNAKE_CASE_  :	Tuple      ,SCREAMING_SNAKE_CASE_  :	str      ,SCREAMING_SNAKE_CASE_  :	Dict      ,SCREAMING_SNAKE_CASE_  :	Tuple      ,SCREAMING_SNAKE_CASE_  :	Optional[Any]      ,SCREAMING_SNAKE_CASE_  :	str			)       ->	str:





				"""simple docstring"""


				lowerCAmelCase												=  self.num_labels
				lowerCAmelCase												=  EsmForTokenClassification(config=SCREAMING_SNAKE_CASE_			)
				model.to(SCREAMING_SNAKE_CASE_			)
				model.eval()
				lowerCAmelCase												=  model(SCREAMING_SNAKE_CASE_      ,attention_mask=SCREAMING_SNAKE_CASE_      ,labels=SCREAMING_SNAKE_CASE_			)
				self.parent.assertEqual(result.logits.shape      ,(self.batch_size, self.seq_length, self.num_labels)			)






			def 					UpperCAmelCase							(self  :	Dict			)       ->	Dict:





				"""simple docstring"""


				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 			lowercase					( lowercase__   ,lowercase__   ,unittest.TestCase		):
			lowercase        =     False

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






			def 					UpperCAmelCase							(self  :	Any			)       ->	Dict:





				"""simple docstring"""


				lowerCAmelCase												=  EsmModelTester(self			)
				lowerCAmelCase												=  ConfigTester(self      ,config_class=SCREAMING_SNAKE_CASE_      ,hidden_size=37			)






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





				"""simple docstring"""


				self.config_tester.run_common_tests()






			def 					UpperCAmelCase							(self  :	Tuple			)       ->	List[str]:





				"""simple docstring"""


				lowerCAmelCase												=  self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_model(*SCREAMING_SNAKE_CASE_			)






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





				"""simple docstring"""


				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(*SCREAMING_SNAKE_CASE_			)






			def 					UpperCAmelCase							(self  :	Any			)       ->	List[str]:





				"""simple docstring"""


				lowerCAmelCase												=  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]			)       ->	int:





				"""simple docstring"""


				lowerCAmelCase												=  self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_for_token_classification(*SCREAMING_SNAKE_CASE_			)






			@slow
			def 					UpperCAmelCase							(self  :	Any			)       ->	str:





				"""simple docstring"""


				for model_name in ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
					lowerCAmelCase												=  EsmModel.from_pretrained(SCREAMING_SNAKE_CASE_			)
					self.assertIsNotNone(SCREAMING_SNAKE_CASE_			)






			def 					UpperCAmelCase							(self  :	List[Any]			)       ->	Tuple:





				"""simple docstring"""


				lowerCAmelCase												=  self.model_tester.prepare_config_and_inputs()[0]
				lowerCAmelCase												=  EsmEmbeddings(config=SCREAMING_SNAKE_CASE_			)

				lowerCAmelCase												=  torch.as_tensor([[12, 31, 13, 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(SCREAMING_SNAKE_CASE_      ,model.padding_idx			)
				self.assertEqual(position_ids.shape      ,expected_positions.shape			)
				self.assertTrue(torch.all(torch.eq(SCREAMING_SNAKE_CASE_      ,SCREAMING_SNAKE_CASE_			)			)			)






			def 					UpperCAmelCase							(self  :	Dict			)       ->	Any:





				"""simple docstring"""


				lowerCAmelCase												=  self.model_tester.prepare_config_and_inputs()[0]
				lowerCAmelCase												=  EsmEmbeddings(config=SCREAMING_SNAKE_CASE_			)

				lowerCAmelCase												=  torch.empty(2      ,4      ,30			)
				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(SCREAMING_SNAKE_CASE_			)
				self.assertEqual(position_ids.shape      ,expected_positions.shape			)
				self.assertTrue(torch.all(torch.eq(SCREAMING_SNAKE_CASE_      ,SCREAMING_SNAKE_CASE_			)			)			)






			@unittest.skip('''Esm does not support embedding resizing'''			)
			def 					UpperCAmelCase							(self  :	Optional[int]			)       ->	Dict:





				"""simple docstring"""


				pass






			@unittest.skip('''Esm does not support embedding resizing'''			)
			def 					UpperCAmelCase							(self  :	List[Any]			)       ->	Union[str, Any]:





				"""simple docstring"""


				pass






			@unittest.skip('''Will be fixed soon by reducing the size of the model used for common tests.'''			)
			def 					UpperCAmelCase							(self  :	Tuple			)       ->	Union[str, Any]:





				"""simple docstring"""


				pass





@require_torch
class 			lowercase					( lowercase__		):






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





				"""simple docstring"""


				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(SCREAMING_SNAKE_CASE_			)[0]

					lowerCAmelCase												=  33

					lowerCAmelCase												=  torch.Size((1, 6, vocab_size)			)
					self.assertEqual(output.shape      ,SCREAMING_SNAKE_CASE_			)

					lowerCAmelCase												=  torch.tensor(
					    [[[8.92_15, -10.58_98, -6.46_71], [-6.39_67, -13.91_14, -1.12_12], [-7.78_12, -13.95_16, -3.74_06]]]			)
					self.assertTrue(torch.allclose(output[:, :3, :3]      ,SCREAMING_SNAKE_CASE_      ,atol=1e-4			)			)






			@slow
			def 					UpperCAmelCase							(self  :	Union[str, Any]			)       ->	List[Any]:





				"""simple docstring"""


				with torch.no_grad():
					lowerCAmelCase												=  EsmModel.from_pretrained('''facebook/esm2_t6_8M_UR50D'''			)
					model.eval()

					lowerCAmelCase												=  torch.tensor([[0, 6, 4, 13, 5, 4, 16, 12, 11, 7, 2]]			)
					lowerCAmelCase												=  model(SCREAMING_SNAKE_CASE_			)[0]
					# compare the actual values for a slice.
					lowerCAmelCase												=  torch.tensor(
					    [[[0.14_44, 0.54_13, 0.32_48], [0.30_34, 0.00_53, 0.31_08], [0.32_28, -0.24_99, 0.34_15]]]			)
					self.assertTrue(torch.allclose(output[:, :3, :3]      ,SCREAMING_SNAKE_CASE_      ,atol=1e-4			)			)




"""simple docstring"""






import warnings

from ...utils import logging
from .image_processing_mobilevit import MobileViTImageProcessor


UpperCAmelCase        =				logging.get_logger(__name__)



class 			lowercase					( lowercase__		):






			def __init__(self  :	Dict      ,*SCREAMING_SNAKE_CASE_  :	str      ,**SCREAMING_SNAKE_CASE_  :	Optional[Any]			)       ->	None:





				"""simple docstring"""


				warnings.warn(
				    '''The class MobileViTFeatureExtractor is deprecated and will be removed in version 5 of Transformers.'''
				    ''' Please use MobileViTImageProcessor instead.'''      ,SCREAMING_SNAKE_CASE_      ,)
				super().__init__(*SCREAMING_SNAKE_CASE_      ,**SCREAMING_SNAKE_CASE_			)


