'''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       :       List[Any]  ,     __a       :       str  ,     ):
         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    _lowercase		(self       :       int					):
         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    _lowercase		(self       :       Optional[int]					):
         (
             (
             UpperCAmelCase_   
         )       ,					(
             UpperCAmelCase_   
         )       ,					(
             UpperCAmelCase_   
         )       ,					(
             UpperCAmelCase_   
         )       ,					(
             UpperCAmelCase_   
         )       ,					(
             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    _lowercase		(self       :       Tuple  ,     __a       :       List[Any]  ,     __a       :       List[Any]  ,     __a       :       int  ,     __a       :       str  ,     __a       :       Optional[Any]  ,     __a       :       Dict					):
         UpperCAmelCase_			       =  TFEsmModel(config=_SCREAMING_SNAKE_CASE					)
         UpperCAmelCase_			       =  {"input_ids": input_ids, "attention_mask": input_mask}
         UpperCAmelCase_			       =  model(_SCREAMING_SNAKE_CASE					)

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

         UpperCAmelCase_			       =  model(_SCREAMING_SNAKE_CASE					)

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



    def    _lowercase		(self       :       int  ,     __a       :       Optional[int]  ,     __a       :       Optional[Any]  ,     __a       :       Any  ,     __a       :       Any  ,     __a       :       int  ,     __a       :       Any  ,     __a       :       List[str]  ,     __a       :       Dict  ,     ):
         UpperCAmelCase_			       =  True

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

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

         # Also check the case where encoder outputs are not passed
         UpperCAmelCase_			       =  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       :       Union[str, Any]  ,     __a       :       Optional[int]  ,     __a       :       List[Any]  ,     __a       :       Dict  ,     __a       :       Tuple  ,     __a       :       int  ,     __a       :       Optional[int]					):
         UpperCAmelCase_			       =  TFEsmForMaskedLM(config=_SCREAMING_SNAKE_CASE					)
         UpperCAmelCase_			       =  model([input_ids, input_mask]					)
         self.parent.assertEqual(result.logits.shape  ,     (self.batch_size, self.seq_length, self.vocab_size)					)



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






    def    _lowercase		(self       :       Union[str, Any]					):
         UpperCAmelCase_			       =  self.prepare_config_and_inputs()
         (
             (
             UpperCAmelCase_   
         )       ,					(
             UpperCAmelCase_   
         )       ,					(
             UpperCAmelCase_   
         )       ,					(
             UpperCAmelCase_   
         )       ,					(
             UpperCAmelCase_   
         )       ,					(
             UpperCAmelCase_   
         )       ,					   
         )			       =  config_and_inputs
         UpperCAmelCase_			       =  {"input_ids": input_ids, "attention_mask": input_mask}
         return config, inputs_dict




@require_tf
class 			__A (					__lowerCAmelCase						,      __lowerCAmelCase						,      unittest.TestCase      ):
    a__			:		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__			:		str			  =    False
    a__			:		Optional[Any]			  =    False



    def    _lowercase		(self       :       Dict					):
         UpperCAmelCase_			       =  TFEsmModelTester(self					)
         UpperCAmelCase_			       =  ConfigTester(self  ,     config_class=_SCREAMING_SNAKE_CASE  ,     hidden_size=37					)



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



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



    def    _lowercase		(self       :       int					):
         UpperCAmelCase_			       =  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					):
         UpperCAmelCase_			       =  self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_masked_lm(*_SCREAMING_SNAKE_CASE					)



    def    _lowercase		(self       :       Dict					):
         UpperCAmelCase_			       =  self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_token_classification(*_SCREAMING_SNAKE_CASE					)



    @slow
    def    _lowercase		(self       :       Dict					):
         for model_name in TF_ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
              UpperCAmelCase_			       =  TFEsmModel.from_pretrained(_SCREAMING_SNAKE_CASE					)
              self.assertIsNotNone(_SCREAMING_SNAKE_CASE					)



    @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       :       Tuple					):
         pass






    def    _lowercase		(self       :       Optional[Any]					):
         UpperCAmelCase_       ,					UpperCAmelCase_			       =  self.model_tester.prepare_config_and_inputs_for_common()

         for model_class in self.all_model_classes:
              UpperCAmelCase_			       =  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
                   UpperCAmelCase_			       =  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:
                   UpperCAmelCase_			       =  model.get_output_embeddings()
                   assert x is None
                   UpperCAmelCase_			       =  model.get_bias()
                   assert name is None




@require_tf
class 			__A (					unittest.TestCase      ):



    @slow
    def    _lowercase		(self       :       str					):
         UpperCAmelCase_			       =  TFEsmForMaskedLM.from_pretrained("facebook/esm2_t6_8M_UR50D"					)

         UpperCAmelCase_			       =  tf.constant([[0, 1, 2, 3, 4, 5]]					)
         UpperCAmelCase_			       =  model(_SCREAMING_SNAKE_CASE					)[0]
         UpperCAmelCase_			       =  [1, 6, 33]
         self.assertEqual(list(output.numpy().shape					)  ,     _SCREAMING_SNAKE_CASE					)
         # compare the actual values for a slice.
         UpperCAmelCase_			       =  tf.constant(
             [
                 [
                     [8.92_15_18, -10.589814, -6.4_67_13_07],
                     [-6.3_96_71_56, -13.911377, -1.1_21_19_15],
                     [-7.78_12_47, -13.951557, -3.74_05_92],
                 ]
             ]					)
         self.assertTrue(numpy.allclose(output[:, :3, :3].numpy()  ,     expected_slice.numpy()  ,     atol=1E-2					)					)






    @slow
    def    _lowercase		(self       :       Any					):
         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(_SCREAMING_SNAKE_CASE					)[0]
         # compare the actual values for a slice.
         UpperCAmelCase_			       =  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					)					)










'''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),
              ]       )


