
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					:



		def __init__( self				,      A_				,      )  ->					List[str]:
				"""simple docstring"""
				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       )  ->					List[str]:
				"""simple docstring"""
				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       )  ->					Tuple:
				"""simple docstring"""
				(
				    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				,      A_				,      A_				,      A_				,      A_				,      A_				,      A_       )  ->					Dict:
				"""simple docstring"""
				UpperCamelCase		=     TFEsmModel(config=lowercase_       )
				UpperCamelCase		=     {"input_ids": input_ids, "attention_mask": input_mask}
				UpperCamelCase		=     model(lowercase_       )

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

				UpperCamelCase		=     model(lowercase_       )

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



		def __UpperCamelCase ( self				,      A_				,      A_				,      A_				,      A_				,      A_				,      A_				,      A_				,      A_				,      )  ->					str:
				"""simple docstring"""
				UpperCamelCase		=     True

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

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

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

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



		def __UpperCamelCase ( self				,      A_				,      A_				,      A_				,      A_				,      A_				,      A_       )  ->					str:
				"""simple docstring"""
				UpperCamelCase		=     TFEsmForMaskedLM(config=lowercase_       )
				UpperCamelCase		=     model([input_ids, input_mask]       )
				self.parent.assertEqual(result.logits.shape				,      (self.batch_size, self.seq_length, self.vocab_size)       )



		def __UpperCamelCase ( self				,      A_				,      A_				,      A_				,      A_				,      A_				,      A_       )  ->					int:
				"""simple docstring"""
				UpperCamelCase		=     self.num_labels
				UpperCamelCase		=     TFEsmForTokenClassification(config=lowercase_       )
				UpperCamelCase		=     {"input_ids": input_ids, "attention_mask": input_mask}
				UpperCamelCase		=     model(lowercase_       )
				self.parent.assertEqual(result.logits.shape				,      (self.batch_size, self.seq_length, self.num_labels)       )



		def __UpperCamelCase ( self       )  ->					Optional[Any]:
				"""simple docstring"""
				UpperCamelCase		=     self.prepare_config_and_inputs()
				(
				    UpperCamelCase   
				)		=     config_and_inputs
				UpperCamelCase		=     {"input_ids": input_ids, "attention_mask": input_mask}
				return config, inputs_dict






@require_tf
class 				lowercase					(		lowercase__  ,   lowercase__  ,   unittest.TestCase   ):
		__lowercase       :       Any	     =				(
		    (
		        TFEsmModel,
		        TFEsmForMaskedLM,
		        TFEsmForSequenceClassification,
		        TFEsmForTokenClassification,
		    )
		    if is_tf_available()
		    else ()
		)
		__lowercase       :       Optional[Any]	     =				(
		    {
		        """feature-extraction""": TFEsmModel,
		        """fill-mask""": TFEsmForMaskedLM,
		        """text-classification""": TFEsmForSequenceClassification,
		        """token-classification""": TFEsmForTokenClassification,
		        """zero-shot""": TFEsmForSequenceClassification,
		    }
		    if is_tf_available()
		    else {}
		)
		__lowercase       :       int	     =				False
		__lowercase       :       List[str]	     =				False



		def __UpperCamelCase ( self       )  ->					Optional[int]:
				"""simple docstring"""
				UpperCamelCase		=     TFEsmModelTester(self       )
				UpperCamelCase		=     ConfigTester(self				,      config_class=lowercase_				,      hidden_size=37       )



		def __UpperCamelCase ( self       )  ->					Optional[int]:
				"""simple docstring"""
				self.config_tester.run_common_tests()



		def __UpperCamelCase ( self       )  ->					Optional[Any]:
				"""simple docstring"""
				UpperCamelCase		=     self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_model(*lowercase_       )



		def __UpperCamelCase ( self       )  ->					Optional[int]:
				"""simple docstring"""
				UpperCamelCase		=     self.model_tester.prepare_config_and_inputs_for_decoder()
				self.model_tester.create_and_check_model_as_decoder(*lowercase_       )



		def __UpperCamelCase ( self       )  ->					int:
				"""simple docstring"""
				UpperCamelCase		=     self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_for_masked_lm(*lowercase_       )



		def __UpperCamelCase ( self       )  ->					Dict:
				"""simple docstring"""
				UpperCamelCase		=     self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_for_token_classification(*lowercase_       )



		@slow
		def __UpperCamelCase ( self       )  ->					List[Any]:
				"""simple docstring"""
				for model_name in TF_ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
						UpperCamelCase		=     TFEsmModel.from_pretrained(lowercase_       )
						self.assertIsNotNone(lowercase_       )



		@unittest.skip('Protein models do not support embedding resizing.'       )
		def __UpperCamelCase ( self       )  ->					int:
				"""simple docstring"""
				pass



		@unittest.skip('Protein models do not support embedding resizing.'       )
		def __UpperCamelCase ( self       )  ->					Tuple:
				"""simple docstring"""
				pass



		def __UpperCamelCase ( self       )  ->					Union[str, Any]:
				"""simple docstring"""
				UpperCamelCase		=     self.model_tester.prepare_config_and_inputs_for_common()

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






@require_tf
class 				lowercase					(		unittest.TestCase   ):



		@slow
		def __UpperCamelCase ( self       )  ->					str:
				"""simple docstring"""
				UpperCamelCase		=     TFEsmForMaskedLM.from_pretrained('facebook/esm2_t6_8M_UR50D'       )

				UpperCamelCase		=     tf.constant([[0, 1, 2, 3, 4, 5]]       )
				UpperCamelCase		=     model(lowercase_       )[0]
				UpperCamelCase		=     [1, 6, 33]
				self.assertEqual(list(output.numpy().shape       )				,      lowercase_       )
				# 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       )  ->					Tuple:
				"""simple docstring"""
				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(lowercase_       )[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       )       )


"""simple docstring"""





from unittest import TestCase

from datasets import Dataset
from minhash_deduplication import deduplicate_dataset, make_duplicate_clusters







def 	__a     (		):
   UpperCAmelCase_					:					List[Any]			     =    {
       "repo_name": ["test_repo1", "test_repo2", "test_repo3"],
       "path": ["test_1.py", "test_2.py", "unit_test.py"],
       "content": ["a " * 20, "a " * 30, "b " * 7],
   }
   UpperCAmelCase_					:					Optional[int]			     =    Dataset.from_dict(__lowerCamelCase  )
   return dataset


class A_	(lowercase__      ):




  '''simple docstring'''







  def        UpperCamelCase__					(		self  ):

     """simple docstring"""
     UpperCAmelCase_					:					Union[str, Any]			     =    get_dataset()
     UpperCAmelCase_					:					Any			     =    make_duplicate_clusters(lowercase_				,   0.85  )
     self.assertEqual(len(duplicate_clusters[0]  )				,   2  )







  def        UpperCamelCase__					(		self  ):

     """simple docstring"""
     UpperCAmelCase_					:					List[str]			     =    get_dataset()
     UpperCAmelCase_     ,   UpperCAmelCase_					:					List[str]			     =    deduplicate_dataset(lowercase_  )
     self.assertEqual(len(lowercase_  )				,   2  )
     print(lowercase_  )
     self.assertEqual(duplicate_clusters[0][0]["copies"]				,   2  )
     self.assertEqual(duplicate_clusters[0][0]["is_extreme"]				,   lowercase_  )


