



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






								"""simple docstring"""


								A__	  =     parent
								A__	  =     13
								A__	  =     7
								A__	  =     True
								A__	  =     True
								A__	  =     True
								A__	  =     99
								A__	  =     32
								A__	  =     2
								A__	  =     4
								A__	  =     37
								A__	  =     'gelu'
								A__	  =     0.1
								A__	  =     0.1
								A__	  =     5_12
								A__	  =     16
								A__	  =     2
								A__	  =     0.02
								A__	  =     3
								A__	  =     4
								A__	  =     None



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






								"""simple docstring"""


								A__	  =     ids_tensor([self.batch_size, self.seq_length]					,       self.vocab_size       )

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

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

								A__	  =     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      _a				(  self       :							Optional[Any]       ):






								"""simple docstring"""


								(
								    (
								    A__
								)  ,			(
								    A__
								)  ,			(
								    A__
								)  ,			(
								    A__
								)  ,			(
								    A__
								)  ,			(
								    A__
								)  ,			
								)	  =     self.prepare_config_and_inputs()

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






								"""simple docstring"""


								A__	  =     TFEsmModel(config=_snake_case       )
								A__	  =     {'input_ids': input_ids, 'attention_mask': input_mask}
								A__	  =     model(_snake_case       )

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

								A__	  =     model(_snake_case       )

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



					def      _a				(  self       :							str					,       _snake_case       :							int					,       _snake_case       :							Optional[Any]					,       _snake_case       :							Union[str, Any]					,       _snake_case       :							List[Any]					,       _snake_case       :							List[Any]					,       _snake_case       :							Union[str, Any]					,       _snake_case       :							List[Any]					,       _snake_case       :							List[str]					,       ):






								"""simple docstring"""


								A__	  =     True

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

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

								# Also check the case where encoder outputs are not passed
								A__	  =     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      _a				(  self       :							str					,       _snake_case       :							Any					,       _snake_case       :							Union[str, Any]					,       _snake_case       :							Union[str, Any]					,       _snake_case       :							Dict					,       _snake_case       :							Optional[int]					,       _snake_case       :							List[Any]       ):






								"""simple docstring"""


								A__	  =     TFEsmForMaskedLM(config=_snake_case       )
								A__	  =     model([input_ids, input_mask]       )
								self.parent.assertEqual(result.logits.shape					,       (self.batch_size, self.seq_length, self.vocab_size)       )



					def      _a				(  self       :							int					,       _snake_case       :							Optional[Any]					,       _snake_case       :							List[Any]					,       _snake_case       :							int					,       _snake_case       :							Tuple					,       _snake_case       :							Any					,       _snake_case       :							List[str]       ):






								"""simple docstring"""


								A__	  =     self.num_labels
								A__	  =     TFEsmForTokenClassification(config=_snake_case       )
								A__	  =     {'input_ids': input_ids, 'attention_mask': input_mask}
								A__	  =     model(_snake_case       )
								self.parent.assertEqual(result.logits.shape					,       (self.batch_size, self.seq_length, self.num_labels)       )



					def      _a				(  self       :							Dict       ):






								"""simple docstring"""


								A__	  =     self.prepare_config_and_inputs()
								(
								    (
								    A__
								)  ,			(
								    A__
								)  ,			(
								    A__
								)  ,			(
								    A__
								)  ,			(
								    A__
								)  ,			(
								    A__
								)  ,			
								)	  =     config_and_inputs
								A__	  =     {'input_ids': input_ids, 'attention_mask': input_mask}
								return config, inputs_dict







@require_tf
class   __lowerCAmelCase  (		UpperCAmelCase_				,  UpperCAmelCase_				,  unittest.TestCase   ):
					"""simple docstring"""






					A__					:		int      							=     (
					    (
					        TFEsmModel,
					        TFEsmForMaskedLM,
					        TFEsmForSequenceClassification,
					        TFEsmForTokenClassification,
					    )
					    if is_tf_available()
					    else ()
					)
					A__					:		Dict      							=     (
					    {
					        "feature-extraction": TFEsmModel,
					        "fill-mask": TFEsmForMaskedLM,
					        "text-classification": TFEsmForSequenceClassification,
					        "token-classification": TFEsmForTokenClassification,
					        "zero-shot": TFEsmForSequenceClassification,
					    }
					    if is_tf_available()
					    else {}
					)
					A__					:		Tuple      							=     False
					A__					:		Optional[int]      							=     False



					def      _a				(  self       :							Dict       ):






								"""simple docstring"""


								A__	  =     TFEsmModelTester(self       )
								A__	  =     ConfigTester(self					,       config_class=_snake_case					,       hidden_size=37       )



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






								"""simple docstring"""


								self.config_tester.run_common_tests()



					def      _a				(  self       :							int       ):






								"""simple docstring"""


								A__	  =     self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_model(*_snake_case       )



					def      _a				(  self       :							str       ):






								"""simple docstring"""


								A__	  =     self.model_tester.prepare_config_and_inputs_for_decoder()
								self.model_tester.create_and_check_model_as_decoder(*_snake_case       )



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






								"""simple docstring"""


								A__	  =     self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_masked_lm(*_snake_case       )



					def      _a				(  self       :							Optional[Any]       ):






								"""simple docstring"""


								A__	  =     self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_token_classification(*_snake_case       )



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






								"""simple docstring"""


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



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






								"""simple docstring"""


								pass



					@unittest.skip('Protein models do not support embedding resizing.'       )
					def      _a				(  self       :							Union[str, Any]       ):






								"""simple docstring"""


								pass



					def      _a				(  self       :							int       ):






								"""simple docstring"""


								A__  ,			A__	  =     self.model_tester.prepare_config_and_inputs_for_common()

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







@require_tf
class   __lowerCAmelCase  (		unittest.TestCase   ):
					"""simple docstring"""






					@slow
					def      _a				(  self       :							int       ):






								"""simple docstring"""


								A__	  =     TFEsmForMaskedLM.from_pretrained('facebook/esm2_t6_8M_UR50D'       )

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






								"""simple docstring"""


								A__	  =     TFEsmModel.from_pretrained('facebook/esm2_t6_8M_UR50D'       )

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

import tensorflow as tf
from keras_nlp.tokenizers import BytePairTokenizer
from tensorflow_text import pad_model_inputs

from .tokenization_gpta import GPTaTokenizer







class      SCREAMING_SNAKE_CASE    (						tf.keras.layers.Layer				):






      def __init__( self   :  List[Any]						,						lowercase__   :  Dict[str, int]						,						lowercase__   :  List[str]						,						lowercase__   :  int = None						,						lowercase__   :  int = None       ):


           '''simple docstring'''

           super().__init__()
           a_     :							Optional[Any]     						=  pad_token_id
           a_     :							str     						=  max_length
           a_     :							List[str]     						=  vocab
           a_     :							Union[str, Any]     						=  merges
           a_     :							Union[str, Any]     						=  BytePairTokenizer(lowercase__						,						lowercase__						,						sequence_length=lowercase__       )






      @classmethod
      def 							lowercase_						( cls   :  List[Any]						,						lowercase__   :  GPTaTokenizer						,						*lowercase__   :  Optional[Any]						,						**lowercase__   :  int       ):


           '''simple docstring'''

           a_     :							str     						=  [""" """.join(lowercase__       ) for m in tokenizer.bpe_ranks.keys()]
           a_     :							Any     						=  tokenizer.get_vocab()
           return cls(lowercase__						,						lowercase__						,						*lowercase__						,						**lowercase__       )






      @classmethod
      def 							lowercase_						( cls   :  Dict						,						lowercase__   :  Union[str, os.PathLike]						,						*lowercase__   :  str						,						**lowercase__   :  Optional[int]       ):


           '''simple docstring'''

           a_     :							Optional[int]     						=  GPTaTokenizer.from_pretrained(lowercase__						,						*lowercase__						,						**lowercase__       )
           return cls.from_tokenizer(lowercase__						,						*lowercase__						,						**lowercase__       )






      @classmethod
      def 							lowercase_						( cls   :  Dict						,						lowercase__   :  Optional[Any]       ):


           '''simple docstring'''

           return cls(**lowercase__       )






      def 							lowercase_						( self   :  Optional[Any]       ):


           '''simple docstring'''

           return {
               "vocab": self.vocab,
               "merges": self.merges,
               "max_length": self.max_length,
               "pad_token_id": self.pad_token_id,
           }






      def 							lowercase_						( self   :  Dict						,						lowercase__   :  Union[str, Any]						,						lowercase__   :  int = None       ):


           '''simple docstring'''

           a_     :							List[str]     						=  self.tf_tokenizer(lowercase__       )
           a_     :							int     						=  tf.ones_like(lowercase__       )

           if self.pad_token_id is not None:
                # pad the tokens up to max length
                a_     :							int     						=  max_length if max_length is not None else self.max_length

                if max_length is not None:
                     a_						,     a_     :							int     						=  pad_model_inputs(
                         lowercase__						,						max_seq_length=lowercase__						,						pad_value=self.pad_token_id       )

           return {"attention_mask": attention_mask, "input_ids": input_ids}
