'''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    a :



  """simple docstring"""







  def __init__(							self   ,    snake_case_   ,    ):

    '''simple docstring'''







    __UpperCAmelCase:					Optional[Any]		= parent
    __UpperCAmelCase:					List[str]		= 13
    __UpperCAmelCase:					int		= 7
    __UpperCAmelCase:					int		= True
    __UpperCAmelCase:					List[Any]		= True
    __UpperCAmelCase:					List[str]		= True
    __UpperCAmelCase:					str		= 99
    __UpperCAmelCase:					List[str]		= 32
    __UpperCAmelCase:					Dict		= 2
    __UpperCAmelCase:					List[str]		= 4
    __UpperCAmelCase:					Dict		= 37
    __UpperCAmelCase:					Tuple		= """gelu"""
    __UpperCAmelCase:					Optional[Any]		= 0.1
    __UpperCAmelCase:					Tuple		= 0.1
    __UpperCAmelCase:					List[Any]		= 512
    __UpperCAmelCase:					Optional[int]		= 16
    __UpperCAmelCase:					int		= 2
    __UpperCAmelCase:					int		= 0.0_2
    __UpperCAmelCase:					List[Any]		= 3
    __UpperCAmelCase:					Optional[Any]		= 4
    __UpperCAmelCase:					Optional[Any]		= None







  def    lowercase_  (							self					):

    '''simple docstring'''







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

    __UpperCAmelCase:					Any		= None
    if self.use_input_mask:
      __UpperCAmelCase:					List[Any]		= random_attention_mask([self.batch_size, self.seq_length]					)

    __UpperCAmelCase:					Any		= None
    __UpperCAmelCase:					int		= None
    __UpperCAmelCase:					str		= None
    if self.use_labels:
      __UpperCAmelCase:					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:					Optional[Any]		= ids_tensor([self.batch_size]   ,    self.num_choices					)

    __UpperCAmelCase:					Optional[Any]		= 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					):

    '''simple docstring'''







    (
        (
        __UpperCAmelCase
    ), (
        __UpperCAmelCase
    ), (
        __UpperCAmelCase
    ), (
        __UpperCAmelCase
    ), (
        __UpperCAmelCase
    ), (
        __UpperCAmelCase
    ), 
    ):					Union[str, Any]		= self.prepare_config_and_inputs()

    __UpperCAmelCase:					Optional[int]		= True
    __UpperCAmelCase:					Tuple		= floats_tensor([self.batch_size, self.seq_length, self.hidden_size]					)
    __UpperCAmelCase:					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    lowercase_  (							self   ,    snake_case_   ,    snake_case_   ,    snake_case_   ,    snake_case_   ,    snake_case_   ,    snake_case_					):

    '''simple docstring'''







    __UpperCAmelCase:					int		= TFEsmModel(config=snake_case_					)
    __UpperCAmelCase:					Tuple		= {"""input_ids""": input_ids, """attention_mask""": input_mask}
    __UpperCAmelCase:					str		= model(snake_case_					)

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

    __UpperCAmelCase:					Dict		= model(snake_case_					)

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







  def    lowercase_  (							self   ,    snake_case_   ,    snake_case_   ,    snake_case_   ,    snake_case_   ,    snake_case_   ,    snake_case_   ,    snake_case_   ,    snake_case_   ,    ):

    '''simple docstring'''







    __UpperCAmelCase:					str		= True

    __UpperCAmelCase:					Tuple		= TFEsmModel(config=snake_case_					)
    __UpperCAmelCase:					List[str]		= {
        """input_ids""": input_ids,
        """attention_mask""": input_mask,
        """encoder_hidden_states""": encoder_hidden_states,
        """encoder_attention_mask""": encoder_attention_mask,
    }
    __UpperCAmelCase:					Any		= model(snake_case_					)

    __UpperCAmelCase:					Tuple		= [input_ids, input_mask]
    __UpperCAmelCase:					int		= model(snake_case_   ,    encoder_hidden_states=snake_case_					)

    # Also check the case where encoder outputs are not passed
    __UpperCAmelCase:					Tuple		= 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    lowercase_  (							self   ,    snake_case_   ,    snake_case_   ,    snake_case_   ,    snake_case_   ,    snake_case_   ,    snake_case_					):

    '''simple docstring'''







    __UpperCAmelCase:					int		= TFEsmForMaskedLM(config=snake_case_					)
    __UpperCAmelCase:					Dict		= model([input_ids, input_mask]					)
    self.parent.assertEqual(result.logits.shape   ,    (self.batch_size, self.seq_length, self.vocab_size)					)







  def    lowercase_  (							self   ,    snake_case_   ,    snake_case_   ,    snake_case_   ,    snake_case_   ,    snake_case_   ,    snake_case_					):

    '''simple docstring'''







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







  def    lowercase_  (							self					):

    '''simple docstring'''







    __UpperCAmelCase:					str		= self.prepare_config_and_inputs()
    (
        (
        __UpperCAmelCase
    ), (
        __UpperCAmelCase
    ), (
        __UpperCAmelCase
    ), (
        __UpperCAmelCase
    ), (
        __UpperCAmelCase
    ), (
        __UpperCAmelCase
    ), 
    ):					Optional[Any]		= config_and_inputs
    __UpperCAmelCase:					Optional[int]		= {"""input_ids""": input_ids, """attention_mask""": input_mask}
    return config, inputs_dict




@require_tf
class    a (       __lowerCAmelCase						,  __lowerCAmelCase						,  unittest.TestCase				):



  """simple docstring"""
  __lowerCAmelCase				      =					(
      (
          TFEsmModel,
          TFEsmForMaskedLM,
          TFEsmForSequenceClassification,
          TFEsmForTokenClassification,
      )
      if is_tf_available()
      else ()
  )
  __lowerCAmelCase				      =					(
      {
          """feature-extraction""": TFEsmModel,
          """fill-mask""": TFEsmForMaskedLM,
          """text-classification""": TFEsmForSequenceClassification,
          """token-classification""": TFEsmForTokenClassification,
          """zero-shot""": TFEsmForSequenceClassification,
      }
      if is_tf_available()
      else {}
  )
  __lowerCAmelCase				      =					False
  __lowerCAmelCase				      =					False







  def    lowercase_  (							self					):

    '''simple docstring'''







    __UpperCAmelCase:					Union[str, Any]		= TFEsmModelTester(self					)
    __UpperCAmelCase:					List[Any]		= ConfigTester(self   ,    config_class=snake_case_   ,    hidden_size=37					)







  def    lowercase_  (							self					):

    '''simple docstring'''







    self.config_tester.run_common_tests()







  def    lowercase_  (							self					):

    '''simple docstring'''







    __UpperCAmelCase:					str		= self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_model(*snake_case_					)







  def    lowercase_  (							self					):

    '''simple docstring'''







    __UpperCAmelCase:					List[str]		= self.model_tester.prepare_config_and_inputs_for_decoder()
    self.model_tester.create_and_check_model_as_decoder(*snake_case_					)







  def    lowercase_  (							self					):

    '''simple docstring'''







    __UpperCAmelCase:					Optional[int]		= self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_for_masked_lm(*snake_case_					)







  def    lowercase_  (							self					):

    '''simple docstring'''







    __UpperCAmelCase:					int		= self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_for_token_classification(*snake_case_					)







  @slow
  def    lowercase_  (							self					):

    '''simple docstring'''







    for model_name in TF_ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
      __UpperCAmelCase:					Union[str, Any]		= TFEsmModel.from_pretrained(snake_case_					)
      self.assertIsNotNone(snake_case_					)







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

    '''simple docstring'''







    pass







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

    '''simple docstring'''







    pass







  def    lowercase_  (							self					):

    '''simple docstring'''







    __UpperCAmelCase, __UpperCAmelCase:					Optional[Any]		= self.model_tester.prepare_config_and_inputs_for_common()

    for model_class in self.all_model_classes:
      __UpperCAmelCase:					Optional[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:					Any		= 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:					Any		= model.get_bias()
        assert name is None




@require_tf
class    a (       unittest.TestCase				):



  """simple docstring"""







  @slow
  def    lowercase_  (							self					):

    '''simple docstring'''







    __UpperCAmelCase:					Optional[int]		= TFEsmForMaskedLM.from_pretrained("""facebook/esm2_t6_8M_UR50D"""					)

    __UpperCAmelCase:					Any		= tf.constant([[0, 1, 2, 3, 4, 5]]					)
    __UpperCAmelCase:					Optional[Any]		= model(snake_case_					)[0]
    __UpperCAmelCase:					Union[str, Any]		= [1, 6, 33]
    self.assertEqual(list(output.numpy().shape					)   ,    snake_case_					)
    # compare the actual values for a slice.
    __UpperCAmelCase:					Any		= tf.constant(
        [
            [
                [8.9_2_1_5_1_8, -1_0.5_8_9_8_1_4, -6.4_6_7_1_3_0_7],
                [-6.3_9_6_7_1_5_6, -1_3.9_1_1_3_7_7, -1.1_2_1_1_9_1_5],
                [-7.7_8_1_2_4_7, -1_3.9_5_1_5_5_7, -3.7_4_0_5_9_2],
            ]
        ]					)
    self.assertTrue(numpy.allclose(output[:, :3, :3].numpy()   ,    expected_slice.numpy()   ,    atol=1e-2					)					)







  @slow
  def    lowercase_  (							self					):

    '''simple docstring'''







    __UpperCAmelCase:					Tuple		= TFEsmModel.from_pretrained("""facebook/esm2_t6_8M_UR50D"""					)

    __UpperCAmelCase:					Optional[Any]		= 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:					Any		= tf.constant(
        [
            [
                [0.1_4_4_4_3_0_9_2, 0.5_4_1_2_5_3_2_7, 0.3_2_4_7_7_3_9],
                [0.3_0_3_4_0_4_8_4, 0.0_0_5_2_6_6_7_6, 0.3_1_0_7_7_7_2_2],
                [0.3_2_2_7_8_0_4_3, -0.2_4_9_8_7_0_9_6, 0.3_4_1_4_6_2_8],
            ]
        ]					)
    self.assertTrue(numpy.allclose(output[:, :3, :3].numpy()   ,    expected_slice.numpy()   ,    atol=1e-4					)					)
'''simple docstring'''






from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_flax_available,
    is_tf_available,
    is_tokenizers_available,
    is_torch_available,
)


SCREAMING_SNAKE_CASE_	   =			{
    'configuration_electra': ['ELECTRA_PRETRAINED_CONFIG_ARCHIVE_MAP', 'ElectraConfig', 'ElectraOnnxConfig'],
    'tokenization_electra': ['ElectraTokenizer'],
}

try:
    if not is_tokenizers_available():
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
    pass
else:
    SCREAMING_SNAKE_CASE_	   =			['ElectraTokenizerFast']

try:
    if not is_torch_available():
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
    pass
else:
    SCREAMING_SNAKE_CASE_	   =			[
        'ELECTRA_PRETRAINED_MODEL_ARCHIVE_LIST',
        'ElectraForCausalLM',
        'ElectraForMaskedLM',
        'ElectraForMultipleChoice',
        'ElectraForPreTraining',
        'ElectraForQuestionAnswering',
        'ElectraForSequenceClassification',
        'ElectraForTokenClassification',
        'ElectraModel',
        'ElectraPreTrainedModel',
        'load_tf_weights_in_electra',
    ]

try:
    if not is_tf_available():
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
    pass
else:
    SCREAMING_SNAKE_CASE_	   =			[
        'TF_ELECTRA_PRETRAINED_MODEL_ARCHIVE_LIST',
        'TFElectraForMaskedLM',
        'TFElectraForMultipleChoice',
        'TFElectraForPreTraining',
        'TFElectraForQuestionAnswering',
        'TFElectraForSequenceClassification',
        'TFElectraForTokenClassification',
        'TFElectraModel',
        'TFElectraPreTrainedModel',
    ]

try:
    if not is_flax_available():
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
    pass
else:
    SCREAMING_SNAKE_CASE_	   =			[
        'FlaxElectraForCausalLM',
        'FlaxElectraForMaskedLM',
        'FlaxElectraForMultipleChoice',
        'FlaxElectraForPreTraining',
        'FlaxElectraForQuestionAnswering',
        'FlaxElectraForSequenceClassification',
        'FlaxElectraForTokenClassification',
        'FlaxElectraModel',
        'FlaxElectraPreTrainedModel',
    ]


if TYPE_CHECKING:
    from .configuration_electra import ELECTRA_PRETRAINED_CONFIG_ARCHIVE_MAP, ElectraConfig, ElectraOnnxConfig
    from .tokenization_electra import ElectraTokenizer

    try:
        if not is_tokenizers_available():
            raise OptionalDependencyNotAvailable()
    except OptionalDependencyNotAvailable:
        pass
    else:
        from .tokenization_electra_fast import ElectraTokenizerFast

    try:
        if not is_torch_available():
            raise OptionalDependencyNotAvailable()
    except OptionalDependencyNotAvailable:
        pass
    else:
        from .modeling_electra import (
            ELECTRA_PRETRAINED_MODEL_ARCHIVE_LIST,
            ElectraForCausalLM,
            ElectraForMaskedLM,
            ElectraForMultipleChoice,
            ElectraForPreTraining,
            ElectraForQuestionAnswering,
            ElectraForSequenceClassification,
            ElectraForTokenClassification,
            ElectraModel,
            ElectraPreTrainedModel,
            load_tf_weights_in_electra,
        )

    try:
        if not is_tf_available():
            raise OptionalDependencyNotAvailable()
    except OptionalDependencyNotAvailable:
        pass
    else:
        from .modeling_tf_electra import (
            TF_ELECTRA_PRETRAINED_MODEL_ARCHIVE_LIST,
            TFElectraForMaskedLM,
            TFElectraForMultipleChoice,
            TFElectraForPreTraining,
            TFElectraForQuestionAnswering,
            TFElectraForSequenceClassification,
            TFElectraForTokenClassification,
            TFElectraModel,
            TFElectraPreTrainedModel,
        )

    try:
        if not is_flax_available():
            raise OptionalDependencyNotAvailable()
    except OptionalDependencyNotAvailable:
        pass
    else:
        from .modeling_flax_electra import (
            FlaxElectraForCausalLM,
            FlaxElectraForMaskedLM,
            FlaxElectraForMultipleChoice,
            FlaxElectraForPreTraining,
            FlaxElectraForQuestionAnswering,
            FlaxElectraForSequenceClassification,
            FlaxElectraForTokenClassification,
            FlaxElectraModel,
            FlaxElectraPreTrainedModel,
        )

else:
    import sys

    SCREAMING_SNAKE_CASE_	   =			_LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)