




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






     def __init__(					self   ,_SCREAMING_SNAKE_CASE   ,)				->     str:
          _snake_case				= parent
          _snake_case				= 13
          _snake_case				= 7
          _snake_case				= True
          _snake_case				= True
          _snake_case				= True
          _snake_case				= 99
          _snake_case				= 32
          _snake_case				= 2
          _snake_case				= 4
          _snake_case				= 37
          _snake_case				= "gelu"
          _snake_case				= 0.1
          _snake_case				= 0.1
          _snake_case				= 512
          _snake_case				= 16
          _snake_case				= 2
          _snake_case				= 0.0_2
          _snake_case				= 3
          _snake_case				= 4
          _snake_case				= None






     def 							_lowercase      (					self       )				->     str:
          _snake_case				= ids_tensor([self.batch_size, self.seq_length]   ,self.vocab_size       )

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

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

          _snake_case				= 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       )				->     Any:
          (
              (
              _snake_case
          )				,   (
              _snake_case
          )				,   (
              _snake_case
          )				,   (
              _snake_case
          )				,   (
              _snake_case
          )				,   (
              _snake_case
          )				,   
          )				= self.prepare_config_and_inputs()

          _snake_case				= True
          _snake_case				= floats_tensor([self.batch_size, self.seq_length, self.hidden_size]       )
          _snake_case				= 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   ,_SCREAMING_SNAKE_CASE   ,_SCREAMING_SNAKE_CASE   ,_SCREAMING_SNAKE_CASE   ,_SCREAMING_SNAKE_CASE   ,_SCREAMING_SNAKE_CASE   ,_SCREAMING_SNAKE_CASE       )				->     Any:
          _snake_case				= TFEsmModel(config=_SCREAMING_SNAKE_CASE       )
          _snake_case				= {"input_ids": input_ids, "attention_mask": input_mask}
          _snake_case				= model(_SCREAMING_SNAKE_CASE       )

          _snake_case				= [input_ids, input_mask]
          _snake_case				= model(_SCREAMING_SNAKE_CASE       )

          _snake_case				= model(_SCREAMING_SNAKE_CASE       )

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






     def 							_lowercase      (					self   ,_SCREAMING_SNAKE_CASE   ,_SCREAMING_SNAKE_CASE   ,_SCREAMING_SNAKE_CASE   ,_SCREAMING_SNAKE_CASE   ,_SCREAMING_SNAKE_CASE   ,_SCREAMING_SNAKE_CASE   ,_SCREAMING_SNAKE_CASE   ,_SCREAMING_SNAKE_CASE   ,)				->     List[str]:
          _snake_case				= True

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

          _snake_case				= [input_ids, input_mask]
          _snake_case				= model(_SCREAMING_SNAKE_CASE   ,encoder_hidden_states=_SCREAMING_SNAKE_CASE       )

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

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






     def 							_lowercase      (					self   ,_SCREAMING_SNAKE_CASE   ,_SCREAMING_SNAKE_CASE   ,_SCREAMING_SNAKE_CASE   ,_SCREAMING_SNAKE_CASE   ,_SCREAMING_SNAKE_CASE   ,_SCREAMING_SNAKE_CASE       )				->     Any:
          _snake_case				= TFEsmForMaskedLM(config=_SCREAMING_SNAKE_CASE       )
          _snake_case				= model([input_ids, input_mask]       )
          self.parent.assertEqual(result.logits.shape   ,(self.batch_size, self.seq_length, self.vocab_size)       )






     def 							_lowercase      (					self   ,_SCREAMING_SNAKE_CASE   ,_SCREAMING_SNAKE_CASE   ,_SCREAMING_SNAKE_CASE   ,_SCREAMING_SNAKE_CASE   ,_SCREAMING_SNAKE_CASE   ,_SCREAMING_SNAKE_CASE       )				->     Optional[int]:
          _snake_case				= self.num_labels
          _snake_case				= TFEsmForTokenClassification(config=_SCREAMING_SNAKE_CASE       )
          _snake_case				= {"input_ids": input_ids, "attention_mask": input_mask}
          _snake_case				= model(_SCREAMING_SNAKE_CASE       )
          self.parent.assertEqual(result.logits.shape   ,(self.batch_size, self.seq_length, self.num_labels)       )






     def 							_lowercase      (					self       )				->     int:
          _snake_case				= self.prepare_config_and_inputs()
          (
              (
              _snake_case
          )				,   (
              _snake_case
          )				,   (
              _snake_case
          )				,   (
              _snake_case
          )				,   (
              _snake_case
          )				,   (
              _snake_case
          )				,   
          )				= config_and_inputs
          _snake_case				= {"input_ids": input_ids, "attention_mask": input_mask}
          return config, inputs_dict







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






     def 							_lowercase      (					self       )				->     int:
          _snake_case				= TFEsmModelTester(self       )
          _snake_case				= ConfigTester(self   ,config_class=_SCREAMING_SNAKE_CASE   ,hidden_size=37       )






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






     def 							_lowercase      (					self       )				->     Any:
          _snake_case				= self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_model(*_SCREAMING_SNAKE_CASE       )






     def 							_lowercase      (					self       )				->     int:
          _snake_case				= self.model_tester.prepare_config_and_inputs_for_decoder()
          self.model_tester.create_and_check_model_as_decoder(*_SCREAMING_SNAKE_CASE       )






     def 							_lowercase      (					self       )				->     str:
          _snake_case				= self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_for_masked_lm(*_SCREAMING_SNAKE_CASE       )






     def 							_lowercase      (					self       )				->     List[Any]:
          _snake_case				= self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_for_token_classification(*_SCREAMING_SNAKE_CASE       )






     @slow
     def 							_lowercase      (					self       )				->     Optional[Any]:
          for model_name in TF_ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
               _snake_case				= TFEsmModel.from_pretrained(_SCREAMING_SNAKE_CASE       )
               self.assertIsNotNone(_SCREAMING_SNAKE_CASE       )






     @unittest.skip("Protein models do not support embedding resizing."       )
     def 							_lowercase      (					self       )				->     Optional[Any]:
          pass






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






     def 							_lowercase      (					self       )				->     int:
          _snake_case				,   _snake_case				= self.model_tester.prepare_config_and_inputs_for_common()

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







@require_tf
class   _a							(		unittest.TestCase				):






     @slow
     def 							_lowercase      (					self       )				->     Optional[int]:
          _snake_case				= TFEsmForMaskedLM.from_pretrained("facebook/esm2_t6_8M_UR50D"       )

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

          _snake_case				= tf.constant([[0, 6, 4, 13, 5, 4, 16, 12, 11, 7, 2]]       )
          _snake_case				= model(_SCREAMING_SNAKE_CASE       )[0]
          # compare the actual values for a slice.
          _snake_case				= 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 collections import OrderedDict
from typing import Mapping

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig
from ...utils import logging


UpperCamelCase_		:     Optional[int]          =		logging.get_logger(__name__)

UpperCamelCase_		:     str          =		{
    '''kssteven/ibert-roberta-base''': '''https://huggingface.co/kssteven/ibert-roberta-base/resolve/main/config.json''',
    '''kssteven/ibert-roberta-large''': '''https://huggingface.co/kssteven/ibert-roberta-large/resolve/main/config.json''',
    '''kssteven/ibert-roberta-large-mnli''': (
        '''https://huggingface.co/kssteven/ibert-roberta-large-mnli/resolve/main/config.json'''
    ),
}

class   _a							(		__lowerCAmelCase				):
     SCREAMING_SNAKE_CASE_				:       List[Any]       							=		"""ibert"""






     def __init__(					self   ,_SCREAMING_SNAKE_CASE=30_522   ,_SCREAMING_SNAKE_CASE=768   ,_SCREAMING_SNAKE_CASE=12   ,_SCREAMING_SNAKE_CASE=12   ,_SCREAMING_SNAKE_CASE=3_072   ,_SCREAMING_SNAKE_CASE="gelu"   ,_SCREAMING_SNAKE_CASE=0.1   ,_SCREAMING_SNAKE_CASE=0.1   ,_SCREAMING_SNAKE_CASE=512   ,_SCREAMING_SNAKE_CASE=2   ,_SCREAMING_SNAKE_CASE=0.0_2   ,_SCREAMING_SNAKE_CASE=1e-12   ,_SCREAMING_SNAKE_CASE=1   ,_SCREAMING_SNAKE_CASE=0   ,_SCREAMING_SNAKE_CASE=2   ,_SCREAMING_SNAKE_CASE="absolute"   ,_SCREAMING_SNAKE_CASE=False   ,_SCREAMING_SNAKE_CASE="none"   ,**_SCREAMING_SNAKE_CASE   ,)				->     Dict:
          super().__init__(pad_token_id=_SCREAMING_SNAKE_CASE   ,bos_token_id=_SCREAMING_SNAKE_CASE   ,eos_token_id=_SCREAMING_SNAKE_CASE   ,**_SCREAMING_SNAKE_CASE       )

          _snake_case				= vocab_size
          _snake_case				= hidden_size
          _snake_case				= num_hidden_layers
          _snake_case				= num_attention_heads
          _snake_case				= hidden_act
          _snake_case				= intermediate_size
          _snake_case				= hidden_dropout_prob
          _snake_case				= attention_probs_dropout_prob
          _snake_case				= max_position_embeddings
          _snake_case				= type_vocab_size
          _snake_case				= initializer_range
          _snake_case				= layer_norm_eps
          _snake_case				= position_embedding_type
          _snake_case				= quant_mode
          _snake_case				= force_dequant







class   _a							(		__lowerCAmelCase				):






     @property
     def 							_lowercase      (					self       )				->     Mapping[str, Mapping[int, str]]:
          if self.task == "multiple-choice":
               _snake_case				= {0: "batch", 1: "choice", 2: "sequence"}
          else:
               _snake_case				= {0: "batch", 1: "sequence"}
          return OrderedDict(
              [
                  ("input_ids", dynamic_axis),
                  ("attention_mask", dynamic_axis),
              ]       )


