
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					:






			'''simple docstring'''

			def __init__(							self			:							Dict    ,						snake_case_			:							int    ,						):
							UpperCamelCase_:				Dict						=    parent
							UpperCamelCase_:				Any						=    13
							UpperCamelCase_:				Dict						=    7
							UpperCamelCase_:				Union[str, Any]						=    True
							UpperCamelCase_:				Optional[Any]						=    True
							UpperCamelCase_:				Dict						=    True
							UpperCamelCase_:				List[str]						=    99
							UpperCamelCase_:				Optional[Any]						=    32
							UpperCamelCase_:				Union[str, Any]						=    2
							UpperCamelCase_:				Optional[Any]						=    4
							UpperCamelCase_:				List[str]						=    37
							UpperCamelCase_:				Any						=    'gelu'
							UpperCamelCase_:				Dict						=    0.1
							UpperCamelCase_:				List[str]						=    0.1
							UpperCamelCase_:				Union[str, Any]						=    512
							UpperCamelCase_:				Any						=    16
							UpperCamelCase_:				Optional[int]						=    2
							UpperCamelCase_:				str						=    0.02
							UpperCamelCase_:				str						=    3
							UpperCamelCase_:				int						=    4
							UpperCamelCase_:				Optional[int]						=    None

			def    lowerCAmelCase__						(							self			:							Any   ):
							UpperCamelCase_:				Any						=    ids_tensor([self.batch_size, self.seq_length]    ,						self.vocab_size   )

							UpperCamelCase_:				Union[str, Any]						=    None
							if self.use_input_mask:
											UpperCamelCase_:				Any						=    random_attention_mask([self.batch_size, self.seq_length]   )

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

							UpperCamelCase_:				List[str]						=    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			:							List[Any]   ):
							(
							    UpperCamelCase_							
							):				List[Any]						=    self.prepare_config_and_inputs()

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

							UpperCamelCase_:				Union[str, Any]						=    [input_ids, input_mask]
							UpperCamelCase_:				Any						=    model(snake_case_   )

							UpperCamelCase_:				Union[str, Any]						=    model(snake_case_   )

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

			def    lowerCAmelCase__						(							self			:							int    ,						snake_case_			:							str    ,						snake_case_			:							Optional[int]    ,						snake_case_			:							Tuple    ,						snake_case_			:							Any    ,						snake_case_			:							Tuple    ,						snake_case_			:							Union[str, Any]    ,						snake_case_			:							Union[str, Any]    ,						snake_case_			:							Tuple    ,						):
							UpperCamelCase_:				Optional[Any]						=    True

							UpperCamelCase_:				Any						=    TFEsmModel(config=snake_case_   )
							UpperCamelCase_:				Optional[Any]						=    {
							    'input_ids': input_ids,
							    'attention_mask': input_mask,
							    'encoder_hidden_states': encoder_hidden_states,
							    'encoder_attention_mask': encoder_attention_mask,
							}
							UpperCamelCase_:				Optional[int]						=    model(snake_case_   )

							UpperCamelCase_:				List[Any]						=    [input_ids, input_mask]
							UpperCamelCase_:				str						=    model(snake_case_    ,						encoder_hidden_states=snake_case_   )

							# Also check the case where encoder outputs are not passed
							UpperCamelCase_:				Union[str, 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			:							str    ,						snake_case_			:							Dict    ,						snake_case_			:							Tuple    ,						snake_case_			:							Tuple    ,						snake_case_			:							List[str]    ,						snake_case_			:							Optional[Any]    ,						snake_case_			:							Union[str, Any]   ):
							UpperCamelCase_:				List[Any]						=    TFEsmForMaskedLM(config=snake_case_   )
							UpperCamelCase_:				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			:							List[str]    ,						snake_case_			:							Optional[int]    ,						snake_case_			:							str    ,						snake_case_			:							Any    ,						snake_case_			:							Any    ,						snake_case_			:							str    ,						snake_case_			:							Tuple   ):
							UpperCamelCase_:				List[str]						=    self.num_labels
							UpperCamelCase_:				Tuple						=    TFEsmForTokenClassification(config=snake_case_   )
							UpperCamelCase_:				Optional[Any]						=    {'input_ids': input_ids, 'attention_mask': input_mask}
							UpperCamelCase_:				Any						=    model(snake_case_   )
							self.parent.assertEqual(result.logits.shape    ,						(self.batch_size, self.seq_length, self.num_labels)   )




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




@require_tf
class 							_UpperCamelCase					(       SCREAMING_SNAKE_CASE_  ,							SCREAMING_SNAKE_CASE_  ,							unittest.TestCase	):






			'''simple docstring'''

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

			def    lowerCAmelCase__						(							self			:							List[Any]   ):
							UpperCamelCase_:				int						=    TFEsmModelTester(self   )
							UpperCamelCase_:				Optional[int]						=    ConfigTester(self    ,						config_class=snake_case_    ,						hidden_size=37   )

			def    lowerCAmelCase__						(							self			:							List[str]   ):
							self.config_tester.run_common_tests()

			def    lowerCAmelCase__						(							self			:							Dict   ):
							UpperCamelCase_:				List[str]						=    self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_model(*snake_case_   )

			def    lowerCAmelCase__						(							self			:							Tuple   ):
							UpperCamelCase_:				List[str]						=    self.model_tester.prepare_config_and_inputs_for_decoder()
							self.model_tester.create_and_check_model_as_decoder(*snake_case_   )

			def    lowerCAmelCase__						(							self			:							List[str]   ):
							UpperCamelCase_:				Any						=    self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_masked_lm(*snake_case_   )

			def    lowerCAmelCase__						(							self			:							List[Any]   ):
							UpperCamelCase_:				Optional[Any]						=    self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_token_classification(*snake_case_   )

			@slow
			def    lowerCAmelCase__						(							self			:							Tuple   ):
							for model_name in TF_ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
											UpperCamelCase_:				Any						=    TFEsmModel.from_pretrained(snake_case_   )
											self.assertIsNotNone(snake_case_   )

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

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




			def    lowerCAmelCase__						(							self			:							Optional[int]   ):
							UpperCamelCase_:				List[str]						=    self.model_tester.prepare_config_and_inputs_for_common()

							for model_class in self.all_model_classes:
											UpperCamelCase_:				Union[str, 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
															UpperCamelCase_:				Tuple						=    model.get_bias()
															assert isinstance(snake_case_    ,						snake_case_   )
															for k, v in name.items():
																			assert isinstance(snake_case_    ,						tf.Variable   )
											else:
															UpperCamelCase_:				Union[str, Any]						=    model.get_output_embeddings()
															assert x is None
															UpperCamelCase_:				Dict						=    model.get_bias()
															assert name is None




@require_tf
class 							_UpperCamelCase					(       unittest.TestCase	):






			'''simple docstring'''

			@slow
			def    lowerCAmelCase__						(							self			:							str   ):
							UpperCamelCase_:				str						=    TFEsmForMaskedLM.from_pretrained("""facebook/esm2_t6_8M_UR50D"""   )

							UpperCamelCase_:				Tuple						=    tf.constant([[0, 1, 2, 3, 4, 5]]   )
							UpperCamelCase_:				str						=    model(snake_case_   )[0]
							UpperCamelCase_:				int						=    [1, 6, 33]
							self.assertEqual(list(output.numpy().shape   )    ,						snake_case_   )
							# compare the actual values for a slice.
							UpperCamelCase_:				Dict						=    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    lowerCAmelCase__						(							self			:							List[Any]   ):
							UpperCamelCase_:				int						=    TFEsmModel.from_pretrained("""facebook/esm2_t6_8M_UR50D"""   )

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











"""simple docstring"""







from math import ceil, sqrt

def     UpperCAmelCase_					(				__a   :			int = 1_00_00_00						):




			'''simple docstring'''



			_lowerCamelCase					:					Tuple	       =    0

			for outer_width in range(3      ,							(limit // 4) + 2						):
						if outer_width**2 > limit:
									_lowerCamelCase					:					Optional[int]	       =    max(ceil(sqrt(outer_width**2 - limit						)						)      ,							1						)
						else:
									_lowerCamelCase					:					Any	       =    1
						if (outer_width - hole_width_lower_bound) % 2:
									hole_width_lower_bound += 1

						answer += (outer_width - hole_width_lower_bound - 2) // 2 + 1

			return answer


if __name__ == "__main__":
	print(F"{solution() = }")





