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







							def __init__(   self      :     List[Any] ,			_A      :     Optional[int] ,			):
												_UpperCamelCase			     =     parent
												_UpperCamelCase			     =     13
												_UpperCamelCase			     =     7
												_UpperCamelCase			     =     True
												_UpperCamelCase			     =     True
												_UpperCamelCase			     =     True
												_UpperCamelCase			     =     99
												_UpperCamelCase			     =     32
												_UpperCamelCase			     =     2
												_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







							def  UpperCamelCase_						(   self      :     Optional[int]		):
												_UpperCamelCase			     =     ids_tensor([self.batch_size, self.seq_length] ,			self.vocab_size		)

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

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

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

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

												_UpperCamelCase			     =     [input_ids, input_mask]
												_UpperCamelCase			     =     model(_A		)

												_UpperCamelCase			     =     model(_A		)

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







							def  UpperCamelCase_						(   self      :     List[str] ,			_A      :     List[Any] ,			_A      :     Optional[Any] ,			_A      :     Optional[int] ,			_A      :     Tuple ,			_A      :     Dict ,			_A      :     Tuple ,			_A      :     Tuple ,			_A      :     List[Any] ,			):
												_UpperCamelCase			     =     True

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

												_UpperCamelCase			     =     [input_ids, input_mask]
												_UpperCamelCase			     =     model(_A ,			encoder_hidden_states=_A		)

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

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







							def  UpperCamelCase_						(   self      :     int ,			_A      :     List[Any] ,			_A      :     str ,			_A      :     Optional[Any] ,			_A      :     Optional[int] ,			_A      :     Optional[Any] ,			_A      :     Tuple		):
												_UpperCamelCase			     =     TFEsmForMaskedLM(config=_A		)
												_UpperCamelCase			     =     model([input_ids, input_mask]		)
												self.parent.assertEqual(result.logits.shape ,			(self.batch_size, self.seq_length, self.vocab_size)		)







							def  UpperCamelCase_						(   self      :     List[Any] ,			_A      :     List[str] ,			_A      :     List[str] ,			_A      :     Any ,			_A      :     Dict ,			_A      :     str ,			_A      :     Any		):
												_UpperCamelCase			     =     self.num_labels
												_UpperCamelCase			     =     TFEsmForTokenClassification(config=_A		)
												_UpperCamelCase			     =     {'''input_ids''': input_ids, '''attention_mask''': input_mask}
												_UpperCamelCase			     =     model(_A		)
												self.parent.assertEqual(result.logits.shape ,			(self.batch_size, self.seq_length, self.num_labels)		)







							def  UpperCamelCase_						(   self      :     Optional[Any]		):
												_UpperCamelCase			     =     self.prepare_config_and_inputs()
												(
												    (
												    _UpperCamelCase
												)     ,    (
												    _UpperCamelCase
												)     ,    (
												    _UpperCamelCase
												)     ,    (
												    _UpperCamelCase
												)     ,    (
												    _UpperCamelCase
												)     ,    (
												    _UpperCamelCase
												)     ,    
												)			     =     config_and_inputs
												_UpperCamelCase			     =     {'''input_ids''': input_ids, '''attention_mask''': input_mask}
												return config, inputs_dict




@require_tf
class     lowerCAmelCase_     (       __lowercase,  __lowercase,  unittest.TestCase ):
							UpperCAmelCase				   =     (
							    (
							        TFEsmModel,
							        TFEsmForMaskedLM,
							        TFEsmForSequenceClassification,
							        TFEsmForTokenClassification,
							    )
							    if is_tf_available()
							    else ()
							)
							UpperCAmelCase				   =     (
							    {
							        "feature-extraction": TFEsmModel,
							        "fill-mask": TFEsmForMaskedLM,
							        "text-classification": TFEsmForSequenceClassification,
							        "token-classification": TFEsmForTokenClassification,
							        "zero-shot": TFEsmForSequenceClassification,
							    }
							    if is_tf_available()
							    else {}
							)
							UpperCAmelCase				   =     False
							UpperCAmelCase				   =     False







							def  UpperCamelCase_						(   self      :     str		):
												_UpperCamelCase			     =     TFEsmModelTester(self		)
												_UpperCamelCase			     =     ConfigTester(self ,			config_class=_A ,			hidden_size=37		)







							def  UpperCamelCase_						(   self      :     List[Any]		):
												self.config_tester.run_common_tests()







							def  UpperCamelCase_						(   self      :     Any		):
												_UpperCamelCase			     =     self.model_tester.prepare_config_and_inputs()
												self.model_tester.create_and_check_model(*_A		)







							def  UpperCamelCase_						(   self      :     str		):
												_UpperCamelCase			     =     self.model_tester.prepare_config_and_inputs_for_decoder()
												self.model_tester.create_and_check_model_as_decoder(*_A		)







							def  UpperCamelCase_						(   self      :     Tuple		):
												_UpperCamelCase			     =     self.model_tester.prepare_config_and_inputs()
												self.model_tester.create_and_check_for_masked_lm(*_A		)







							def  UpperCamelCase_						(   self      :     Optional[int]		):
												_UpperCamelCase			     =     self.model_tester.prepare_config_and_inputs()
												self.model_tester.create_and_check_for_token_classification(*_A		)







							@slow
							def  UpperCamelCase_						(   self      :     List[Any]		):
												for model_name in TF_ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																	_UpperCamelCase			     =     TFEsmModel.from_pretrained(_A		)
																	self.assertIsNotNone(_A		)







							@unittest.skip('''Protein models do not support embedding resizing.'''		)
							def  UpperCamelCase_						(   self      :     List[Any]		):
												pass







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







							def  UpperCamelCase_						(   self      :     str		):
												_UpperCamelCase     ,    _UpperCamelCase			     =     self.model_tester.prepare_config_and_inputs_for_common()

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




@require_tf
class     lowerCAmelCase_     (       unittest.TestCase ):







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

												_UpperCamelCase			     =     tf.constant([[0, 1, 2, 3, 4, 5]]		)
												_UpperCamelCase			     =     model(_A		)[0]
												_UpperCamelCase			     =     [1, 6, 33]
												self.assertEqual(list(output.numpy().shape		) ,			_A		)
												# compare the actual values for a slice.
												_UpperCamelCase			     =     tf.constant(
												    [
												        [
												            [8.92_1518, -10.58_9814, -6.467_1307],
												            [-6.396_7156, -13.91_1377, -1.121_1915],
												            [-7.78_1247, -13.95_1557, -3.74_0592],
												        ]
												    ]		)
												self.assertTrue(numpy.allclose(output[:, :3, :3].numpy() ,			expected_slice.numpy() ,			atol=1e-2		)		)







							@slow
							def  UpperCamelCase_						(   self      :     Dict		):
												_UpperCamelCase			     =     TFEsmModel.from_pretrained('''facebook/esm2_t6_8M_UR50D'''		)

												_UpperCamelCase			     =     tf.constant([[0, 6, 4, 13, 5, 4, 16, 12, 11, 7, 2]]		)
												_UpperCamelCase			     =     model(_A		)[0]
												# compare the actual values for a slice.
												_UpperCamelCase			     =     tf.constant(
												    [
												        [
												            [0.1444_3092, 0.5412_5327, 0.324_7739],
												            [0.3034_0484, 0.0052_6676, 0.3107_7722],
												            [0.3227_8043, -0.2498_7096, 0.341_4628],
												        ]
												    ]		)
												self.assertTrue(numpy.allclose(output[:, :3, :3].numpy() ,			expected_slice.numpy() ,			atol=1e-4		)		)

from typing import List

from .keymap import KEYMAP, get_character
def 					_snake_case	(  __snake_case						):

					def decorator(__snake_case						):
										_UpperCamelCase			     =     getattr(__snake_case  ,     '''handle_key'''  ,     []						)
										handle += [key]
										setattr(__snake_case  ,     '''handle_key'''  ,     __snake_case						)
										return func

					return decorator
def 					_snake_case	(  *__snake_case						):

					def decorator(__snake_case						):
										_UpperCamelCase			     =     getattr(__snake_case  ,     '''handle_key'''  ,     []						)
										handle += keys
										setattr(__snake_case  ,     '''handle_key'''  ,     __snake_case						)
										return func

					return decorator
class     lowerCAmelCase_     (       __lowercase ):







							def __new__(   cls      :     Optional[Any] ,			_A      :     Optional[Any] ,			_A      :     Optional[int] ,			_A      :     Union[str, Any]		):
												_UpperCamelCase			     =     super().__new__(cls ,			_A ,			_A ,			_A		)
												if not hasattr(_A ,			'''key_handler'''		):
																	setattr(_A ,			'''key_handler''' ,			{}		)
												setattr(_A ,			'''handle_input''' ,			KeyHandler.handle_input		)

												for value in attrs.values():
																	_UpperCamelCase			     =     getattr(_A ,			'''handle_key''' ,			[]		)
																	for key in handled_keys:
																						_UpperCamelCase			     =     value
												return new_cls







							@staticmethod
							def  UpperCamelCase_						(   cls      :     str		):
												_UpperCamelCase			     =     get_character()
												if char != KEYMAP["undefined"]:
																	_UpperCamelCase			     =     ord(_A		)
												_UpperCamelCase			     =     cls.key_handler.get(_A		)
												if handler:
																	_UpperCamelCase			     =     char
																	return handler(cls		)
												else:
																	return None




def 					_snake_case	(  cls						):
					return KeyHandler(cls.__name__  ,     cls.__bases__  ,     cls.__dict__.copy()						)
