
"""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   _lowercase						:


     '''simple docstring'''





     def __init__(      self ,			__UpperCamelCase ,			)->					str:
          UpperCAmelCase__							: Optional[Any] 							=						parent
          UpperCAmelCase__							: Optional[int] 							=						13
          UpperCAmelCase__							: Union[str, Any] 							=						7
          UpperCAmelCase__							: str 							=						True
          UpperCAmelCase__							: Optional[int] 							=						True
          UpperCAmelCase__							: Dict 							=						True
          UpperCAmelCase__							: List[str] 							=						99
          UpperCAmelCase__							: Optional[int] 							=						32
          UpperCAmelCase__							: Dict 							=						2
          UpperCAmelCase__							: List[str] 							=						4
          UpperCAmelCase__							: List[str] 							=						37
          UpperCAmelCase__							: List[Any] 							=						"gelu"
          UpperCAmelCase__							: List[Any] 							=						0.1
          UpperCAmelCase__							: List[Any] 							=						0.1
          UpperCAmelCase__							: List[Any] 							=						5_12
          UpperCAmelCase__							: Tuple 							=						16
          UpperCAmelCase__							: int 							=						2
          UpperCAmelCase__							: Dict 							=						0.02
          UpperCAmelCase__							: Optional[Any] 							=						3
          UpperCAmelCase__							: str 							=						4
          UpperCAmelCase__							: List[Any] 							=						None



     def  lowerCAmelCase__       (      self  )->					List[str]:
          UpperCAmelCase__							: Tuple 							=						ids_tensor([self.batch_size, self.seq_length] ,			self.vocab_size  )

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

          UpperCAmelCase__							: Optional[Any] 							=						None
          UpperCAmelCase__							: Optional[int] 							=						None
          UpperCAmelCase__							: Tuple 							=						None
          if self.use_labels:
               UpperCAmelCase__							: Any 							=						ids_tensor([self.batch_size] ,			self.type_sequence_label_size  )
               UpperCAmelCase__							: Optional[int] 							=						ids_tensor([self.batch_size, self.seq_length] ,			self.num_labels  )
               UpperCAmelCase__							: List[str] 							=						ids_tensor([self.batch_size] ,			self.num_choices  )

          UpperCAmelCase__							: Optional[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  lowerCAmelCase__       (      self  )->					Optional[int]:
          (
              (
              UpperCAmelCase__
          )							,	(
              UpperCAmelCase__
          )							,	(
              UpperCAmelCase__
          )							,	(
              UpperCAmelCase__
          )							,	(
              UpperCAmelCase__
          )							,	(
              UpperCAmelCase__
          )							,	
          )							: List[str] 							=						self.prepare_config_and_inputs()

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

          UpperCAmelCase__							: List[str] 							=						[input_ids, input_mask]
          UpperCAmelCase__							: Tuple 							=						model(__UpperCamelCase  )

          UpperCAmelCase__							: Optional[Any] 							=						model(__UpperCamelCase  )

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



     def  lowerCAmelCase__       (      self ,			__UpperCamelCase ,			__UpperCamelCase ,			__UpperCamelCase ,			__UpperCamelCase ,			__UpperCamelCase ,			__UpperCamelCase ,			__UpperCamelCase ,			__UpperCamelCase ,			)->					str:
          UpperCAmelCase__							: List[str] 							=						True

          UpperCAmelCase__							: List[str] 							=						TFEsmModel(config=__UpperCamelCase  )
          UpperCAmelCase__							: Union[str, Any] 							=						{
              "input_ids": input_ids,
              "attention_mask": input_mask,
              "encoder_hidden_states": encoder_hidden_states,
              "encoder_attention_mask": encoder_attention_mask,
          }
          UpperCAmelCase__							: Dict 							=						model(__UpperCamelCase  )

          UpperCAmelCase__							: Dict 							=						[input_ids, input_mask]
          UpperCAmelCase__							: Tuple 							=						model(__UpperCamelCase ,			encoder_hidden_states=__UpperCamelCase  )

          # Also check the case where encoder outputs are not passed
          UpperCAmelCase__							: List[Any] 							=						model(__UpperCamelCase ,			attention_mask=__UpperCamelCase  )

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



     def  lowerCAmelCase__       (      self ,			__UpperCamelCase ,			__UpperCamelCase ,			__UpperCamelCase ,			__UpperCamelCase ,			__UpperCamelCase ,			__UpperCamelCase  )->					Dict:
          UpperCAmelCase__							: str 							=						TFEsmForMaskedLM(config=__UpperCamelCase  )
          UpperCAmelCase__							: Any 							=						model([input_ids, input_mask]  )
          self.parent.assertEqual(result.logits.shape ,			(self.batch_size, self.seq_length, self.vocab_size)  )



     def  lowerCAmelCase__       (      self ,			__UpperCamelCase ,			__UpperCamelCase ,			__UpperCamelCase ,			__UpperCamelCase ,			__UpperCamelCase ,			__UpperCamelCase  )->					Any:
          UpperCAmelCase__							: Dict 							=						self.num_labels
          UpperCAmelCase__							: Union[str, Any] 							=						TFEsmForTokenClassification(config=__UpperCamelCase  )
          UpperCAmelCase__							: int 							=						{"input_ids": input_ids, "attention_mask": input_mask}
          UpperCAmelCase__							: Tuple 							=						model(__UpperCamelCase  )
          self.parent.assertEqual(result.logits.shape ,			(self.batch_size, self.seq_length, self.num_labels)  )



     def  lowerCAmelCase__       (      self  )->					int:
          UpperCAmelCase__							: Union[str, Any] 							=						self.prepare_config_and_inputs()
          (
              (
              UpperCAmelCase__
          )							,	(
              UpperCAmelCase__
          )							,	(
              UpperCAmelCase__
          )							,	(
              UpperCAmelCase__
          )							,	(
              UpperCAmelCase__
          )							,	(
              UpperCAmelCase__
          )							,	
          )							: str 							=						config_and_inputs
          UpperCAmelCase__							: str 							=						{"input_ids": input_ids, "attention_mask": input_mask}
          return config, inputs_dict







@require_tf
class   _lowercase						(       lowerCAmelCase_	,   lowerCAmelCase_	,   unittest.TestCase				):


     '''simple docstring'''





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



     def  lowerCAmelCase__       (      self  )->					Dict:
          UpperCAmelCase__							: Optional[int] 							=						TFEsmModelTester(self  )
          UpperCAmelCase__							: List[Any] 							=						ConfigTester(self ,			config_class=__UpperCamelCase ,			hidden_size=37  )



     def  lowerCAmelCase__       (      self  )->					Optional[int]:
          self.config_tester.run_common_tests()



     def  lowerCAmelCase__       (      self  )->					str:
          UpperCAmelCase__							: Union[str, Any] 							=						self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_model(*__UpperCamelCase  )



     def  lowerCAmelCase__       (      self  )->					int:
          UpperCAmelCase__							: Optional[int] 							=						self.model_tester.prepare_config_and_inputs_for_decoder()
          self.model_tester.create_and_check_model_as_decoder(*__UpperCamelCase  )



     def  lowerCAmelCase__       (      self  )->					Tuple:
          UpperCAmelCase__							: Union[str, Any] 							=						self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_for_masked_lm(*__UpperCamelCase  )



     def  lowerCAmelCase__       (      self  )->					str:
          UpperCAmelCase__							: Dict 							=						self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_for_token_classification(*__UpperCamelCase  )



     @slow
     def  lowerCAmelCase__       (      self  )->					Union[str, Any]:
          for model_name in TF_ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
               UpperCAmelCase__							: Optional[int] 							=						TFEsmModel.from_pretrained(__UpperCamelCase  )
               self.assertIsNotNone(__UpperCamelCase  )



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



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



     def  lowerCAmelCase__       (      self  )->					Any:
          UpperCAmelCase__							,	UpperCAmelCase__							: List[str] 							=						self.model_tester.prepare_config_and_inputs_for_common()

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







@require_tf
class   _lowercase						(       unittest.TestCase				):


     '''simple docstring'''





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

          UpperCAmelCase__							: Dict 							=						tf.constant([[0, 1, 2, 3, 4, 5]]  )
          UpperCAmelCase__							: Tuple 							=						model(__UpperCamelCase  )[0]
          UpperCAmelCase__							: Optional[Any] 							=						[1, 6, 33]
          self.assertEqual(list(output.numpy().shape  ) ,			__UpperCamelCase  )
          # compare the actual values for a slice.
          UpperCAmelCase__							: Optional[int] 							=						tf.constant(
              [
                  [
                      [8.92_1518, -10.58_9814, -6.467_1307],
                      [-6.396_7156, -13.91_1377, -1.121_1915],
                      [-7.78_1247, -13.95_1557, -3.74_0592],
                  ]
              ]  )
          self.assertTrue(numpy.allclose(output[:, :3, :3].numpy() ,			expected_slice.numpy() ,			atol=1E-2  )  )



     @slow
     def  lowerCAmelCase__       (      self  )->					List[str]:
          UpperCAmelCase__							: Tuple 							=						TFEsmModel.from_pretrained("facebook/esm2_t6_8M_UR50D"  )

          UpperCAmelCase__							: Dict 							=						tf.constant([[0, 6, 4, 13, 5, 4, 16, 12, 11, 7, 2]]  )
          UpperCAmelCase__							: Optional[int] 							=						model(__UpperCamelCase  )[0]
          # compare the actual values for a slice.
          UpperCAmelCase__							: str 							=						tf.constant(
              [
                  [
                      [0.1444_3092, 0.5412_5327, 0.324_7739],
                      [0.3034_0484, 0.0052_6676, 0.3107_7722],
                      [0.3227_8043, -0.2498_7096, 0.341_4628],
                  ]
              ]  )
          self.assertTrue(numpy.allclose(output[:, :3, :3].numpy() ,			expected_slice.numpy() ,			atol=1E-4  )  )





"""simple docstring"""




import timeit

import numpy as np

import datasets
from datasets.arrow_writer import ArrowWriter
from datasets.features.features import _ArrayXD



def 							a__ (				lowerCAmelCase	:       List[str]		):






     '''simple docstring'''
     def wrapper(*lowerCAmelCase	:       Any       ,    **lowerCAmelCase	:       Tuple		):
          UpperCAmelCase__							: Optional[int] 							=						timeit.default_timer()
          UpperCAmelCase__							: int 							=						func(*lowerCAmelCase       ,    **lowerCAmelCase		)
          UpperCAmelCase__							: List[Any] 							=						timeit.default_timer() - starttime
          return delta

     UpperCAmelCase__							: int 							=						func.__name__

     return wrapper



def 							a__ (				lowerCAmelCase	:       dict       ,    lowerCAmelCase	:       Optional[int]=100       ,    lowerCAmelCase	:       List[str]=None		):






     '''simple docstring'''
     UpperCAmelCase__							: str 							=						[]
     UpperCAmelCase__							: Optional[Any] 							=						seq_shapes or {}
     for i in range(lowerCAmelCase		):
          UpperCAmelCase__							: int 							=						{}
          for col_id, (k, v) in enumerate(features.items()		):
               if isinstance(lowerCAmelCase       ,    _ArrayXD		):
                    UpperCAmelCase__							: List[str] 							=						np.random.rand(*v.shape		).astype(v.dtype		)
               elif isinstance(lowerCAmelCase       ,    datasets.Value		):
                    if v.dtype == "string":
                         UpperCAmelCase__							: Dict 							=						"The small grey turtle was surprisingly fast when challenged."
                    else:
                         UpperCAmelCase__							: str 							=						np.random.randint(10       ,    size=1		).astype(v.dtype		).item()
               elif isinstance(lowerCAmelCase       ,    datasets.Sequence		):
                    while isinstance(lowerCAmelCase       ,    datasets.Sequence		):
                         UpperCAmelCase__							: List[str] 							=						v.feature
                    UpperCAmelCase__							: Optional[int] 							=						seq_shapes[k]
                    UpperCAmelCase__							: Optional[int] 							=						np.random.rand(*lowerCAmelCase		).astype(v.dtype		)
               UpperCAmelCase__							: Union[str, Any] 							=						data

          dummy_data.append((i, example)		)

     return dummy_data



def 							a__ (				lowerCAmelCase	:       List[str]       ,    lowerCAmelCase	:       Tuple       ,    lowerCAmelCase	:       List[str]=100       ,    lowerCAmelCase	:       Optional[int]=None		):






     '''simple docstring'''
     UpperCAmelCase__							: int 							=						generate_examples(lowerCAmelCase       ,    num_examples=lowerCAmelCase       ,    seq_shapes=lowerCAmelCase		)

     with ArrowWriter(features=lowerCAmelCase       ,    path=lowerCAmelCase		) as writer:
          for key, record in dummy_data:
               UpperCAmelCase__							: List[Any] 							=						features.encode_example(lowerCAmelCase		)
               writer.write(lowerCAmelCase		)

          UpperCAmelCase__							,	UpperCAmelCase__							: Optional[Any] 							=						writer.finalize()

     if not num_final_examples == num_examples:
          raise ValueError(
              F"Error writing the dataset, wrote {num_final_examples} examples but should have written {num_examples}."		)

     UpperCAmelCase__							: Optional[int] 							=						datasets.Dataset.from_file(filename=lowerCAmelCase       ,    info=datasets.DatasetInfo(features=lowerCAmelCase		)		)

     return dataset



