




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  __magic_name__						:
				'''simple docstring'''







				def __init__(							self:   str		,					_lowerCamelCase:   Optional[Any]		,					):
											SCREAMING_SNAKE_CASE_   	=    parent
											SCREAMING_SNAKE_CASE_   	=    13
											SCREAMING_SNAKE_CASE_   	=    7
											SCREAMING_SNAKE_CASE_   	=    True
											SCREAMING_SNAKE_CASE_   	=    True
											SCREAMING_SNAKE_CASE_   	=    True
											SCREAMING_SNAKE_CASE_   	=    99
											SCREAMING_SNAKE_CASE_   	=    32
											SCREAMING_SNAKE_CASE_   	=    2
											SCREAMING_SNAKE_CASE_   	=    4
											SCREAMING_SNAKE_CASE_   	=    37
											SCREAMING_SNAKE_CASE_   	=    '''gelu'''
											SCREAMING_SNAKE_CASE_   	=    0.1
											SCREAMING_SNAKE_CASE_   	=    0.1
											SCREAMING_SNAKE_CASE_   	=    5_12
											SCREAMING_SNAKE_CASE_   	=    16
											SCREAMING_SNAKE_CASE_   	=    2
											SCREAMING_SNAKE_CASE_   	=    0.02
											SCREAMING_SNAKE_CASE_   	=    3
											SCREAMING_SNAKE_CASE_   	=    4
											SCREAMING_SNAKE_CASE_   	=    None


				def     _A							(							self:   str		):
											SCREAMING_SNAKE_CASE_   	=    ids_tensor([self.batch_size, self.seq_length]		,					self.vocab_size		)

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

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

											SCREAMING_SNAKE_CASE_   	=    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:   Optional[int]		):
											(
											    (
											    SCREAMING_SNAKE_CASE_
											)	,	(
											    SCREAMING_SNAKE_CASE_
											)	,	(
											    SCREAMING_SNAKE_CASE_
											)	,	(
											    SCREAMING_SNAKE_CASE_
											)	,	(
											    SCREAMING_SNAKE_CASE_
											)	,	(
											    SCREAMING_SNAKE_CASE_
											)	,	
											)   	=    self.prepare_config_and_inputs()

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

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

											SCREAMING_SNAKE_CASE_   	=    model(_lowerCamelCase		)

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


				def     _A							(							self:   str		,					_lowerCamelCase:   List[Any]		,					_lowerCamelCase:   str		,					_lowerCamelCase:   Union[str, Any]		,					_lowerCamelCase:   Dict		,					_lowerCamelCase:   Any		,					_lowerCamelCase:   Union[str, Any]		,					_lowerCamelCase:   Union[str, Any]		,					_lowerCamelCase:   str		,					):
											SCREAMING_SNAKE_CASE_   	=    True

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

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

											# Also check the case where encoder outputs are not passed
											SCREAMING_SNAKE_CASE_   	=    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:   str		,					_lowerCamelCase:   Union[str, Any]		,					_lowerCamelCase:   List[str]		,					_lowerCamelCase:   List[str]		,					_lowerCamelCase:   List[Any]		,					_lowerCamelCase:   str		,					_lowerCamelCase:   Dict		):
											SCREAMING_SNAKE_CASE_   	=    TFEsmForMaskedLM(config=_lowerCamelCase		)
											SCREAMING_SNAKE_CASE_   	=    model([input_ids, input_mask]		)
											self.parent.assertEqual(result.logits.shape		,					(self.batch_size, self.seq_length, self.vocab_size)		)


				def     _A							(							self:   Dict		,					_lowerCamelCase:   Dict		,					_lowerCamelCase:   Optional[Any]		,					_lowerCamelCase:   Any		,					_lowerCamelCase:   Optional[int]		,					_lowerCamelCase:   Tuple		,					_lowerCamelCase:   int		):
											SCREAMING_SNAKE_CASE_   	=    self.num_labels
											SCREAMING_SNAKE_CASE_   	=    TFEsmForTokenClassification(config=_lowerCamelCase		)
											SCREAMING_SNAKE_CASE_   	=    {'''input_ids''': input_ids, '''attention_mask''': input_mask}
											SCREAMING_SNAKE_CASE_   	=    model(_lowerCamelCase		)
											self.parent.assertEqual(result.logits.shape		,					(self.batch_size, self.seq_length, self.num_labels)		)





				def     _A							(							self:   Dict		):
											SCREAMING_SNAKE_CASE_   	=    self.prepare_config_and_inputs()
											(
											    (
											    SCREAMING_SNAKE_CASE_
											)	,	(
											    SCREAMING_SNAKE_CASE_
											)	,	(
											    SCREAMING_SNAKE_CASE_
											)	,	(
											    SCREAMING_SNAKE_CASE_
											)	,	(
											    SCREAMING_SNAKE_CASE_
											)	,	(
											    SCREAMING_SNAKE_CASE_
											)	,	
											)   	=    config_and_inputs
											SCREAMING_SNAKE_CASE_   	=    {'''input_ids''': input_ids, '''attention_mask''': input_mask}
											return config, inputs_dict






@require_tf
class  __magic_name__						(  __UpperCAmelCase    ,  __UpperCAmelCase    ,  unittest.TestCase):
				'''simple docstring'''







				SCREAMING_SNAKE_CASE__					:							List[Any]	      =					(
				    (
				        TFEsmModel,
				        TFEsmForMaskedLM,
				        TFEsmForSequenceClassification,
				        TFEsmForTokenClassification,
				    )
				    if is_tf_available()
				    else ()
				)
				SCREAMING_SNAKE_CASE__					:							str	      =					(
				    {
				        "feature-extraction": TFEsmModel,
				        "fill-mask": TFEsmForMaskedLM,
				        "text-classification": TFEsmForSequenceClassification,
				        "token-classification": TFEsmForTokenClassification,
				        "zero-shot": TFEsmForSequenceClassification,
				    }
				    if is_tf_available()
				    else {}
				)
				SCREAMING_SNAKE_CASE__					:							Optional[int]	      =					False
				SCREAMING_SNAKE_CASE__					:							Any	      =					False


				def     _A							(							self:   Any		):
											SCREAMING_SNAKE_CASE_   	=    TFEsmModelTester(self		)
											SCREAMING_SNAKE_CASE_   	=    ConfigTester(self		,					config_class=_lowerCamelCase		,					hidden_size=37		)


				def     _A							(							self:   Optional[Any]		):
											self.config_tester.run_common_tests()


				def     _A							(							self:   Tuple		):
											SCREAMING_SNAKE_CASE_   	=    self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_model(*_lowerCamelCase		)


				def     _A							(							self:   Union[str, Any]		):
											SCREAMING_SNAKE_CASE_   	=    self.model_tester.prepare_config_and_inputs_for_decoder()
											self.model_tester.create_and_check_model_as_decoder(*_lowerCamelCase		)


				def     _A							(							self:   Dict		):
											SCREAMING_SNAKE_CASE_   	=    self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_for_masked_lm(*_lowerCamelCase		)


				def     _A							(							self:   Tuple		):
											SCREAMING_SNAKE_CASE_   	=    self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_for_token_classification(*_lowerCamelCase		)


				@slow
				def     _A							(							self:   Union[str, Any]		):
											for model_name in TF_ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																		SCREAMING_SNAKE_CASE_   	=    TFEsmModel.from_pretrained(_lowerCamelCase		)
																		self.assertIsNotNone(_lowerCamelCase		)


				@unittest.skip('''Protein models do not support embedding resizing.'''		)
				def     _A							(							self:   Optional[int]		):
											pass


				@unittest.skip('''Protein models do not support embedding resizing.'''		)
				def     _A							(							self:   Any		):
											pass





				def     _A							(							self:   List[Any]		):
											SCREAMING_SNAKE_CASE_	,	SCREAMING_SNAKE_CASE_   	=    self.model_tester.prepare_config_and_inputs_for_common()

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






@require_tf
class  __magic_name__						(  unittest.TestCase):
				'''simple docstring'''







				@slow
				def     _A							(							self:   List[Any]		):
											SCREAMING_SNAKE_CASE_   	=    TFEsmForMaskedLM.from_pretrained('''facebook/esm2_t6_8M_UR50D'''		)

											SCREAMING_SNAKE_CASE_   	=    tf.constant([[0, 1, 2, 3, 4, 5]]		)
											SCREAMING_SNAKE_CASE_   	=    model(_lowerCamelCase		)[0]
											SCREAMING_SNAKE_CASE_   	=    [1, 6, 33]
											self.assertEqual(list(output.numpy().shape		)		,					_lowerCamelCase		)
											# compare the actual values for a slice.
											SCREAMING_SNAKE_CASE_   	=    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     _A							(							self:   Union[str, Any]		):
											SCREAMING_SNAKE_CASE_   	=    TFEsmModel.from_pretrained('''facebook/esm2_t6_8M_UR50D'''		)

											SCREAMING_SNAKE_CASE_   	=    tf.constant([[0, 6, 4, 13, 5, 4, 16, 12, 11, 7, 2]]		)
											SCREAMING_SNAKE_CASE_   	=    model(_lowerCamelCase		)[0]
											# compare the actual values for a slice.
											SCREAMING_SNAKE_CASE_   	=    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


__SCREAMING_SNAKE_CASE 						=logging.get_logger(__name__)

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






class  __magic_name__						(  __UpperCAmelCase):
				'''simple docstring'''







				SCREAMING_SNAKE_CASE__					:							Optional[Any]	      =					"swin2sr"

				SCREAMING_SNAKE_CASE__					:							str	      =					{
				    "hidden_size": "embed_dim",
				    "num_attention_heads": "num_heads",
				    "num_hidden_layers": "num_layers",
				}


				def __init__(							self:   Optional[int]		,					_lowerCamelCase:   Dict=64		,					_lowerCamelCase:   str=1		,					_lowerCamelCase:   Optional[int]=3		,					_lowerCamelCase:   List[Any]=1_80		,					_lowerCamelCase:   Optional[int]=[6, 6, 6, 6, 6, 6]		,					_lowerCamelCase:   Union[str, Any]=[6, 6, 6, 6, 6, 6]		,					_lowerCamelCase:   Tuple=8		,					_lowerCamelCase:   List[Any]=2.0		,					_lowerCamelCase:   Dict=True		,					_lowerCamelCase:   Dict=0.0		,					_lowerCamelCase:   int=0.0		,					_lowerCamelCase:   Optional[int]=0.1		,					_lowerCamelCase:   Any="gelu"		,					_lowerCamelCase:   int=False		,					_lowerCamelCase:   Tuple=0.02		,					_lowerCamelCase:   List[str]=1E-5		,					_lowerCamelCase:   List[Any]=2		,					_lowerCamelCase:   str=1.0		,					_lowerCamelCase:   List[Any]="1conv"		,					_lowerCamelCase:   str="pixelshuffle"		,					**_lowerCamelCase:   List[str]		,					):
											super().__init__(**_lowerCamelCase		)

											SCREAMING_SNAKE_CASE_   	=    image_size
											SCREAMING_SNAKE_CASE_   	=    patch_size
											SCREAMING_SNAKE_CASE_   	=    num_channels
											SCREAMING_SNAKE_CASE_   	=    embed_dim
											SCREAMING_SNAKE_CASE_   	=    depths
											SCREAMING_SNAKE_CASE_   	=    len(_lowerCamelCase		)
											SCREAMING_SNAKE_CASE_   	=    num_heads
											SCREAMING_SNAKE_CASE_   	=    window_size
											SCREAMING_SNAKE_CASE_   	=    mlp_ratio
											SCREAMING_SNAKE_CASE_   	=    qkv_bias
											SCREAMING_SNAKE_CASE_   	=    hidden_dropout_prob
											SCREAMING_SNAKE_CASE_   	=    attention_probs_dropout_prob
											SCREAMING_SNAKE_CASE_   	=    drop_path_rate
											SCREAMING_SNAKE_CASE_   	=    hidden_act
											SCREAMING_SNAKE_CASE_   	=    use_absolute_embeddings
											SCREAMING_SNAKE_CASE_   	=    layer_norm_eps
											SCREAMING_SNAKE_CASE_   	=    initializer_range
											SCREAMING_SNAKE_CASE_   	=    upscale
											SCREAMING_SNAKE_CASE_   	=    img_range
											SCREAMING_SNAKE_CASE_   	=    resi_connection
											SCREAMING_SNAKE_CASE_   	=    upsampler


