

"""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				,       snake_case				,       snake_case=13				,       snake_case=7				,       snake_case=False				,       snake_case=True				,       snake_case=False				,       snake_case=True				,       snake_case=33				,       snake_case=32				,       snake_case=5				,       snake_case=4				,       snake_case=37				,       snake_case="gelu"				,       snake_case=0.1				,       snake_case=0.1				,       snake_case=512				,       snake_case=16				,       snake_case=2				,       snake_case=0.02				,       snake_case=3				,       snake_case=4				,       snake_case=None				,       ):
											snake_case_          =   parent
											snake_case_          =   batch_size
											snake_case_          =   seq_length
											snake_case_          =   is_training
											snake_case_          =   use_input_mask
											snake_case_          =   use_token_type_ids
											snake_case_          =   use_labels
											snake_case_          =   vocab_size
											snake_case_          =   hidden_size
											snake_case_          =   num_hidden_layers
											snake_case_          =   num_attention_heads
											snake_case_          =   intermediate_size
											snake_case_          =   hidden_act
											snake_case_          =   hidden_dropout_prob
											snake_case_          =   attention_probs_dropout_prob
											snake_case_          =   max_position_embeddings
											snake_case_          =   type_vocab_size
											snake_case_          =   type_sequence_label_size
											snake_case_          =   initializer_range
											snake_case_          =   num_labels
											snake_case_          =   num_choices
											snake_case_          =   scope




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

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

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

											snake_case_          =   self.get_config()

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




				def    a							( 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    a							( self				,       snake_case				,       snake_case				,       snake_case				,       snake_case				,       snake_case				,       snake_case	):
											snake_case_          =   EsmModel(config=snake_case	)
											model.to(snake_case	)
											model.eval()
											snake_case_          =   model(snake_case				,       attention_mask=snake_case	)
											snake_case_          =   model(snake_case	)
											snake_case_          =   model(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    a							( self				,       snake_case				,       snake_case				,       snake_case				,       snake_case				,       snake_case				,       snake_case	):
											snake_case_          =   EsmForMaskedLM(config=snake_case	)
											model.to(snake_case	)
											model.eval()
											snake_case_          =   model(snake_case				,       attention_mask=snake_case				,       labels=snake_case	)
											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	):
											snake_case_          =   self.num_labels
											snake_case_          =   EsmForTokenClassification(config=snake_case	)
											model.to(snake_case	)
											model.eval()
											snake_case_          =   model(snake_case				,       attention_mask=snake_case				,       labels=snake_case	)
											self.parent.assertEqual(result.logits.shape				,       (self.batch_size, self.seq_length, self.num_labels)	)




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







@require_torch
class 			lowercase    (	lowercase_     ,  lowercase_     ,  unittest.TestCase      ):
				__SCREAMING_SNAKE_CASE :			List[Any]	      =      False

				__SCREAMING_SNAKE_CASE :			Tuple	      =      (
				    (
				        EsmForMaskedLM,
				        EsmModel,
				        EsmForSequenceClassification,
				        EsmForTokenClassification,
				    )
				    if is_torch_available()
				    else ()
				)
				__SCREAMING_SNAKE_CASE :			int	      =      ()
				__SCREAMING_SNAKE_CASE :			Dict	      =      (
				    {
				        '''feature-extraction''': EsmModel,
				        '''fill-mask''': EsmForMaskedLM,
				        '''text-classification''': EsmForSequenceClassification,
				        '''token-classification''': EsmForTokenClassification,
				        '''zero-shot''': EsmForSequenceClassification,
				    }
				    if is_torch_available()
				    else {}
				)
				__SCREAMING_SNAKE_CASE :			List[str]	      =      True




				def    a							( self	):
											snake_case_          =   EsmModelTester(self	)
											snake_case_          =   ConfigTester(self				,       config_class=snake_case				,       hidden_size=37	)




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




				def    a							( self	):
											snake_case_          =   self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_model(*snake_case	)




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




				def    a							( self	):
											snake_case_          =   self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_for_masked_lm(*snake_case	)




				def    a							( self	):
											snake_case_          =   self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_for_token_classification(*snake_case	)




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




				def    a							( self	):
											snake_case_          =   self.model_tester.prepare_config_and_inputs()[0]
											snake_case_          =   EsmEmbeddings(config=snake_case	)

											snake_case_          =   torch.as_tensor([[12, 31, 13, model.padding_idx]]	)
											snake_case_          =   torch.as_tensor(
											    [
											        [
											            0 + model.padding_idx + 1,
											            1 + model.padding_idx + 1,
											            2 + model.padding_idx + 1,
											            model.padding_idx,
											        ]
											    ]	)
											snake_case_          =   create_position_ids_from_input_ids(snake_case				,       model.padding_idx	)
											self.assertEqual(position_ids.shape				,       expected_positions.shape	)
											self.assertTrue(torch.all(torch.eq(snake_case				,       snake_case	)	)	)




				def    a							( self	):
											snake_case_          =   self.model_tester.prepare_config_and_inputs()[0]
											snake_case_          =   EsmEmbeddings(config=snake_case	)

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




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




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




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







@require_torch
class 			lowercase    (	lowercase_      ):




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

																		snake_case_          =   33

																		snake_case_          =   torch.Size((1, 6, vocab_size)	)
																		self.assertEqual(output.shape				,       snake_case	)

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




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

																		snake_case_          =   torch.tensor([[0, 6, 4, 13, 5, 4, 16, 12, 11, 7, 2]]	)
																		snake_case_          =   model(snake_case	)[0]
																		# compare the actual values for a slice.
																		snake_case_          =   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]				,       snake_case				,       atol=1e-4	)	)







from ...utils import (
    OptionalDependencyNotAvailable,
    is_torch_available,
    is_transformers_available,
    is_transformers_version,
)


try:
					if not (is_transformers_available() and is_torch_available() and is_transformers_version(""">=""", """4.25.0""")):
										raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
					from ...utils.dummy_torch_and_transformers_objects import UnCLIPImageVariationPipeline, UnCLIPPipeline
else:
					from .pipeline_unclip import UnCLIPPipeline
					from .pipeline_unclip_image_variation import UnCLIPImageVariationPipeline
					from .text_proj import UnCLIPTextProjModel




