





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[str]      ,				UpperCamelCase      :					Any      ,				):


    '''simple docstring'''





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







  def UpperCamelCase_ (      self      :					Optional[Any]							):


    '''simple docstring'''





    _snake_case   :				int														=	ids_tensor([self.batch_size, self.seq_length]      ,				self.vocab_size							)

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

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

    _snake_case   :				int														=	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[int]							):


    '''simple docstring'''





    (
        (
        _snake_case
    )							,					(
        _snake_case
    )							,					(
        _snake_case
    )							,					(
        _snake_case
    )							,					(
        _snake_case
    )							,					(
        _snake_case
    )							,					
    )   :				Tuple														=	self.prepare_config_and_inputs()

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


    '''simple docstring'''





    _snake_case   :				int														=	TFEsmModel(config=UpperCamelCase							)
    _snake_case   :				int														=	{'input_ids': input_ids, 'attention_mask': input_mask}
    _snake_case   :				List[Any]														=	model(UpperCamelCase							)

    _snake_case   :				Any														=	[input_ids, input_mask]
    _snake_case   :				str														=	model(UpperCamelCase							)

    _snake_case   :				Any														=	model(UpperCamelCase							)

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







  def UpperCamelCase_ (      self      :					str      ,				UpperCamelCase      :					List[str]      ,				UpperCamelCase      :					List[Any]      ,				UpperCamelCase      :					Dict      ,				UpperCamelCase      :					Dict      ,				UpperCamelCase      :					List[Any]      ,				UpperCamelCase      :					Dict      ,				UpperCamelCase      :					Tuple      ,				UpperCamelCase      :					Tuple      ,				):


    '''simple docstring'''





    _snake_case   :				Union[str, Any]														=	True

    _snake_case   :				List[Any]														=	TFEsmModel(config=UpperCamelCase							)
    _snake_case   :				Tuple														=	{
        'input_ids': input_ids,
        'attention_mask': input_mask,
        'encoder_hidden_states': encoder_hidden_states,
        'encoder_attention_mask': encoder_attention_mask,
    }
    _snake_case   :				List[Any]														=	model(UpperCamelCase							)

    _snake_case   :				Tuple														=	[input_ids, input_mask]
    _snake_case   :				Union[str, Any]														=	model(UpperCamelCase      ,				encoder_hidden_states=UpperCamelCase							)

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

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







  def UpperCamelCase_ (      self      :					List[str]      ,				UpperCamelCase      :					Union[str, Any]      ,				UpperCamelCase      :					List[Any]      ,				UpperCamelCase      :					Optional[int]      ,				UpperCamelCase      :					str      ,				UpperCamelCase      :					Any      ,				UpperCamelCase      :					Dict							):


    '''simple docstring'''





    _snake_case   :				Any														=	TFEsmForMaskedLM(config=UpperCamelCase							)
    _snake_case   :				Dict														=	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]      ,				UpperCamelCase      :					str      ,				UpperCamelCase      :					Optional[Any]      ,				UpperCamelCase      :					Dict      ,				UpperCamelCase      :					int      ,				UpperCamelCase      :					Dict      ,				UpperCamelCase      :					Union[str, Any]							):


    '''simple docstring'''





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







  def UpperCamelCase_ (      self      :					Union[str, Any]							):


    '''simple docstring'''





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




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


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







  def UpperCamelCase_ (      self      :					Union[str, Any]							):


    '''simple docstring'''





    _snake_case   :				Any														=	TFEsmModelTester(self							)
    _snake_case   :				Union[str, Any]														=	ConfigTester(self      ,				config_class=UpperCamelCase      ,				hidden_size=37							)







  def UpperCamelCase_ (      self      :					List[str]							):


    '''simple docstring'''





    self.config_tester.run_common_tests()







  def UpperCamelCase_ (      self      :					List[str]							):


    '''simple docstring'''





    _snake_case   :				Dict														=	self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_model(*UpperCamelCase							)







  def UpperCamelCase_ (      self      :					Any							):


    '''simple docstring'''





    _snake_case   :				int														=	self.model_tester.prepare_config_and_inputs_for_decoder()
    self.model_tester.create_and_check_model_as_decoder(*UpperCamelCase							)







  def UpperCamelCase_ (      self      :					Optional[int]							):


    '''simple docstring'''





    _snake_case   :				Optional[int]														=	self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_for_masked_lm(*UpperCamelCase							)







  def UpperCamelCase_ (      self      :					Optional[int]							):


    '''simple docstring'''





    _snake_case   :				Optional[int]														=	self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_for_token_classification(*UpperCamelCase							)







  @slow
  def UpperCamelCase_ (      self      :					Dict							):


    '''simple docstring'''





    for model_name in TF_ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
      _snake_case   :				Tuple														=	TFEsmModel.from_pretrained(UpperCamelCase							)
      self.assertIsNotNone(UpperCamelCase							)







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


    '''simple docstring'''





    pass







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


    '''simple docstring'''





    pass







  def UpperCamelCase_ (      self      :					Optional[Any]							):


    '''simple docstring'''





    _snake_case							,					_snake_case   :				Tuple														=	self.model_tester.prepare_config_and_inputs_for_common()

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




@require_tf
class        _lowerCAmelCase    (			unittest.TestCase							):


  '''simple docstring'''







  @slow
  def UpperCamelCase_ (      self      :					List[Any]							):


    '''simple docstring'''





    _snake_case   :				int														=	TFEsmForMaskedLM.from_pretrained('facebook/esm2_t6_8M_UR50D'							)

    _snake_case   :				Optional[int]														=	tf.constant([[0, 1, 2, 3, 4, 5]]							)
    _snake_case   :				Tuple														=	model(UpperCamelCase							)[0]
    _snake_case   :				Tuple														=	[1, 6, 33]
    self.assertEqual(list(output.numpy().shape							)      ,				UpperCamelCase							)
    # compare the actual values for a slice.
    _snake_case   :				str														=	tf.constant(
        [
            [
                [8.92_15_18, -10.58_98_14, -6.4_67_13_07],
                [-6.3_96_71_56, -13.91_13_77, -1.1_21_19_15],
                [-7.78_12_47, -13.95_15_57, -3.74_05_92],
            ]
        ]							)
    self.assertTrue(numpy.allclose(output[:, :3, :3].numpy()      ,				expected_slice.numpy()      ,				atol=1e-2							)							)







  @slow
  def UpperCamelCase_ (      self      :					Optional[Any]							):


    '''simple docstring'''





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

    _snake_case   :				Dict														=	tf.constant([[0, 6, 4, 13, 5, 4, 16, 12, 11, 7, 2]]							)
    _snake_case   :				Any														=	model(UpperCamelCase							)[0]
    # compare the actual values for a slice.
    _snake_case   :				Any														=	tf.constant(
        [
            [
                [0.14_44_30_92, 0.54_12_53_27, 0.3_24_77_39],
                [0.30_34_04_84, 0.00_52_66_76, 0.31_07_77_22],
                [0.32_27_80_43, -0.24_98_70_96, 0.3_41_46_28],
            ]
        ]							)
    self.assertTrue(numpy.allclose(output[:, :3, :3].numpy()      ,				expected_slice.numpy()      ,				atol=1e-4							)							)











from unittest import TestCase

from datasets import Dataset
from minhash_deduplication import deduplicate_dataset, make_duplicate_clusters


def     lowerCamelCase_	(      )->		Any:
  _snake_case   :				List[str]														=	{
      '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],
  }
  _snake_case   :				Optional[Any]														=	Dataset.from_dict(lowerCAmelCase				)
  return dataset







class        _lowerCAmelCase    (			UpperCAmelCase_							):


  '''simple docstring'''







  def UpperCamelCase_ (      self      :					Dict							):


    '''simple docstring'''





    _snake_case   :				Union[str, Any]														=	get_dataset()
    _snake_case   :				Tuple														=	make_duplicate_clusters(UpperCamelCase      ,				0.85							)
    self.assertEqual(len(duplicate_clusters[0]							)      ,				2							)







  def UpperCamelCase_ (      self      :					Dict							):


    '''simple docstring'''





    _snake_case   :				List[str]														=	get_dataset()
    _snake_case							,					_snake_case   :				str														=	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							)




