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



	"""simple docstring"""




	def __init__(       self  :			Tuple       ,							SCREAMING_SNAKE_CASE_  :			Union[str, Any]       ,							)       ->     Any:
					lowercase_       =					parent
					lowercase_       =					1_3
					lowercase_       =					7
					lowercase_       =					True
					lowercase_       =					True
					lowercase_       =					True
					lowercase_       =					9_9
					lowercase_       =					3_2
					lowercase_       =					2
					lowercase_       =					4
					lowercase_       =					3_7
					lowercase_       =					'''gelu'''
					lowercase_       =					0.1
					lowercase_       =					0.1
					lowercase_       =					5_1_2
					lowercase_       =					1_6
					lowercase_       =					2
					lowercase_       =					0.02
					lowercase_       =					3
					lowercase_       =					4
					lowercase_       =					None
	def 						_lowercase  (       self  :			Optional[int]       )       ->     Tuple:
					lowercase_       =					ids_tensor([self.batch_size, self.seq_length]       ,							self.vocab_size       )

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

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

					lowercase_       =					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 						_lowercase  (       self  :			Any       )       ->     Optional[Any]:
					(
					    (
					    lowercase_
					)  ,	(
					    lowercase_
					)  ,	(
					    lowercase_
					)  ,	(
					    lowercase_
					)  ,	(
					    lowercase_
					)  ,	(
					    lowercase_
					)  ,	
					)       =					self.prepare_config_and_inputs()

					lowercase_       =					True
					lowercase_       =					floats_tensor([self.batch_size, self.seq_length, self.hidden_size]       )
					lowercase_       =					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 						_lowercase  (       self  :			Tuple       ,							SCREAMING_SNAKE_CASE_  :			Any       ,							SCREAMING_SNAKE_CASE_  :			List[str]       ,							SCREAMING_SNAKE_CASE_  :			str       ,							SCREAMING_SNAKE_CASE_  :			List[str]       ,							SCREAMING_SNAKE_CASE_  :			Optional[Any]       ,							SCREAMING_SNAKE_CASE_  :			str       )       ->     List[str]:
					lowercase_       =					TFEsmModel(config=SCREAMING_SNAKE_CASE_       )
					lowercase_       =					{'''input_ids''': input_ids, '''attention_mask''': input_mask}
					lowercase_       =					model(SCREAMING_SNAKE_CASE_       )

					lowercase_       =					[input_ids, input_mask]
					lowercase_       =					model(SCREAMING_SNAKE_CASE_       )

					lowercase_       =					model(SCREAMING_SNAKE_CASE_       )

					self.parent.assertEqual(result.last_hidden_state.shape       ,							(self.batch_size, self.seq_length, self.hidden_size)       )
	def 						_lowercase  (       self  :			List[Any]       ,							SCREAMING_SNAKE_CASE_  :			Optional[int]       ,							SCREAMING_SNAKE_CASE_  :			List[Any]       ,							SCREAMING_SNAKE_CASE_  :			Dict       ,							SCREAMING_SNAKE_CASE_  :			Optional[Any]       ,							SCREAMING_SNAKE_CASE_  :			List[Any]       ,							SCREAMING_SNAKE_CASE_  :			int       ,							SCREAMING_SNAKE_CASE_  :			int       ,							SCREAMING_SNAKE_CASE_  :			str       ,							)       ->     Union[str, Any]:
					lowercase_       =					True

					lowercase_       =					TFEsmModel(config=SCREAMING_SNAKE_CASE_       )
					lowercase_       =					{
					    '''input_ids''': input_ids,
					    '''attention_mask''': input_mask,
					    '''encoder_hidden_states''': encoder_hidden_states,
					    '''encoder_attention_mask''': encoder_attention_mask,
					}
					lowercase_       =					model(SCREAMING_SNAKE_CASE_       )

					lowercase_       =					[input_ids, input_mask]
					lowercase_       =					model(SCREAMING_SNAKE_CASE_       ,							encoder_hidden_states=SCREAMING_SNAKE_CASE_       )

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

					self.parent.assertEqual(result.last_hidden_state.shape       ,							(self.batch_size, self.seq_length, self.hidden_size)       )
	def 						_lowercase  (       self  :			Optional[int]       ,							SCREAMING_SNAKE_CASE_  :			List[str]       ,							SCREAMING_SNAKE_CASE_  :			str       ,							SCREAMING_SNAKE_CASE_  :			Dict       ,							SCREAMING_SNAKE_CASE_  :			Optional[Any]       ,							SCREAMING_SNAKE_CASE_  :			Any       ,							SCREAMING_SNAKE_CASE_  :			Any       )       ->     Any:
					lowercase_       =					TFEsmForMaskedLM(config=SCREAMING_SNAKE_CASE_       )
					lowercase_       =					model([input_ids, input_mask]       )
					self.parent.assertEqual(result.logits.shape       ,							(self.batch_size, self.seq_length, self.vocab_size)       )
	def 						_lowercase  (       self  :			List[Any]       ,							SCREAMING_SNAKE_CASE_  :			List[Any]       ,							SCREAMING_SNAKE_CASE_  :			Optional[int]       ,							SCREAMING_SNAKE_CASE_  :			Union[str, Any]       ,							SCREAMING_SNAKE_CASE_  :			Optional[int]       ,							SCREAMING_SNAKE_CASE_  :			Optional[Any]       ,							SCREAMING_SNAKE_CASE_  :			str       )       ->     int:
					lowercase_       =					self.num_labels
					lowercase_       =					TFEsmForTokenClassification(config=SCREAMING_SNAKE_CASE_       )
					lowercase_       =					{'''input_ids''': input_ids, '''attention_mask''': input_mask}
					lowercase_       =					model(SCREAMING_SNAKE_CASE_       )
					self.parent.assertEqual(result.logits.shape       ,							(self.batch_size, self.seq_length, self.num_labels)       )




	def 						_lowercase  (       self  :			Tuple       )       ->     str:
					lowercase_       =					self.prepare_config_and_inputs()
					(
					    (
					    lowercase_
					)  ,	(
					    lowercase_
					)  ,	(
					    lowercase_
					)  ,	(
					    lowercase_
					)  ,	(
					    lowercase_
					)  ,	(
					    lowercase_
					)  ,	
					)       =					config_and_inputs
					lowercase_       =					{'''input_ids''': input_ids, '''attention_mask''': input_mask}
					return config, inputs_dict

@require_tf
class  lowercase__(				UpperCAmelCase ,						UpperCAmelCase ,						unittest.TestCase							):



	"""simple docstring"""




	a			:str         =	(
	    (
	        TFEsmModel,
	        TFEsmForMaskedLM,
	        TFEsmForSequenceClassification,
	        TFEsmForTokenClassification,
	    )
	    if is_tf_available()
	    else ()
	)
	a			:Union[str, Any]         =	(
	    {
	        '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			:Dict         =	False
	def 						_lowercase  (       self  :			List[str]       )       ->     List[Any]:
					lowercase_       =					TFEsmModelTester(self       )
					lowercase_       =					ConfigTester(self       ,							config_class=SCREAMING_SNAKE_CASE_       ,							hidden_size=3_7       )
	def 						_lowercase  (       self  :			Optional[int]       )       ->     str:
					self.config_tester.run_common_tests()
	def 						_lowercase  (       self  :			str       )       ->     int:
					lowercase_       =					self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_model(*SCREAMING_SNAKE_CASE_       )
	def 						_lowercase  (       self  :			Any       )       ->     Union[str, Any]:
					lowercase_       =					self.model_tester.prepare_config_and_inputs_for_decoder()
					self.model_tester.create_and_check_model_as_decoder(*SCREAMING_SNAKE_CASE_       )
	def 						_lowercase  (       self  :			str       )       ->     Optional[int]:
					lowercase_       =					self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_for_masked_lm(*SCREAMING_SNAKE_CASE_       )
	def 						_lowercase  (       self  :			Optional[Any]       )       ->     List[str]:
					lowercase_       =					self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_for_token_classification(*SCREAMING_SNAKE_CASE_       )
	@slow
	def 						_lowercase  (       self  :			int       )       ->     List[str]:
					for model_name in TF_ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
									lowercase_       =					TFEsmModel.from_pretrained(SCREAMING_SNAKE_CASE_       )
									self.assertIsNotNone(SCREAMING_SNAKE_CASE_       )
	@unittest.skip('''Protein models do not support embedding resizing.'''       )
	def 						_lowercase  (       self  :			str       )       ->     Tuple:
					pass
	@unittest.skip('''Protein models do not support embedding resizing.'''       )
	def 						_lowercase  (       self  :			List[Any]       )       ->     Any:
					pass




	def 						_lowercase  (       self  :			List[str]       )       ->     Any:
					lowercase_  ,	lowercase_       =					self.model_tester.prepare_config_and_inputs_for_common()

					for model_class in self.all_model_classes:
									lowercase_       =					model_class(SCREAMING_SNAKE_CASE_       )
									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
													lowercase_       =					model.get_bias()
													assert isinstance(SCREAMING_SNAKE_CASE_       ,							SCREAMING_SNAKE_CASE_       )
													for k, v in name.items():
																	assert isinstance(SCREAMING_SNAKE_CASE_       ,							tf.Variable       )
									else:
													lowercase_       =					model.get_output_embeddings()
													assert x is None
													lowercase_       =					model.get_bias()
													assert name is None

@require_tf
class  lowercase__(				unittest.TestCase							):



	"""simple docstring"""




	@slow
	def 						_lowercase  (       self  :			str       )       ->     Optional[int]:
					lowercase_       =					TFEsmForMaskedLM.from_pretrained('''facebook/esm2_t6_8M_UR50D'''       )

					lowercase_       =					tf.constant([[0, 1, 2, 3, 4, 5]]       )
					lowercase_       =					model(SCREAMING_SNAKE_CASE_       )[0]
					lowercase_       =					[1, 6, 3_3]
					self.assertEqual(list(output.numpy().shape       )       ,							SCREAMING_SNAKE_CASE_       )
					# compare the actual values for a slice.
					lowercase_       =					tf.constant(
					    [
					        [
					            [8.92_15_18, -10.58_98_14, -6.4_67_13_07],
					            [-6.3_96_71_56, -13.91_13_77, -1.1_21_19_15],
					            [-7.78_12_47, -13.95_15_57, -3.74_05_92],
					        ]
					    ]       )
					self.assertTrue(numpy.allclose(output[:, :3, :3].numpy()       ,							expected_slice.numpy()       ,							atol=1e-2       )       )




	@slow
	def 						_lowercase  (       self  :			int       )       ->     Union[str, Any]:
					lowercase_       =					TFEsmModel.from_pretrained('''facebook/esm2_t6_8M_UR50D'''       )

					lowercase_       =					tf.constant([[0, 6, 4, 1_3, 5, 4, 1_6, 1_2, 1_1, 7, 2]]       )
					lowercase_       =					model(SCREAMING_SNAKE_CASE_       )[0]
					# compare the actual values for a slice.
					lowercase_       =					tf.constant(
					    [
					        [
					            [0.14_44_30_92, 0.54_12_53_27, 0.3_24_77_39],
					            [0.30_34_04_84, 0.00_52_66_76, 0.31_07_77_22],
					            [0.32_27_80_43, -0.24_98_70_96, 0.3_41_46_28],
					        ]
					    ]       )
					self.assertTrue(numpy.allclose(output[:, :3, :3].numpy()       ,							expected_slice.numpy()       ,							atol=1e-4       )       )











from ...configuration_utils import PretrainedConfig
from ...utils import logging


UpperCAmelCase	:	Dict		     =       logging.get_logger(__name__)

UpperCAmelCase	:	Tuple		     =       {
    """caidas/swin2sr-classicalsr-x2-64""": (
        """https://huggingface.co/caidas/swin2sr-classicalsr-x2-64/resolve/main/config.json"""
    ),
}







class 				__lowerCAmelCase       (    UpperCamelCase__):
      _lowercase				:					Any						=			"""swin2sr"""

      _lowercase				:					Tuple						=			{
          """hidden_size""": """embed_dim""",
          """num_attention_heads""": """num_heads""",
          """num_hidden_layers""": """num_layers""",
      }







      def __init__(    self		,      lowerCAmelCase__=6_4		,      lowerCAmelCase__=1		,      lowerCAmelCase__=3		,      lowerCAmelCase__=1_8_0		,      lowerCAmelCase__=[6, 6, 6, 6, 6, 6]		,      lowerCAmelCase__=[6, 6, 6, 6, 6, 6]		,      lowerCAmelCase__=8		,      lowerCAmelCase__=2.0		,      lowerCAmelCase__=True		,      lowerCAmelCase__=0.0		,      lowerCAmelCase__=0.0		,      lowerCAmelCase__=0.1		,      lowerCAmelCase__="gelu"		,      lowerCAmelCase__=False		,      lowerCAmelCase__=0.02		,      lowerCAmelCase__=1E-5		,      lowerCAmelCase__=2		,      lowerCAmelCase__=1.0		,      lowerCAmelCase__="1conv"		,      lowerCAmelCase__="pixelshuffle"		,      **lowerCAmelCase__		,      )   ->						int:



        '''simple docstring'''




        super().__init__(**lowerCAmelCase__							)

        a__	:					Optional[Any]										=image_size
        a__	:					Dict										=patch_size
        a__	:					Tuple										=num_channels
        a__	:					Union[str, Any]										=embed_dim
        a__	:					Optional[Any]										=depths
        a__	:					List[str]										=len(lowerCAmelCase__							)
        a__	:					Any										=num_heads
        a__	:					Any										=window_size
        a__	:					str										=mlp_ratio
        a__	:					List[str]										=qkv_bias
        a__	:					Dict										=hidden_dropout_prob
        a__	:					List[str]										=attention_probs_dropout_prob
        a__	:					Dict										=drop_path_rate
        a__	:					Optional[Any]										=hidden_act
        a__	:					Union[str, Any]										=use_absolute_embeddings
        a__	:					Optional[Any]										=layer_norm_eps
        a__	:					List[Any]										=initializer_range
        a__	:					int										=upscale
        a__	:					Optional[int]										=img_range
        a__	:					Any										=resi_connection
        a__	:					Optional[Any]										=upsampler






