"""simple docstring"""


from __future__ import annotations

import unittest

from transformers import EsmConfig, is_tf_available
from transformers.testing_utils import require_tf, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
	import numpy
	import tensorflow as tf

	from transformers.models.esm.modeling_tf_esm import (
	    TF_ESM_PRETRAINED_MODEL_ARCHIVE_LIST,
	    TFEsmForMaskedLM,
	    TFEsmForSequenceClassification,
	    TFEsmForTokenClassification,
	    TFEsmModel,
	)






class   __UpperCAmelCase		:

				def __init__( self      ,					_lowerCamelCase      ,					):
										lowerCamelCase__   =parent
										lowerCamelCase__   =13
										lowerCamelCase__   =7
										lowerCamelCase__   =True
										lowerCamelCase__   =True
										lowerCamelCase__   =True
										lowerCamelCase__   =99
										lowerCamelCase__   =32
										lowerCamelCase__   =2
										lowerCamelCase__   =4
										lowerCamelCase__   =37
										lowerCamelCase__   ="gelu"
										lowerCamelCase__   =0.1
										lowerCamelCase__   =0.1
										lowerCamelCase__   =512
										lowerCamelCase__   =16
										lowerCamelCase__   =2
										lowerCamelCase__   =0.0_2
										lowerCamelCase__   =3
										lowerCamelCase__   =4
										lowerCamelCase__   =None

				def 						_a      ( self ):
										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__   =EsmConfig(
										    vocab_size=self.vocab_size      ,					hidden_size=self.hidden_size      ,					num_hidden_layers=self.num_hidden_layers      ,					pad_token_id=1      ,					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      ,					)

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

				def 						_a      ( self ):
										(
										    (
										    lowerCamelCase__	
										)   , (
										    lowerCamelCase__	
										)   , (
										    lowerCamelCase__	
										)   , (
										    lowerCamelCase__	
										)   , (
										    lowerCamelCase__	
										)   , (
										    lowerCamelCase__	
										)   , 	
										)   =self.prepare_config_and_inputs()

										lowerCamelCase__   =True
										lowerCamelCase__   =floats_tensor([self.batch_size, self.seq_length, self.hidden_size] )
										lowerCamelCase__   =ids_tensor([self.batch_size, self.seq_length]      ,					vocab_size=2 )

										return (
										    config,
										    input_ids,
										    input_mask,
										    sequence_labels,
										    token_labels,
										    choice_labels,
										    encoder_hidden_states,
										    encoder_attention_mask,
										)

				def 						_a      ( self      ,					_lowerCamelCase      ,					_lowerCamelCase      ,					_lowerCamelCase      ,					_lowerCamelCase      ,					_lowerCamelCase      ,					_lowerCamelCase ):
										lowerCamelCase__   =TFEsmModel(config=lowerCAmelCase__ )
										lowerCamelCase__   ={"input_ids": input_ids, "attention_mask": input_mask}
										lowerCamelCase__   =model(lowerCAmelCase__ )

										lowerCamelCase__   =[input_ids, input_mask]
										lowerCamelCase__   =model(lowerCAmelCase__ )

										lowerCamelCase__   =model(lowerCAmelCase__ )

										self.parent.assertEqual(result.last_hidden_state.shape      ,					(self.batch_size, self.seq_length, self.hidden_size) )

				def 						_a      ( self      ,					_lowerCamelCase      ,					_lowerCamelCase      ,					_lowerCamelCase      ,					_lowerCamelCase      ,					_lowerCamelCase      ,					_lowerCamelCase      ,					_lowerCamelCase      ,					_lowerCamelCase      ,					):
										lowerCamelCase__   =True

										lowerCamelCase__   =TFEsmModel(config=lowerCAmelCase__ )
										lowerCamelCase__   ={
										    "input_ids": input_ids,
										    "attention_mask": input_mask,
										    "encoder_hidden_states": encoder_hidden_states,
										    "encoder_attention_mask": encoder_attention_mask,
										}
										lowerCamelCase__   =model(lowerCAmelCase__ )

										lowerCamelCase__   =[input_ids, input_mask]
										lowerCamelCase__   =model(lowerCAmelCase__      ,					encoder_hidden_states=lowerCAmelCase__ )

										# Also check the case where encoder outputs are not passed
										lowerCamelCase__   =model(lowerCAmelCase__      ,					attention_mask=lowerCAmelCase__ )

										self.parent.assertEqual(result.last_hidden_state.shape      ,					(self.batch_size, self.seq_length, self.hidden_size) )

				def 						_a      ( self      ,					_lowerCamelCase      ,					_lowerCamelCase      ,					_lowerCamelCase      ,					_lowerCamelCase      ,					_lowerCamelCase      ,					_lowerCamelCase ):
										lowerCamelCase__   =TFEsmForMaskedLM(config=lowerCAmelCase__ )
										lowerCamelCase__   =model([input_ids, input_mask] )
										self.parent.assertEqual(result.logits.shape      ,					(self.batch_size, self.seq_length, self.vocab_size) )

				def 						_a      ( self      ,					_lowerCamelCase      ,					_lowerCamelCase      ,					_lowerCamelCase      ,					_lowerCamelCase      ,					_lowerCamelCase      ,					_lowerCamelCase ):
										lowerCamelCase__   =self.num_labels
										lowerCamelCase__   =TFEsmForTokenClassification(config=lowerCAmelCase__ )
										lowerCamelCase__   ={"input_ids": input_ids, "attention_mask": input_mask}
										lowerCamelCase__   =model(lowerCAmelCase__ )
										self.parent.assertEqual(result.logits.shape      ,					(self.batch_size, self.seq_length, self.num_labels) )






				def 						_a      ( self ):
										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_tf
class   __UpperCAmelCase		(      _UpperCAmelCase , _UpperCAmelCase , unittest.TestCase      ):
				A__ :  List[str]    				=			(
				    (
				        TFEsmModel,
				        TFEsmForMaskedLM,
				        TFEsmForSequenceClassification,
				        TFEsmForTokenClassification,
				    )
				    if is_tf_available()
				    else ()
				)
				A__ :  str    				=			(
				    {
				        '''feature-extraction''': TFEsmModel,
				        '''fill-mask''': TFEsmForMaskedLM,
				        '''text-classification''': TFEsmForSequenceClassification,
				        '''token-classification''': TFEsmForTokenClassification,
				        '''zero-shot''': TFEsmForSequenceClassification,
				    }
				    if is_tf_available()
				    else {}
				)
				A__ :  Union[str, Any]    				=			False
				A__ :  Any    				=			False

				def 						_a      ( self ):
										lowerCamelCase__   =TFEsmModelTester(self )
										lowerCamelCase__   =ConfigTester(self      ,					config_class=lowerCAmelCase__      ,					hidden_size=37 )

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

				def 						_a      ( self ):
										lowerCamelCase__   =self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_model(*lowerCAmelCase__ )

				def 						_a      ( self ):
										lowerCamelCase__   =self.model_tester.prepare_config_and_inputs_for_decoder()
										self.model_tester.create_and_check_model_as_decoder(*lowerCAmelCase__ )

				def 						_a      ( self ):
										lowerCamelCase__   =self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_for_masked_lm(*lowerCAmelCase__ )

				def 						_a      ( self ):
										lowerCamelCase__   =self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_for_token_classification(*lowerCAmelCase__ )

				@slow
				def 						_a      ( self ):


										for model_name in TF_ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																lowerCamelCase__   =TFEsmModel.from_pretrained(lowerCAmelCase__ )
																self.assertIsNotNone(lowerCAmelCase__ )

				@unittest.skip("Protein models do not support embedding resizing." )
				def 						_a      ( self ):
										pass

				@unittest.skip("Protein models do not support embedding resizing." )
				def 						_a      ( self ):
										pass






				def 						_a      ( self ):
										lowerCamelCase__   , lowerCamelCase__   =self.model_tester.prepare_config_and_inputs_for_common()

										for model_class in self.all_model_classes:
																lowerCamelCase__   =model_class(lowerCAmelCase__ )
																assert isinstance(model.get_input_embeddings()      ,					tf.keras.layers.Layer )
																if model_class is TFEsmForMaskedLM:
																						# Output embedding test differs from the main test because they're a matrix, not a layer
																						lowerCamelCase__   =model.get_bias()
																						assert isinstance(lowerCAmelCase__      ,					lowerCAmelCase__ )
																						for k, v in name.items():
																												assert isinstance(lowerCAmelCase__      ,					tf.Variable )
																else:
																						lowerCamelCase__   =model.get_output_embeddings()
																						assert x is None
																						lowerCamelCase__   =model.get_bias()
																						assert name is None






@require_tf
class   __UpperCAmelCase		(      unittest.TestCase      ):

				@slow
				def 						_a      ( self ):
										lowerCamelCase__   =TFEsmForMaskedLM.from_pretrained("facebook/esm2_t6_8M_UR50D" )

										lowerCamelCase__   =tf.constant([[0, 1, 2, 3, 4, 5]] )
										lowerCamelCase__   =model(lowerCAmelCase__ )[0]
										lowerCamelCase__   =[1, 6, 33]
										self.assertEqual(list(output.numpy().shape )      ,					lowerCAmelCase__ )
										# compare the actual values for a slice.
										lowerCamelCase__   =tf.constant(
										    [
										        [
										            [8.9_2_1_5_1_8, -1_0.5_8_9_8_1_4, -6.4_6_7_1_3_0_7],
										            [-6.3_9_6_7_1_5_6, -1_3.9_1_1_3_7_7, -1.1_2_1_1_9_1_5],
										            [-7.7_8_1_2_4_7, -1_3.9_5_1_5_5_7, -3.7_4_0_5_9_2],
										        ]
										    ] )
										self.assertTrue(numpy.allclose(output[:, :3, :3].numpy()      ,					expected_slice.numpy()      ,					atol=1E-2 ) )






				@slow
				def 						_a      ( self ):
										lowerCamelCase__   =TFEsmModel.from_pretrained("facebook/esm2_t6_8M_UR50D" )

										lowerCamelCase__   =tf.constant([[0, 6, 4, 13, 5, 4, 16, 12, 11, 7, 2]] )
										lowerCamelCase__   =model(lowerCAmelCase__ )[0]
										# compare the actual values for a slice.
										lowerCamelCase__   =tf.constant(
										    [
										        [
										            [0.1_4_4_4_3_0_9_2, 0.5_4_1_2_5_3_2_7, 0.3_2_4_7_7_3_9],
										            [0.3_0_3_4_0_4_8_4, 0.0_0_5_2_6_6_7_6, 0.3_1_0_7_7_7_2_2],
										            [0.3_2_2_7_8_0_4_3, -0.2_4_9_8_7_0_9_6, 0.3_4_1_4_6_2_8],
										        ]
										    ] )
										self.assertTrue(numpy.allclose(output[:, :3, :3].numpy()      ,					expected_slice.numpy()      ,					atol=1E-4 ) )



import torch
from torch import nn

from ...configuration_utils import ConfigMixin, register_to_config
from ...models import ModelMixin




class   _UpperCamelCase (		_UpperCAmelCase		,_UpperCAmelCase					):


			"""simple docstring"""






			@register_to_config
			def __init__(  self ,  *,
			    lowerCAmelCase__ = 4 ,  lowerCAmelCase__ = 7_68 ,  lowerCAmelCase__ ,  lowerCAmelCase__ ,  )	->   Dict:

							'''simple docstring'''


							super().__init__()

							__lowercase					     =				nn.Parameter(torch.zeros(lowerCAmelCase__		)		)

							# parameters for additional clip time embeddings
							__lowercase					     =				nn.Linear(lowerCAmelCase__ ,  lowerCAmelCase__		)
							__lowercase					     =				nn.Linear(lowerCAmelCase__ ,  lowerCAmelCase__		)

							# parameters for encoder hidden states
							__lowercase					     =				clip_extra_context_tokens
							__lowercase					     =				nn.Linear(
							    lowerCAmelCase__ ,  self.clip_extra_context_tokens * cross_attention_dim		)
							__lowercase					     =				nn.Linear(lowerCAmelCase__ ,  lowerCAmelCase__		)
							__lowercase					     =				nn.LayerNorm(lowerCAmelCase__		)





			def 			_SCREAMING_SNAKE_CASE     (  self ,  *, lowerCAmelCase__ ,  lowerCAmelCase__ ,  lowerCAmelCase__ ,  lowerCAmelCase__		)	->   List[str]:

							'''simple docstring'''


							if do_classifier_free_guidance:
											# Add the classifier free guidance embeddings to the image embeddings
											__lowercase					     =				image_embeddings.shape[0]
											__lowercase					     =				self.learned_classifier_free_guidance_embeddings.unsqueeze(0		)
											__lowercase					     =				classifier_free_guidance_embeddings.expand(
											    lowerCAmelCase__ ,  -1		)
											__lowercase					     =				torch.cat([classifier_free_guidance_embeddings, image_embeddings] ,  dim=0		)

							# The image embeddings batch size and the text embeddings batch size are equal
							assert image_embeddings.shape[0] == prompt_embeds.shape[0]

							__lowercase					     =				prompt_embeds.shape[0]

							# "Specifically, we modify the architecture described in Nichol et al. (2021) by projecting and
							# adding CLIP embeddings to the existing timestep embedding, ...
							__lowercase					     =				self.embedding_proj(lowerCAmelCase__		)
							__lowercase					     =				self.clip_image_embeddings_project_to_time_embeddings(lowerCAmelCase__		)
							__lowercase					     =				time_projected_image_embeddings + time_projected_prompt_embeds

							# ... and by projecting CLIP embeddings into four
							# extra tokens of context that are concatenated to the sequence of outputs from the GLIDE text encoder"
							__lowercase					     =				self.clip_extra_context_tokens_proj(lowerCAmelCase__		)
							__lowercase					     =				clip_extra_context_tokens.reshape(lowerCAmelCase__ ,  -1 ,  self.clip_extra_context_tokens		)
							__lowercase					     =				clip_extra_context_tokens.permute(0 ,  2 ,  1		)

							__lowercase					     =				self.encoder_hidden_states_proj(lowerCAmelCase__		)
							__lowercase					     =				self.text_encoder_hidden_states_norm(lowerCAmelCase__		)
							__lowercase					     =				torch.cat([clip_extra_context_tokens, text_encoder_hidden_states] ,  dim=1		)

							return text_encoder_hidden_states, additive_clip_time_embeddings