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 ,							lowercase_ ,							lowercase_=13 ,							lowercase_=7 ,							lowercase_=False ,							lowercase_=True ,							lowercase_=False ,							lowercase_=True ,							lowercase_=33 ,							lowercase_=32 ,							lowercase_=5 ,							lowercase_=4 ,							lowercase_=37 ,							lowercase_="gelu" ,							lowercase_=0.1 ,							lowercase_=0.1 ,							lowercase_=512 ,							lowercase_=16 ,							lowercase_=2 ,							lowercase_=0.02 ,							lowercase_=3 ,							lowercase_=4 ,							lowercase_=None ,							):
													_snake_case		:	Union[str, Any]							=       parent
													_snake_case		:	Any							=       batch_size
													_snake_case		:	Dict							=       seq_length
													_snake_case		:	List[Any]							=       is_training
													_snake_case		:	List[Any]							=       use_input_mask
													_snake_case		:	str							=       use_token_type_ids
													_snake_case		:	Any							=       use_labels
													_snake_case		:	Optional[int]							=       vocab_size
													_snake_case		:	Dict							=       hidden_size
													_snake_case		:	Dict							=       num_hidden_layers
													_snake_case		:	Optional[int]							=       num_attention_heads
													_snake_case		:	str							=       intermediate_size
													_snake_case		:	str							=       hidden_act
													_snake_case		:	Optional[int]							=       hidden_dropout_prob
													_snake_case		:	Union[str, Any]							=       attention_probs_dropout_prob
													_snake_case		:	Any							=       max_position_embeddings
													_snake_case		:	Optional[Any]							=       type_vocab_size
													_snake_case		:	List[Any]							=       type_sequence_label_size
													_snake_case		:	Union[str, Any]							=       initializer_range
													_snake_case		:	Union[str, Any]							=       num_labels
													_snake_case		:	List[str]							=       num_choices
													_snake_case		:	List[Any]							=       scope



							def 			UpperCamelCase			(		self							):
													_snake_case		:	List[str]							=       ids_tensor([self.batch_size, self.seq_length] ,							self.vocab_size							)

													_snake_case		:	int							=       None
													if self.use_input_mask:
																			_snake_case		:	Tuple							=       random_attention_mask([self.batch_size, self.seq_length]							)

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

													_snake_case		:	Optional[int]							=       self.get_config()

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



							def 			UpperCamelCase			(		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 			UpperCamelCase			(		self ,							lowercase_ ,							lowercase_ ,							lowercase_ ,							lowercase_ ,							lowercase_ ,							lowercase_							):
													_snake_case		:	List[Any]							=       EsmModel(config=_lowerCAmelCase							)
													model.to(_lowerCAmelCase							)
													model.eval()
													_snake_case		:	Union[str, Any]							=       model(_lowerCAmelCase ,							attention_mask=_lowerCAmelCase							)
													_snake_case		:	Tuple							=       model(_lowerCAmelCase							)
													_snake_case		:	int							=       model(_lowerCAmelCase							)

													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 ,							lowercase_ ,							lowercase_ ,							lowercase_ ,							lowercase_ ,							lowercase_ ,							lowercase_							):
													_snake_case		:	List[Any]							=       EsmForMaskedLM(config=_lowerCAmelCase							)
													model.to(_lowerCAmelCase							)
													model.eval()
													_snake_case		:	Optional[Any]							=       model(_lowerCAmelCase ,							attention_mask=_lowerCAmelCase ,							labels=_lowerCAmelCase							)
													self.parent.assertEqual(result.logits.shape ,							(self.batch_size, self.seq_length, self.vocab_size)							)



							def 			UpperCamelCase			(		self ,							lowercase_ ,							lowercase_ ,							lowercase_ ,							lowercase_ ,							lowercase_ ,							lowercase_							):
													_snake_case		:	Any							=       self.num_labels
													_snake_case		:	List[Any]							=       EsmForTokenClassification(config=_lowerCAmelCase							)
													model.to(_lowerCAmelCase							)
													model.eval()
													_snake_case		:	int							=       model(_lowerCAmelCase ,							attention_mask=_lowerCAmelCase ,							labels=_lowerCAmelCase							)
													self.parent.assertEqual(result.logits.shape ,							(self.batch_size, self.seq_length, self.num_labels)							)



							def 			UpperCamelCase			(		self							):
													_snake_case		:	str							=       self.prepare_config_and_inputs()
													(
													    _snake_case
													)		:	int							=       config_and_inputs
													_snake_case		:	Tuple							=       {'''input_ids''': input_ids, '''attention_mask''': input_mask}
													return config, inputs_dict




@require_torch
class   lowercase_       (	__lowercase		,					__lowercase		,					unittest.TestCase						):
							_lowerCamelCase		=						False

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



							def 			UpperCamelCase			(		self							):
													_snake_case		:	List[str]							=       EsmModelTester(self							)
													_snake_case		:	List[str]							=       ConfigTester(self ,							config_class=_lowerCAmelCase ,							hidden_size=37							)



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



							def 			UpperCamelCase			(		self							):
													_snake_case		:	Dict							=       self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_model(*_lowerCAmelCase							)



							def 			UpperCamelCase			(		self							):
													_snake_case		:	Optional[int]							=       self.model_tester.prepare_config_and_inputs()
													for type in ["absolute", "relative_key", "relative_key_query"]:
																			_snake_case		:	Any							=       type
																			self.model_tester.create_and_check_model(*_lowerCAmelCase							)



							def 			UpperCamelCase			(		self							):
													_snake_case		:	Optional[int]							=       self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_for_masked_lm(*_lowerCAmelCase							)



							def 			UpperCamelCase			(		self							):
													_snake_case		:	Optional[int]							=       self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_for_token_classification(*_lowerCAmelCase							)



							@slow
							def 			UpperCamelCase			(		self							):




													for model_name in ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																			_snake_case		:	Optional[int]							=       EsmModel.from_pretrained(_lowerCAmelCase							)
																			self.assertIsNotNone(_lowerCAmelCase							)



							def 			UpperCamelCase			(		self							):
													_snake_case		:	Union[str, Any]							=       self.model_tester.prepare_config_and_inputs()[0]
													_snake_case		:	int							=       EsmEmbeddings(config=_lowerCAmelCase							)

													_snake_case		:	Union[str, Any]							=       torch.as_tensor([[12, 31, 13, model.padding_idx]]							)
													_snake_case		:	Optional[Any]							=       torch.as_tensor(
													    [
													        [
													            0 + model.padding_idx + 1,
													            1 + model.padding_idx + 1,
													            2 + model.padding_idx + 1,
													            model.padding_idx,
													        ]
													    ]							)
													_snake_case		:	Any							=       create_position_ids_from_input_ids(_lowerCAmelCase ,							model.padding_idx							)
													self.assertEqual(position_ids.shape ,							expected_positions.shape							)
													self.assertTrue(torch.all(torch.eq(_lowerCAmelCase ,							_lowerCAmelCase							)							)							)



							def 			UpperCamelCase			(		self							):
													_snake_case		:	Tuple							=       self.model_tester.prepare_config_and_inputs()[0]
													_snake_case		:	Dict							=       EsmEmbeddings(config=_lowerCAmelCase							)

													_snake_case		:	Optional[int]							=       torch.empty(2 ,							4 ,							30							)
													_snake_case		:	List[str]							=       [
													    0 + embeddings.padding_idx + 1,
													    1 + embeddings.padding_idx + 1,
													    2 + embeddings.padding_idx + 1,
													    3 + embeddings.padding_idx + 1,
													]
													_snake_case		:	str							=       torch.as_tensor([expected_single_positions, expected_single_positions]							)
													_snake_case		:	Optional[Any]							=       embeddings.create_position_ids_from_inputs_embeds(_lowerCAmelCase							)
													self.assertEqual(position_ids.shape ,							expected_positions.shape							)
													self.assertTrue(torch.all(torch.eq(_lowerCAmelCase ,							_lowerCAmelCase							)							)							)



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



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



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




@require_torch
class   lowercase_       (	__lowercase						):



							@slow
							def 			UpperCamelCase			(		self							):




													with torch.no_grad():
																			_snake_case		:	Optional[int]							=       EsmForMaskedLM.from_pretrained("facebook/esm2_t6_8M_UR50D"							)
																			model.eval()
																			_snake_case		:	Tuple							=       torch.tensor([[0, 1, 2, 3, 4, 5]]							)
																			_snake_case		:	str							=       model(_lowerCAmelCase							)[0]

																			_snake_case		:	Union[str, Any]							=       33

																			_snake_case		:	Dict							=       torch.Size((1, 6, vocab_size)							)
																			self.assertEqual(output.shape ,							_lowerCAmelCase							)

																			_snake_case		:	Optional[int]							=       torch.tensor(
																			    [[[8.9_215, -10.5_898, -6.4_671], [-6.3_967, -13.9_114, -1.1_212], [-7.7_812, -13.9_516, -3.7_406]]]							)
																			self.assertTrue(torch.allclose(output[:, :3, :3] ,							_lowerCAmelCase ,							atol=1e-4							)							)



							@slow
							def 			UpperCamelCase			(		self							):




													with torch.no_grad():
																			_snake_case		:	Dict							=       EsmModel.from_pretrained("facebook/esm2_t6_8M_UR50D"							)
																			model.eval()

																			_snake_case		:	Dict							=       torch.tensor([[0, 6, 4, 13, 5, 4, 16, 12, 11, 7, 2]]							)
																			_snake_case		:	Union[str, Any]							=       model(_lowerCAmelCase							)[0]
																			# compare the actual values for a slice.
																			_snake_case		:	Optional[int]							=       torch.tensor(
																			    [[[0.1_444, 0.5_413, 0.3_248], [0.3_034, 0.0_053, 0.3_108], [0.3_228, -0.2_499, 0.3_415]]]							)
																			self.assertTrue(torch.allclose(output[:, :3, :3] ,							_lowerCAmelCase ,							atol=1e-4							)							)






from collections import OrderedDict
from typing import Mapping

from packaging import version

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig
from ...utils import logging


lowerCAmelCase							:       Optional[int]					 =	logging.get_logger(__name__)

lowerCAmelCase							:       Tuple					 =	{
    """facebook/data2vec-vision-base-ft""": (
        """https://huggingface.co/facebook/data2vec-vision-base-ft/resolve/main/config.json"""
    ),
}





class 				a			(      __lowercase			):
				SCREAMING_SNAKE_CASE__     :   Union[str, Any]		      =     '''data2vec-vision'''






				def __init__(  self			, _lowerCAmelCase=768			, _lowerCAmelCase=12			, _lowerCAmelCase=12			, _lowerCAmelCase=3072			, _lowerCAmelCase="gelu"			, _lowerCAmelCase=0.0			, _lowerCAmelCase=0.0			, _lowerCAmelCase=0.02			, _lowerCAmelCase=1e-12			, _lowerCAmelCase=224			, _lowerCAmelCase=16			, _lowerCAmelCase=3			, _lowerCAmelCase=False			, _lowerCAmelCase=False			, _lowerCAmelCase=False			, _lowerCAmelCase=False			, _lowerCAmelCase=0.1			, _lowerCAmelCase=0.1			, _lowerCAmelCase=True			, _lowerCAmelCase=[3, 5, 7, 11]			, _lowerCAmelCase=[1, 2, 3, 6]			, _lowerCAmelCase=True			, _lowerCAmelCase=0.4			, _lowerCAmelCase=256			, _lowerCAmelCase=1			, _lowerCAmelCase=False			, _lowerCAmelCase=255			, **_lowerCAmelCase			, ):

										"""simple docstring"""




										super().__init__(**_lowerCAmelCase					)

										__SCREAMING_SNAKE_CASE:				str           =   hidden_size
										__SCREAMING_SNAKE_CASE:				Tuple           =   num_hidden_layers
										__SCREAMING_SNAKE_CASE:				Tuple           =   num_attention_heads
										__SCREAMING_SNAKE_CASE:				int           =   intermediate_size
										__SCREAMING_SNAKE_CASE:				Optional[Any]           =   hidden_act
										__SCREAMING_SNAKE_CASE:				int           =   hidden_dropout_prob
										__SCREAMING_SNAKE_CASE:				str           =   attention_probs_dropout_prob
										__SCREAMING_SNAKE_CASE:				Any           =   initializer_range
										__SCREAMING_SNAKE_CASE:				Union[str, Any]           =   layer_norm_eps

										__SCREAMING_SNAKE_CASE:				Optional[Any]           =   image_size
										__SCREAMING_SNAKE_CASE:				List[str]           =   patch_size
										__SCREAMING_SNAKE_CASE:				Optional[int]           =   num_channels
										__SCREAMING_SNAKE_CASE:				List[Any]           =   use_mask_token
										__SCREAMING_SNAKE_CASE:				Tuple           =   use_absolute_position_embeddings
										__SCREAMING_SNAKE_CASE:				Any           =   use_relative_position_bias
										__SCREAMING_SNAKE_CASE:				Dict           =   use_shared_relative_position_bias
										__SCREAMING_SNAKE_CASE:				Any           =   layer_scale_init_value
										__SCREAMING_SNAKE_CASE:				List[Any]           =   drop_path_rate
										__SCREAMING_SNAKE_CASE:				int           =   use_mean_pooling
										# decode head attributes (semantic segmentation)
										__SCREAMING_SNAKE_CASE:				Union[str, Any]           =   out_indices
										__SCREAMING_SNAKE_CASE:				Optional[int]           =   pool_scales
										# auxiliary head attributes (semantic segmentation)
										__SCREAMING_SNAKE_CASE:				Optional[int]           =   use_auxiliary_head
										__SCREAMING_SNAKE_CASE:				Optional[Any]           =   auxiliary_loss_weight
										__SCREAMING_SNAKE_CASE:				Optional[Any]           =   auxiliary_channels
										__SCREAMING_SNAKE_CASE:				Optional[int]           =   auxiliary_num_convs
										__SCREAMING_SNAKE_CASE:				Any           =   auxiliary_concat_input
										__SCREAMING_SNAKE_CASE:				List[str]           =   semantic_loss_ignore_index





class 				a			(      __lowercase			):
				SCREAMING_SNAKE_CASE__     :   List[Any]		      =     version.parse('''1.11'''			)






				@property
				def 			snake_case_    (  self					):

										"""simple docstring"""




										return OrderedDict(
										    [
										        ('''pixel_values''', {0: '''batch''', 1: '''num_channels''', 2: '''height''', 3: '''width'''}),
										    ]					)






				@property
				def 			snake_case_    (  self					):

										"""simple docstring"""




										return 1e-4



