



'''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 							a__:






			def __init__(     self					,   _UpperCAmelCase					,   _UpperCAmelCase=13					,   _UpperCAmelCase=7					,   _UpperCAmelCase=False					,   _UpperCAmelCase=True					,   _UpperCAmelCase=False					,   _UpperCAmelCase=True					,   _UpperCAmelCase=33					,   _UpperCAmelCase=32					,   _UpperCAmelCase=5					,   _UpperCAmelCase=4					,   _UpperCAmelCase=37					,   _UpperCAmelCase="gelu"					,   _UpperCAmelCase=0.1					,   _UpperCAmelCase=0.1					,   _UpperCAmelCase=512					,   _UpperCAmelCase=16					,   _UpperCAmelCase=2					,   _UpperCAmelCase=0.02					,   _UpperCAmelCase=3					,   _UpperCAmelCase=4					,   _UpperCAmelCase=None					,   )							->							List[str]:
					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 						_lowercase  (     self						)							->							str:
					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 						_lowercase  (     self						)							->							str:
					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 						_lowercase  (     self					,   _UpperCAmelCase					,   _UpperCAmelCase					,   _UpperCAmelCase					,   _UpperCAmelCase					,   _UpperCAmelCase					,   _UpperCAmelCase						)							->							Optional[Any]:
					snake_case__	     =EsmModel(config=_UpperCAmelCase						)
					model.to(_UpperCAmelCase						)
					model.eval()
					snake_case__	     =model(_UpperCAmelCase					,   attention_mask=_UpperCAmelCase						)
					snake_case__	     =model(_UpperCAmelCase						)
					snake_case__	     =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 						_lowercase  (     self					,   _UpperCAmelCase					,   _UpperCAmelCase					,   _UpperCAmelCase					,   _UpperCAmelCase					,   _UpperCAmelCase					,   _UpperCAmelCase						)							->							Any:
					snake_case__	     =EsmForMaskedLM(config=_UpperCAmelCase						)
					model.to(_UpperCAmelCase						)
					model.eval()
					snake_case__	     =model(_UpperCAmelCase					,   attention_mask=_UpperCAmelCase					,   labels=_UpperCAmelCase						)
					self.parent.assertEqual(result.logits.shape					,   (self.batch_size, self.seq_length, self.vocab_size)						)






			def 						_lowercase  (     self					,   _UpperCAmelCase					,   _UpperCAmelCase					,   _UpperCAmelCase					,   _UpperCAmelCase					,   _UpperCAmelCase					,   _UpperCAmelCase						)							->							List[str]:
					snake_case__	     =self.num_labels
					snake_case__	     =EsmForTokenClassification(config=_UpperCAmelCase						)
					model.to(_UpperCAmelCase						)
					model.eval()
					snake_case__	     =model(_UpperCAmelCase					,   attention_mask=_UpperCAmelCase					,   labels=_UpperCAmelCase						)
					self.parent.assertEqual(result.logits.shape					,   (self.batch_size, self.seq_length, self.num_labels)						)






			def 						_lowercase  (     self						)							->							List[Any]:
					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 							a__( snake_case__   ,      snake_case__   ,      unittest.TestCase     ):
			a_	: List[Any]									=      False

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






			def 						_lowercase  (     self						)							->							Any:
					snake_case__	     =EsmModelTester(self						)
					snake_case__	     =ConfigTester(self					,   config_class=_UpperCAmelCase					,   hidden_size=37						)






			def 						_lowercase  (     self						)							->							List[str]:
					self.config_tester.run_common_tests()






			def 						_lowercase  (     self						)							->							List[Any]:
					snake_case__	     =self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_model(*_UpperCAmelCase						)






			def 						_lowercase  (     self						)							->							List[str]:
					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(*_UpperCAmelCase						)






			def 						_lowercase  (     self						)							->							Optional[int]:
					snake_case__	     =self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_for_masked_lm(*_UpperCAmelCase						)






			def 						_lowercase  (     self						)							->							str:
					snake_case__	     =self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_for_token_classification(*_UpperCAmelCase						)






			@slow
			def 						_lowercase  (     self						)							->							Optional[int]:
					for model_name in ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
							snake_case__	     =EsmModel.from_pretrained(_UpperCAmelCase						)
							self.assertIsNotNone(_UpperCAmelCase						)






			def 						_lowercase  (     self						)							->							str:
					snake_case__	     =self.model_tester.prepare_config_and_inputs()[0]
					snake_case__	     =EsmEmbeddings(config=_UpperCAmelCase						)

					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(_UpperCAmelCase					,   model.padding_idx						)
					self.assertEqual(position_ids.shape					,   expected_positions.shape						)
					self.assertTrue(torch.all(torch.eq(_UpperCAmelCase					,   _UpperCAmelCase						)						)						)






			def 						_lowercase  (     self						)							->							str:
					snake_case__	     =self.model_tester.prepare_config_and_inputs()[0]
					snake_case__	     =EsmEmbeddings(config=_UpperCAmelCase						)

					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(_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 						_lowercase  (     self						)							->							Union[str, Any]:
					pass






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






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





@require_torch
class 							a__( snake_case__     ):






			@slow
			def 						_lowercase  (     self						)							->							Optional[int]:
					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(_UpperCAmelCase						)[0]

							snake_case__	     =33

							snake_case__	     =torch.Size((1, 6, vocab_size)						)
							self.assertEqual(output.shape					,   _UpperCAmelCase						)

							snake_case__	     =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]					,   _UpperCAmelCase					,   atol=1E-4						)						)






			@slow
			def 						_lowercase  (     self						)							->							List[str]:
					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(_UpperCAmelCase						)[0]
							# compare the actual values for a slice.
							snake_case__	     =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]					,   _UpperCAmelCase					,   atol=1E-4						)						)









'''simple docstring'''

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




