



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    :				Dict       ,	snake_case    :				Any       ,	):



											'''simple docstring'''


											SCREAMING_SNAKE_CASE							:     str						=				parent
											SCREAMING_SNAKE_CASE							:     List[Any]						=				13
											SCREAMING_SNAKE_CASE							:     Dict						=				7
											SCREAMING_SNAKE_CASE							:     Union[str, Any]						=				True
											SCREAMING_SNAKE_CASE							:     int						=				True
											SCREAMING_SNAKE_CASE							:     Any						=				True
											SCREAMING_SNAKE_CASE							:     Tuple						=				99
											SCREAMING_SNAKE_CASE							:     Union[str, Any]						=				32
											SCREAMING_SNAKE_CASE							:     int						=				2
											SCREAMING_SNAKE_CASE							:     int						=				4
											SCREAMING_SNAKE_CASE							:     Optional[Any]						=				37
											SCREAMING_SNAKE_CASE							:     int						=				'gelu'
											SCREAMING_SNAKE_CASE							:     List[str]						=				0.1
											SCREAMING_SNAKE_CASE							:     Dict						=				0.1
											SCREAMING_SNAKE_CASE							:     Any						=				512
											SCREAMING_SNAKE_CASE							:     Optional[int]						=				16
											SCREAMING_SNAKE_CASE							:     List[str]						=				2
											SCREAMING_SNAKE_CASE							:     List[Any]						=				0.02
											SCREAMING_SNAKE_CASE							:     Optional[int]						=				3
											SCREAMING_SNAKE_CASE							:     List[str]						=				4
											SCREAMING_SNAKE_CASE							:     Tuple						=				None







						def       lowerCamelCase_				(  self    :				Dict	):



											'''simple docstring'''


											SCREAMING_SNAKE_CASE							:     Tuple						=				ids_tensor([self.batch_size, self.seq_length]       ,	self.vocab_size	)

											SCREAMING_SNAKE_CASE							:     Tuple						=				None
											if self.use_input_mask:
																SCREAMING_SNAKE_CASE							:     str						=				random_attention_mask([self.batch_size, self.seq_length]	)

											SCREAMING_SNAKE_CASE							:     Any						=				None
											SCREAMING_SNAKE_CASE							:     Optional[Any]						=				None
											SCREAMING_SNAKE_CASE							:     Optional[int]						=				None
											if self.use_labels:
																SCREAMING_SNAKE_CASE							:     List[Any]						=				ids_tensor([self.batch_size]       ,	self.type_sequence_label_size	)
																SCREAMING_SNAKE_CASE							:     str						=				ids_tensor([self.batch_size, self.seq_length]       ,	self.num_labels	)
																SCREAMING_SNAKE_CASE							:     Optional[int]						=				ids_tensor([self.batch_size]       ,	self.num_choices	)

											SCREAMING_SNAKE_CASE							:     Dict						=				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       lowerCamelCase_				(  self    :				Optional[Any]	):



											'''simple docstring'''


											(
											    (
											    SCREAMING_SNAKE_CASE
											)					,  (
											    SCREAMING_SNAKE_CASE
											)					,  (
											    SCREAMING_SNAKE_CASE
											)					,  (
											    SCREAMING_SNAKE_CASE
											)					,  (
											    SCREAMING_SNAKE_CASE
											)					,  (
											    SCREAMING_SNAKE_CASE
											)					,  
											)							:     Optional[int]						=				self.prepare_config_and_inputs()

											SCREAMING_SNAKE_CASE							:     List[str]						=				True
											SCREAMING_SNAKE_CASE							:     List[Any]						=				floats_tensor([self.batch_size, self.seq_length, self.hidden_size]	)
											SCREAMING_SNAKE_CASE							:     List[Any]						=				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       lowerCamelCase_				(  self    :				Dict       ,	snake_case    :				List[str]       ,	snake_case    :				Tuple       ,	snake_case    :				Optional[Any]       ,	snake_case    :				Optional[Any]       ,	snake_case    :				Optional[Any]       ,	snake_case    :				int	):



											'''simple docstring'''


											SCREAMING_SNAKE_CASE							:     str						=				TFEsmModel(config=snake_case	)
											SCREAMING_SNAKE_CASE							:     List[Any]						=				{'input_ids': input_ids, 'attention_mask': input_mask}
											SCREAMING_SNAKE_CASE							:     Tuple						=				model(snake_case	)

											SCREAMING_SNAKE_CASE							:     Tuple						=				[input_ids, input_mask]
											SCREAMING_SNAKE_CASE							:     List[Any]						=				model(snake_case	)

											SCREAMING_SNAKE_CASE							:     List[Any]						=				model(snake_case	)

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







						def       lowerCamelCase_				(  self    :				str       ,	snake_case    :				str       ,	snake_case    :				Tuple       ,	snake_case    :				Union[str, Any]       ,	snake_case    :				str       ,	snake_case    :				Optional[Any]       ,	snake_case    :				Any       ,	snake_case    :				Union[str, Any]       ,	snake_case    :				Union[str, Any]       ,	):



											'''simple docstring'''


											SCREAMING_SNAKE_CASE							:     List[Any]						=				True

											SCREAMING_SNAKE_CASE							:     Union[str, Any]						=				TFEsmModel(config=snake_case	)
											SCREAMING_SNAKE_CASE							:     Optional[int]						=				{
											    'input_ids': input_ids,
											    'attention_mask': input_mask,
											    'encoder_hidden_states': encoder_hidden_states,
											    'encoder_attention_mask': encoder_attention_mask,
											}
											SCREAMING_SNAKE_CASE							:     Optional[int]						=				model(snake_case	)

											SCREAMING_SNAKE_CASE							:     Optional[int]						=				[input_ids, input_mask]
											SCREAMING_SNAKE_CASE							:     Any						=				model(snake_case       ,	encoder_hidden_states=snake_case	)

											# Also check the case where encoder outputs are not passed
											SCREAMING_SNAKE_CASE							:     List[Any]						=				model(snake_case       ,	attention_mask=snake_case	)

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







						def       lowerCamelCase_				(  self    :				List[str]       ,	snake_case    :				Optional[int]       ,	snake_case    :				Optional[int]       ,	snake_case    :				Tuple       ,	snake_case    :				int       ,	snake_case    :				Dict       ,	snake_case    :				int	):



											'''simple docstring'''


											SCREAMING_SNAKE_CASE							:     Optional[Any]						=				TFEsmForMaskedLM(config=snake_case	)
											SCREAMING_SNAKE_CASE							:     Union[str, Any]						=				model([input_ids, input_mask]	)
											self.parent.assertEqual(result.logits.shape       ,	(self.batch_size, self.seq_length, self.vocab_size)	)







						def       lowerCamelCase_				(  self    :				Any       ,	snake_case    :				int       ,	snake_case    :				Optional[Any]       ,	snake_case    :				List[Any]       ,	snake_case    :				Union[str, Any]       ,	snake_case    :				List[str]       ,	snake_case    :				Dict	):



											'''simple docstring'''


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







						def       lowerCamelCase_				(  self    :				int	):



											'''simple docstring'''


											SCREAMING_SNAKE_CASE							:     Union[str, Any]						=				self.prepare_config_and_inputs()
											(
											    (
											    SCREAMING_SNAKE_CASE
											)					,  (
											    SCREAMING_SNAKE_CASE
											)					,  (
											    SCREAMING_SNAKE_CASE
											)					,  (
											    SCREAMING_SNAKE_CASE
											)					,  (
											    SCREAMING_SNAKE_CASE
											)					,  (
											    SCREAMING_SNAKE_CASE
											)					,  
											)							:     Dict						=				config_and_inputs
											SCREAMING_SNAKE_CASE							:     Optional[Any]						=				{'input_ids': input_ids, 'attention_mask': input_mask}
											return config, inputs_dict


@require_tf
class 						lowercase    (							SCREAMING_SNAKE_CASE_      ,    SCREAMING_SNAKE_CASE_      ,    unittest.TestCase):
						'''simple docstring'''


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







						def       lowerCamelCase_				(  self    :				str	):



											'''simple docstring'''


											SCREAMING_SNAKE_CASE							:     Optional[int]						=				TFEsmModelTester(self	)
											SCREAMING_SNAKE_CASE							:     str						=				ConfigTester(self       ,	config_class=snake_case       ,	hidden_size=37	)







						def       lowerCamelCase_				(  self    :				List[Any]	):



											'''simple docstring'''


											self.config_tester.run_common_tests()







						def       lowerCamelCase_				(  self    :				List[Any]	):



											'''simple docstring'''


											SCREAMING_SNAKE_CASE							:     List[Any]						=				self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_model(*snake_case	)







						def       lowerCamelCase_				(  self    :				List[Any]	):



											'''simple docstring'''


											SCREAMING_SNAKE_CASE							:     Tuple						=				self.model_tester.prepare_config_and_inputs_for_decoder()
											self.model_tester.create_and_check_model_as_decoder(*snake_case	)







						def       lowerCamelCase_				(  self    :				int	):



											'''simple docstring'''


											SCREAMING_SNAKE_CASE							:     Optional[Any]						=				self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_for_masked_lm(*snake_case	)







						def       lowerCamelCase_				(  self    :				List[Any]	):



											'''simple docstring'''


											SCREAMING_SNAKE_CASE							:     Optional[int]						=				self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_for_token_classification(*snake_case	)







						@slow
						def       lowerCamelCase_				(  self    :				List[str]	):



											'''simple docstring'''


											for model_name in TF_ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																SCREAMING_SNAKE_CASE							:     List[Any]						=				TFEsmModel.from_pretrained(snake_case	)
																self.assertIsNotNone(snake_case	)







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



											'''simple docstring'''


											pass







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



											'''simple docstring'''


											pass







						def       lowerCamelCase_				(  self    :				Tuple	):



											'''simple docstring'''


											SCREAMING_SNAKE_CASE					,  SCREAMING_SNAKE_CASE							:     Dict						=				self.model_tester.prepare_config_and_inputs_for_common()

											for model_class in self.all_model_classes:
																SCREAMING_SNAKE_CASE							:     Any						=				model_class(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
																					SCREAMING_SNAKE_CASE							:     Optional[int]						=				model.get_bias()
																					assert isinstance(snake_case       ,	snake_case	)
																					for k, v in name.items():
																										assert isinstance(snake_case       ,	tf.Variable	)
																else:
																					SCREAMING_SNAKE_CASE							:     Any						=				model.get_output_embeddings()
																					assert x is None
																					SCREAMING_SNAKE_CASE							:     List[Any]						=				model.get_bias()
																					assert name is None


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


						@slow
						def       lowerCamelCase_				(  self    :				Dict	):



											'''simple docstring'''


											SCREAMING_SNAKE_CASE							:     Tuple						=				TFEsmForMaskedLM.from_pretrained('facebook/esm2_t6_8M_UR50D'	)

											SCREAMING_SNAKE_CASE							:     Optional[int]						=				tf.constant([[0, 1, 2, 3, 4, 5]]	)
											SCREAMING_SNAKE_CASE							:     Optional[Any]						=				model(snake_case	)[0]
											SCREAMING_SNAKE_CASE							:     str						=				[1, 6, 33]
											self.assertEqual(list(output.numpy().shape	)       ,	snake_case	)
											# compare the actual values for a slice.
											SCREAMING_SNAKE_CASE							:     Tuple						=				tf.constant(
											    [
											        [
											            [8.921518, -10.589814, -6.4671307],
											            [-6.3967156, -13.911377, -1.1211915],
											            [-7.781247, -13.951557, -3.740592],
											        ]
											    ]	)
											self.assertTrue(numpy.allclose(output[:, :3, :3].numpy()       ,	expected_slice.numpy()       ,	atol=1E-2	)	)







						@slow
						def       lowerCamelCase_				(  self    :				Tuple	):



											'''simple docstring'''


											SCREAMING_SNAKE_CASE							:     Any						=				TFEsmModel.from_pretrained('facebook/esm2_t6_8M_UR50D'	)

											SCREAMING_SNAKE_CASE							:     Optional[Any]						=				tf.constant([[0, 6, 4, 13, 5, 4, 16, 12, 11, 7, 2]]	)
											SCREAMING_SNAKE_CASE							:     int						=				model(snake_case	)[0]
											# compare the actual values for a slice.
											SCREAMING_SNAKE_CASE							:     Union[str, Any]						=				tf.constant(
											    [
											        [
											            [0.14443092, 0.54125327, 0.3247739],
											            [0.30340484, 0.00526676, 0.31077722],
											            [0.32278043, -0.24987096, 0.3414628],
											        ]
											    ]	)
											self.assertTrue(numpy.allclose(output[:, :3, :3].numpy()       ,	expected_slice.numpy()       ,	atol=1E-4	)	)




def        __a			(   __lowerCAmelCase   )     ->     list[list[float]]:
					SCREAMING_SNAKE_CASE							:     list[list[float]]						=				[]
					for data in source_data:
										for i, el in enumerate(__lowerCAmelCase   ):
															if len(__lowerCAmelCase   ) < i + 1:
																				data_lists.append([]   )
															data_lists[i].append(float(__lowerCAmelCase   )   )
					return data_lists

def        __a			(   __lowerCAmelCase      ,	__lowerCAmelCase   )     ->     list[list[float]]:
					SCREAMING_SNAKE_CASE							:     list[list[float]]						=				[]
					for dlist, weight in zip(__lowerCAmelCase      ,	__lowerCAmelCase   ):
										SCREAMING_SNAKE_CASE							:     str						=				min(__lowerCAmelCase   )
										SCREAMING_SNAKE_CASE							:     int						=				max(__lowerCAmelCase   )

										SCREAMING_SNAKE_CASE							:     list[float]						=				[]
										# for weight 0 score is 1 - actual score
										if weight == 0:
															for item in dlist:
																				try:
																									score.append(1 - ((item - mind) / (maxd - mind))   )
																				except ZeroDivisionError:
																									score.append(1   )

										elif weight == 1:
															for item in dlist:
																				try:
																									score.append((item - mind) / (maxd - mind)   )
																				except ZeroDivisionError:
																									score.append(0   )

        # weight not 0 or 1
										else:
															SCREAMING_SNAKE_CASE							:     Union[str, Any]						=				F'''Invalid weight of {weight:f} provided'''
															raise ValueError(__lowerCAmelCase   )

										score_lists.append(__lowerCAmelCase   )

					return score_lists

def        __a			(   __lowerCAmelCase   )     ->     list[float]:
					SCREAMING_SNAKE_CASE							:     list[float]						=				[0 for i in range(len(score_lists[0]   )   )]

					for slist in score_lists:
										for j, ele in enumerate(__lowerCAmelCase   ):
															SCREAMING_SNAKE_CASE							:     str						=				final_scores[j] + ele

					return final_scores



def        __a			(   __lowerCAmelCase      ,	__lowerCAmelCase   )     ->     list[list[float]]:
					SCREAMING_SNAKE_CASE							:     Tuple						=				get_data(__lowerCAmelCase   )
					SCREAMING_SNAKE_CASE							:     List[str]						=				calculate_each_score(__lowerCAmelCase      ,	__lowerCAmelCase   )
					SCREAMING_SNAKE_CASE							:     str						=				generate_final_scores(__lowerCAmelCase   )

					# append scores to source data
					for i, ele in enumerate(__lowerCAmelCase   ):
										source_data[i].append(__lowerCAmelCase   )

					return source_data