


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







						'''simple docstring'''




						def __init__( self      ,  __lowercase      ,  )    ->			Any:
										lowerCAmelCase_       :					List[str]			 = parent
										lowerCAmelCase_       :					Optional[Any]			 = 1_3
										lowerCAmelCase_       :					Tuple			 = 7
										lowerCAmelCase_       :					Any			 = True
										lowerCAmelCase_       :					Tuple			 = True
										lowerCAmelCase_       :					List[Any]			 = True
										lowerCAmelCase_       :					str			 = 9_9
										lowerCAmelCase_       :					Any			 = 3_2
										lowerCAmelCase_       :					int			 = 2
										lowerCAmelCase_       :					str			 = 4
										lowerCAmelCase_       :					int			 = 3_7
										lowerCAmelCase_       :					Optional[int]			 = '''gelu'''
										lowerCAmelCase_       :					str			 = 0.1
										lowerCAmelCase_       :					Dict			 = 0.1
										lowerCAmelCase_       :					Dict			 = 5_1_2
										lowerCAmelCase_       :					int			 = 1_6
										lowerCAmelCase_       :					Tuple			 = 2
										lowerCAmelCase_       :					List[str]			 = 0.02
										lowerCAmelCase_       :					str			 = 3
										lowerCAmelCase_       :					str			 = 4
										lowerCAmelCase_       :					Dict			 = None




						def   lowercase_       ( self							)    ->			Tuple:
										lowerCAmelCase_       :					str			 = ids_tensor([self.batch_size, self.seq_length]      ,  self.vocab_size							)

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

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

										lowerCAmelCase_       :					Dict			 = 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[Any]:
										(
										    (
										    lowerCAmelCase_	
										) ,     (
										    lowerCAmelCase_	
										) ,     (
										    lowerCAmelCase_	
										) ,     (
										    lowerCAmelCase_	
										) ,     (
										    lowerCAmelCase_	
										) ,     (
										    lowerCAmelCase_	
										) ,     	
										)       :					str			 = self.prepare_config_and_inputs()

										lowerCAmelCase_       :					Optional[int]			 = True
										lowerCAmelCase_       :					List[str]			 = floats_tensor([self.batch_size, self.seq_length, self.hidden_size]							)
										lowerCAmelCase_       :					Optional[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   lowercase_       ( self      ,  __lowercase      ,  __lowercase      ,  __lowercase      ,  __lowercase      ,  __lowercase      ,  __lowercase							)    ->			List[str]:
										lowerCAmelCase_       :					Optional[Any]			 = TFEsmModel(config=SCREAMING_SNAKE_CASE_							)
										lowerCAmelCase_       :					Union[str, Any]			 = {'''input_ids''': input_ids, '''attention_mask''': input_mask}
										lowerCAmelCase_       :					Optional[int]			 = model(SCREAMING_SNAKE_CASE_							)

										lowerCAmelCase_       :					Any			 = [input_ids, input_mask]
										lowerCAmelCase_       :					Dict			 = model(SCREAMING_SNAKE_CASE_							)

										lowerCAmelCase_       :					str			 = model(SCREAMING_SNAKE_CASE_							)

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




						def   lowercase_       ( self      ,  __lowercase      ,  __lowercase      ,  __lowercase      ,  __lowercase      ,  __lowercase      ,  __lowercase      ,  __lowercase      ,  __lowercase      ,  )    ->			Union[str, Any]:
										lowerCAmelCase_       :					Optional[Any]			 = True

										lowerCAmelCase_       :					List[str]			 = TFEsmModel(config=SCREAMING_SNAKE_CASE_							)
										lowerCAmelCase_       :					int			 = {
										    '''input_ids''': input_ids,
										    '''attention_mask''': input_mask,
										    '''encoder_hidden_states''': encoder_hidden_states,
										    '''encoder_attention_mask''': encoder_attention_mask,
										}
										lowerCAmelCase_       :					List[Any]			 = model(SCREAMING_SNAKE_CASE_							)

										lowerCAmelCase_       :					Optional[int]			 = [input_ids, input_mask]
										lowerCAmelCase_       :					Union[str, Any]			 = model(SCREAMING_SNAKE_CASE_      ,  encoder_hidden_states=SCREAMING_SNAKE_CASE_							)

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




						def   lowercase_       ( self      ,  __lowercase      ,  __lowercase      ,  __lowercase      ,  __lowercase      ,  __lowercase      ,  __lowercase							)    ->			int:
										lowerCAmelCase_       :					Union[str, Any]			 = self.num_labels
										lowerCAmelCase_       :					Optional[int]			 = TFEsmForTokenClassification(config=SCREAMING_SNAKE_CASE_							)
										lowerCAmelCase_       :					List[str]			 = {'''input_ids''': input_ids, '''attention_mask''': input_mask}
										lowerCAmelCase_       :					Tuple			 = model(SCREAMING_SNAKE_CASE_							)
										self.parent.assertEqual(result.logits.shape      ,  (self.batch_size, self.seq_length, self.num_labels)							)




						def   lowercase_       ( self							)    ->			str:
										lowerCAmelCase_       :					int			 = self.prepare_config_and_inputs()
										(
										    (
										    lowerCAmelCase_	
										) ,     (
										    lowerCAmelCase_	
										) ,     (
										    lowerCAmelCase_	
										) ,     (
										    lowerCAmelCase_	
										) ,     (
										    lowerCAmelCase_	
										) ,     (
										    lowerCAmelCase_	
										) ,     	
										)       :					Union[str, Any]			 = config_and_inputs
										lowerCAmelCase_       :					List[Any]			 = {'''input_ids''': input_ids, '''attention_mask''': input_mask}
										return config, inputs_dict







@require_tf
class  snake_case__(							UpperCAmelCase__,  UpperCAmelCase__,  unittest.TestCase							):







						'''simple docstring'''




						SCREAMING_SNAKE_CASE__					:     str    =							(
						    (
						        TFEsmModel,
						        TFEsmForMaskedLM,
						        TFEsmForSequenceClassification,
						        TFEsmForTokenClassification,
						    )
						    if is_tf_available()
						    else ()
						)
						SCREAMING_SNAKE_CASE__					:     Union[str, Any]    =							(
						    {
						        'feature-extraction': TFEsmModel,
						        'fill-mask': TFEsmForMaskedLM,
						        'text-classification': TFEsmForSequenceClassification,
						        'token-classification': TFEsmForTokenClassification,
						        'zero-shot': TFEsmForSequenceClassification,
						    }
						    if is_tf_available()
						    else {}
						)
						SCREAMING_SNAKE_CASE__					:     Union[str, Any]    =							False
						SCREAMING_SNAKE_CASE__					:     Dict    =							False




						def   lowercase_       ( self							)    ->			List[Any]:
										lowerCAmelCase_       :					List[Any]			 = TFEsmModelTester(self							)
										lowerCAmelCase_       :					Any			 = ConfigTester(self      ,  config_class=SCREAMING_SNAKE_CASE_      ,  hidden_size=3_7							)




						def   lowercase_       ( self							)    ->			str:
										self.config_tester.run_common_tests()




						def   lowercase_       ( self							)    ->			int:
										lowerCAmelCase_       :					Any			 = self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_model(*SCREAMING_SNAKE_CASE_							)




						def   lowercase_       ( self							)    ->			Union[str, Any]:
										lowerCAmelCase_       :					str			 = self.model_tester.prepare_config_and_inputs_for_decoder()
										self.model_tester.create_and_check_model_as_decoder(*SCREAMING_SNAKE_CASE_							)




						def   lowercase_       ( self							)    ->			Optional[int]:
										lowerCAmelCase_       :					Optional[int]			 = self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_for_masked_lm(*SCREAMING_SNAKE_CASE_							)




						def   lowercase_       ( self							)    ->			List[str]:
										lowerCAmelCase_       :					Union[str, Any]			 = self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_for_token_classification(*SCREAMING_SNAKE_CASE_							)




						@slow
						def   lowercase_       ( self							)    ->			List[str]:
										for model_name in TF_ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
														lowerCAmelCase_       :					List[Any]			 = TFEsmModel.from_pretrained(SCREAMING_SNAKE_CASE_							)
														self.assertIsNotNone(SCREAMING_SNAKE_CASE_							)




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




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




						def   lowercase_       ( self							)    ->			Any:
										lowerCAmelCase_ ,     lowerCAmelCase_       :					Tuple			 = self.model_tester.prepare_config_and_inputs_for_common()

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







@require_tf
class  snake_case__(							unittest.TestCase							):







						'''simple docstring'''




						@slow
						def   lowercase_       ( self							)    ->			Optional[int]:
										lowerCAmelCase_       :					Optional[Any]			 = TFEsmForMaskedLM.from_pretrained('''facebook/esm2_t6_8M_UR50D'''							)

										lowerCAmelCase_       :					List[Any]			 = tf.constant([[0, 1, 2, 3, 4, 5]]							)
										lowerCAmelCase_       :					List[Any]			 = model(SCREAMING_SNAKE_CASE_							)[0]
										lowerCAmelCase_       :					Dict			 = [1, 6, 3_3]
										self.assertEqual(list(output.numpy().shape							)      ,  SCREAMING_SNAKE_CASE_							)
										# compare the actual values for a slice.
										lowerCAmelCase_       :					List[str]			 = tf.constant(
										    [
										        [
										            [8.92_15_18, -10.58_98_14, -6.4_67_13_07],
										            [-6.3_96_71_56, -13.91_13_77, -1.1_21_19_15],
										            [-7.78_12_47, -13.95_15_57, -3.74_05_92],
										        ]
										    ]							)
										self.assertTrue(numpy.allclose(output[:, :3, :3].numpy()      ,  expected_slice.numpy()      ,  atol=1e-2							)							)




						@slow
						def   lowercase_       ( self							)    ->			Union[str, Any]:
										lowerCAmelCase_       :					str			 = TFEsmModel.from_pretrained('''facebook/esm2_t6_8M_UR50D'''							)

										lowerCAmelCase_       :					Dict			 = tf.constant([[0, 6, 4, 1_3, 5, 4, 1_6, 1_2, 1_1, 7, 2]]							)
										lowerCAmelCase_       :					Dict			 = model(SCREAMING_SNAKE_CASE_							)[0]
										# compare the actual values for a slice.
										lowerCAmelCase_       :					Dict			 = 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							)							)

from typing import TYPE_CHECKING

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


__a  			=					{
    'configuration_rembert': ['REMBERT_PRETRAINED_CONFIG_ARCHIVE_MAP', 'RemBertConfig', 'RemBertOnnxConfig']
}

try:
						if not is_sentencepiece_available():
												raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
						pass
else:
						__a  			=					['RemBertTokenizer']

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

try:
						if not is_torch_available():
												raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
						pass
else:
						__a  			=					[
						    'REMBERT_PRETRAINED_MODEL_ARCHIVE_LIST',
						    'RemBertForCausalLM',
						    'RemBertForMaskedLM',
						    'RemBertForMultipleChoice',
						    'RemBertForQuestionAnswering',
						    'RemBertForSequenceClassification',
						    'RemBertForTokenClassification',
						    'RemBertLayer',
						    'RemBertModel',
						    'RemBertPreTrainedModel',
						    'load_tf_weights_in_rembert',
						]


try:
						if not is_tf_available():
												raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
						pass
else:
						__a  			=					[
						    'TF_REMBERT_PRETRAINED_MODEL_ARCHIVE_LIST',
						    'TFRemBertForCausalLM',
						    'TFRemBertForMaskedLM',
						    'TFRemBertForMultipleChoice',
						    'TFRemBertForQuestionAnswering',
						    'TFRemBertForSequenceClassification',
						    'TFRemBertForTokenClassification',
						    'TFRemBertLayer',
						    'TFRemBertModel',
						    'TFRemBertPreTrainedModel',
						]


if TYPE_CHECKING:
						from .configuration_rembert import REMBERT_PRETRAINED_CONFIG_ARCHIVE_MAP, RemBertConfig, RemBertOnnxConfig

						try:
												if not is_sentencepiece_available():
																		raise OptionalDependencyNotAvailable()
						except OptionalDependencyNotAvailable:
												pass
						else:
												from .tokenization_rembert import RemBertTokenizer

						try:
												if not is_tokenizers_available():
																		raise OptionalDependencyNotAvailable()
						except OptionalDependencyNotAvailable:
												pass
						else:
												from .tokenization_rembert_fast import RemBertTokenizerFast

						try:
												if not is_torch_available():
																		raise OptionalDependencyNotAvailable()
						except OptionalDependencyNotAvailable:
												pass
						else:
												from .modeling_rembert import (
												    REMBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
												    RemBertForCausalLM,
												    RemBertForMaskedLM,
												    RemBertForMultipleChoice,
												    RemBertForQuestionAnswering,
												    RemBertForSequenceClassification,
												    RemBertForTokenClassification,
												    RemBertLayer,
												    RemBertModel,
												    RemBertPreTrainedModel,
												    load_tf_weights_in_rembert,
												)

						try:
												if not is_tf_available():
																		raise OptionalDependencyNotAvailable()
						except OptionalDependencyNotAvailable:
												pass
						else:
												from .modeling_tf_rembert import (
												    TF_REMBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
												    TFRemBertForCausalLM,
												    TFRemBertForMaskedLM,
												    TFRemBertForMultipleChoice,
												    TFRemBertForQuestionAnswering,
												    TFRemBertForSequenceClassification,
												    TFRemBertForTokenClassification,
												    TFRemBertLayer,
												    TFRemBertModel,
												    TFRemBertPreTrainedModel,
												)


else:
						import sys

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




