

"""simple docstring"""





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  :







 """simple docstring"""







 def __init__(      self							:      List[Any], UpperCAmelCase__							:      int, ):
    __lowercase								=						parent
    __lowercase								=						1_3
    __lowercase								=						7
    __lowercase								=						True
    __lowercase								=						True
    __lowercase								=						True
    __lowercase								=						9_9
    __lowercase								=						3_2
    __lowercase								=						2
    __lowercase								=						4
    __lowercase								=						3_7
    __lowercase								=						"gelu"
    __lowercase								=						0.1
    __lowercase								=						0.1
    __lowercase								=						5_1_2
    __lowercase								=						1_6
    __lowercase								=						2
    __lowercase								=						0.02
    __lowercase								=						3
    __lowercase								=						4
    __lowercase								=						None


 def  _lowercase			(      self							:      List[str]     ):
    __lowercase								=						ids_tensor([self.batch_size, self.seq_length], self.vocab_size     )

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

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

    __lowercase								=						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  _lowercase			(      self							:      List[Any]     ):
    (
        (
        __lowercase
    )		,(
        __lowercase
    )		,(
        __lowercase
    )		,(
        __lowercase
    )		,(
        __lowercase
    )		,(
        __lowercase
    )		,
    )								=						self.prepare_config_and_inputs()

    __lowercase								=						True
    __lowercase								=						floats_tensor([self.batch_size, self.seq_length, self.hidden_size]     )
    __lowercase								=						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  _lowercase			(      self							:      str, UpperCAmelCase__							:      Dict, UpperCAmelCase__							:      Union[str, Any], UpperCAmelCase__							:      Any, UpperCAmelCase__							:      Optional[int], UpperCAmelCase__							:      Optional[int], UpperCAmelCase__							:      Any     ):
    __lowercase								=						TFEsmModel(config=UpperCAmelCase__     )
    __lowercase								=						{"input_ids": input_ids, "attention_mask": input_mask}
    __lowercase								=						model(UpperCAmelCase__     )

    __lowercase								=						[input_ids, input_mask]
    __lowercase								=						model(UpperCAmelCase__     )

    __lowercase								=						model(UpperCAmelCase__     )

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


 def  _lowercase			(      self							:      str, UpperCAmelCase__							:      List[str], UpperCAmelCase__							:      str, UpperCAmelCase__							:      List[str], UpperCAmelCase__							:      Optional[int], UpperCAmelCase__							:      Union[str, Any], UpperCAmelCase__							:      List[str], UpperCAmelCase__							:      Union[str, Any], UpperCAmelCase__							:      List[str], ):
    __lowercase								=						True

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

    __lowercase								=						[input_ids, input_mask]
    __lowercase								=						model(UpperCAmelCase__, encoder_hidden_states=UpperCAmelCase__     )

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

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


 def  _lowercase			(      self							:      List[str], UpperCAmelCase__							:      List[Any], UpperCAmelCase__							:      Optional[Any], UpperCAmelCase__							:      List[Any], UpperCAmelCase__							:      List[Any], UpperCAmelCase__							:      str, UpperCAmelCase__							:      Any     ):
    __lowercase								=						TFEsmForMaskedLM(config=UpperCAmelCase__     )
    __lowercase								=						model([input_ids, input_mask]     )
    self.parent.assertEqual(result.logits.shape, (self.batch_size, self.seq_length, self.vocab_size)     )


 def  _lowercase			(      self							:      Optional[Any], UpperCAmelCase__							:      Optional[int], UpperCAmelCase__							:      List[str], UpperCAmelCase__							:      Dict, UpperCAmelCase__							:      Any, UpperCAmelCase__							:      List[Any], UpperCAmelCase__							:      Optional[int]     ):
    __lowercase								=						self.num_labels
    __lowercase								=						TFEsmForTokenClassification(config=UpperCAmelCase__     )
    __lowercase								=						{"input_ids": input_ids, "attention_mask": input_mask}
    __lowercase								=						model(UpperCAmelCase__     )
    self.parent.assertEqual(result.logits.shape, (self.batch_size, self.seq_length, self.num_labels)     )



 def  _lowercase			(      self							:      Optional[int]     ):
    __lowercase								=						self.prepare_config_and_inputs()
    (
        (
        __lowercase
    )		,(
        __lowercase
    )		,(
        __lowercase
    )		,(
        __lowercase
    )		,(
        __lowercase
    )		,(
        __lowercase
    )		,
    )								=						config_and_inputs
    __lowercase								=						{"input_ids": input_ids, "attention_mask": input_mask}
    return config, inputs_dict







@require_tf
class  _lowerCAmelCase  (     lowercase					,lowercase					,unittest.TestCase       ):







 """simple docstring"""







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


 def  _lowercase			(      self							:      Optional[int]     ):
    __lowercase								=						TFEsmModelTester(self     )
    __lowercase								=						ConfigTester(self, config_class=UpperCAmelCase__, hidden_size=3_7     )


 def  _lowercase			(      self							:      Tuple     ):
    self.config_tester.run_common_tests()


 def  _lowercase			(      self							:      Optional[int]     ):
    __lowercase								=						self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_model(*UpperCAmelCase__     )


 def  _lowercase			(      self							:      Optional[Any]     ):
    __lowercase								=						self.model_tester.prepare_config_and_inputs_for_decoder()
    self.model_tester.create_and_check_model_as_decoder(*UpperCAmelCase__     )


 def  _lowercase			(      self							:      int     ):
    __lowercase								=						self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_for_masked_lm(*UpperCAmelCase__     )


 def  _lowercase			(      self							:      str     ):
    __lowercase								=						self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_for_token_classification(*UpperCAmelCase__     )


 @slow
 def  _lowercase			(      self							:      Any     ):
    for model_name in TF_ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
       __lowercase								=						TFEsmModel.from_pretrained(UpperCAmelCase__     )
       self.assertIsNotNone(UpperCAmelCase__     )


 @unittest.skip("Protein models do not support embedding resizing."     )
 def  _lowercase			(      self							:      Union[str, Any]     ):
    pass


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



 def  _lowercase			(      self							:      Dict     ):
    __lowercase		,__lowercase								=						self.model_tester.prepare_config_and_inputs_for_common()

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







@require_tf
class  _lowerCAmelCase  (     unittest.TestCase       ):







 """simple docstring"""







 @slow
 def  _lowercase			(      self							:      List[str]     ):
    __lowercase								=						TFEsmForMaskedLM.from_pretrained("facebook/esm2_t6_8M_UR50D"     )

    __lowercase								=						tf.constant([[0, 1, 2, 3, 4, 5]]     )
    __lowercase								=						model(UpperCAmelCase__     )[0]
    __lowercase								=						[1, 6, 3_3]
    self.assertEqual(list(output.numpy().shape     ), UpperCAmelCase__     )
    # compare the actual values for a slice.
    __lowercase								=						tf.constant(
        [
            [
                [8.921_518, -10.589_814, -6.4_671_307],
                [-6.3_967_156, -13.911_377, -1.1_211_915],
                [-7.781_247, -13.951_557, -3.740_592],
            ]
        ]     )
    self.assertTrue(numpy.allclose(output[:, :3, :3].numpy(), expected_slice.numpy(), atol=1E-2     )     )



 @slow
 def  _lowercase			(      self							:      List[str]     ):
    __lowercase								=						TFEsmModel.from_pretrained("facebook/esm2_t6_8M_UR50D"     )

    __lowercase								=						tf.constant([[0, 6, 4, 1_3, 5, 4, 1_6, 1_2, 1_1, 7, 2]]     )
    __lowercase								=						model(UpperCAmelCase__     )[0]
    # compare the actual values for a slice.
    __lowercase								=						tf.constant(
        [
            [
                [0.14_443_092, 0.54_125_327, 0.3_247_739],
                [0.30_340_484, 0.00_526_676, 0.31_077_722],
                [0.32_278_043, -0.24_987_096, 0.3_414_628],
            ]
        ]     )
    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  (  )     ->						Dict:
   '''simple docstring'''


   __lowercase								=						{
       "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],
   }
   __lowercase								=						Dataset.from_dict(UpperCamelCase_)
   return dataset




class  _lowerCAmelCase  (     lowercase       ):







 """simple docstring"""







 def  _lowercase			(      self							:      int     ):
    __lowercase								=						get_dataset()
    __lowercase								=						make_duplicate_clusters(UpperCAmelCase__, 0.85     )
    self.assertEqual(len(duplicate_clusters[0]     ), 2     )



 def  _lowercase			(      self							:      Any     ):
    __lowercase								=						get_dataset()
    __lowercase		,__lowercase								=						deduplicate_dataset(UpperCAmelCase__     )
    self.assertEqual(len(UpperCAmelCase__     ), 2     )
    print(UpperCAmelCase__     )
    self.assertEqual(duplicate_clusters[0][0]["copies"], 2     )
    self.assertEqual(duplicate_clusters[0][0]["is_extreme"], UpperCAmelCase__     )




