






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





	def __init__( self   :					Union[str, Any]      ,      lowerCAmelCase_   :					List[Any]      ,      )		->     Any:
				UpperCAmelCase_	:					Optional[int]     	=   parent
				UpperCAmelCase_	:					str     	=   13
				UpperCAmelCase_	:					Dict     	=   7
				UpperCAmelCase_	:					Dict     	=   True
				UpperCAmelCase_	:					Dict     	=   True
				UpperCAmelCase_	:					Optional[Any]     	=   True
				UpperCAmelCase_	:					Optional[Any]     	=   99
				UpperCAmelCase_	:					Optional[int]     	=   32
				UpperCAmelCase_	:					Optional[int]     	=   2
				UpperCAmelCase_	:					Optional[Any]     	=   4
				UpperCAmelCase_	:					int     	=   37
				UpperCAmelCase_	:					int     	=   "gelu"
				UpperCAmelCase_	:					Any     	=   0.1
				UpperCAmelCase_	:					Any     	=   0.1
				UpperCAmelCase_	:					List[Any]     	=   512
				UpperCAmelCase_	:					List[str]     	=   16
				UpperCAmelCase_	:					str     	=   2
				UpperCAmelCase_	:					List[str]     	=   0.0_2
				UpperCAmelCase_	:					int     	=   3
				UpperCAmelCase_	:					List[str]     	=   4
				UpperCAmelCase_	:					Optional[Any]     	=   None





	def 	_SCREAMING_SNAKE_CASE ( self   :					Optional[int]		)		->     Optional[Any]:
				UpperCAmelCase_	:					Dict     	=   ids_tensor([self.batch_size, self.seq_length]      ,      self.vocab_size		)

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

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

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

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

				UpperCAmelCase_	:					Union[str, Any]     	=   [input_ids, input_mask]
				UpperCAmelCase_	:					Dict     	=   model(lowerCAmelCase_		)

				UpperCAmelCase_	:					List[Any]     	=   model(lowerCAmelCase_		)

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





	def 	_SCREAMING_SNAKE_CASE ( self   :					List[Any]      ,      lowerCAmelCase_   :					Dict      ,      lowerCAmelCase_   :					List[Any]      ,      lowerCAmelCase_   :					Union[str, Any]      ,      lowerCAmelCase_   :					Any      ,      lowerCAmelCase_   :					List[str]      ,      lowerCAmelCase_   :					Union[str, Any]      ,      lowerCAmelCase_   :					List[str]      ,      lowerCAmelCase_   :					Optional[Any]      ,      )		->     Union[str, Any]:
				UpperCAmelCase_	:					Optional[Any]     	=   True

				UpperCAmelCase_	:					Optional[int]     	=   TFEsmModel(config=lowerCAmelCase_		)
				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_	:					Optional[Any]     	=   model(lowerCAmelCase_		)

				UpperCAmelCase_	:					Tuple     	=   [input_ids, input_mask]
				UpperCAmelCase_	:					Optional[int]     	=   model(lowerCAmelCase_      ,      encoder_hidden_states=lowerCAmelCase_		)

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

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





	def 	_SCREAMING_SNAKE_CASE ( self   :					List[str]      ,      lowerCAmelCase_   :					int      ,      lowerCAmelCase_   :					Optional[Any]      ,      lowerCAmelCase_   :					List[str]      ,      lowerCAmelCase_   :					str      ,      lowerCAmelCase_   :					str      ,      lowerCAmelCase_   :					int		)		->     Dict:
				UpperCAmelCase_	:					Dict     	=   TFEsmForMaskedLM(config=lowerCAmelCase_		)
				UpperCAmelCase_	:					Optional[int]     	=   model([input_ids, input_mask]		)
				self.parent.assertEqual(result.logits.shape      ,      (self.batch_size, self.seq_length, self.vocab_size)		)





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





	def 	_SCREAMING_SNAKE_CASE ( self   :					str		)		->     Dict:
				UpperCAmelCase_	:					int     	=   self.prepare_config_and_inputs()
				(
				    (
				    UpperCAmelCase_
				)	,      (
				    UpperCAmelCase_
				)	,      (
				    UpperCAmelCase_
				)	,      (
				    UpperCAmelCase_
				)	,      (
				    UpperCAmelCase_
				)	,      (
				    UpperCAmelCase_
				)	,      
				)	:					Optional[int]     	=   config_and_inputs
				UpperCAmelCase_	:					List[str]     	=   {"input_ids": input_ids, "attention_mask": input_mask}
				return config, inputs_dict







@require_tf
class      UpperCamelCase_			(__A							,		__A							,		unittest.TestCase	):
	__magic_name__       =							(
	    (
	        TFEsmModel,
	        TFEsmForMaskedLM,
	        TFEsmForSequenceClassification,
	        TFEsmForTokenClassification,
	    )
	    if is_tf_available()
	    else ()
	)
	__magic_name__       =							(
	    {
	        '''feature-extraction''': TFEsmModel,
	        '''fill-mask''': TFEsmForMaskedLM,
	        '''text-classification''': TFEsmForSequenceClassification,
	        '''token-classification''': TFEsmForTokenClassification,
	        '''zero-shot''': TFEsmForSequenceClassification,
	    }
	    if is_tf_available()
	    else {}
	)
	__magic_name__       =							False
	__magic_name__       =							False





	def 	_SCREAMING_SNAKE_CASE ( self   :					Optional[Any]		)		->     Optional[int]:
				UpperCAmelCase_	:					Any     	=   TFEsmModelTester(self		)
				UpperCAmelCase_	:					Optional[Any]     	=   ConfigTester(self      ,      config_class=lowerCAmelCase_      ,      hidden_size=37		)





	def 	_SCREAMING_SNAKE_CASE ( self   :					str		)		->     List[Any]:
				self.config_tester.run_common_tests()





	def 	_SCREAMING_SNAKE_CASE ( self   :					Dict		)		->     Union[str, Any]:
				UpperCAmelCase_	:					Tuple     	=   self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_model(*lowerCAmelCase_		)





	def 	_SCREAMING_SNAKE_CASE ( self   :					Tuple		)		->     List[str]:
				UpperCAmelCase_	:					List[str]     	=   self.model_tester.prepare_config_and_inputs_for_decoder()
				self.model_tester.create_and_check_model_as_decoder(*lowerCAmelCase_		)





	def 	_SCREAMING_SNAKE_CASE ( self   :					str		)		->     Tuple:
				UpperCAmelCase_	:					int     	=   self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_for_masked_lm(*lowerCAmelCase_		)





	def 	_SCREAMING_SNAKE_CASE ( self   :					Optional[int]		)		->     List[Any]:
				UpperCAmelCase_	:					Union[str, Any]     	=   self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_for_token_classification(*lowerCAmelCase_		)





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





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





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





	def 	_SCREAMING_SNAKE_CASE ( self   :					Tuple		)		->     int:
				UpperCAmelCase_	,      UpperCAmelCase_	:					Union[str, Any]     	=   self.model_tester.prepare_config_and_inputs_for_common()

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







@require_tf
class      UpperCamelCase_			(unittest.TestCase	):





	@slow
	def 	_SCREAMING_SNAKE_CASE ( self   :					Union[str, Any]		)		->     List[Any]:
				UpperCAmelCase_	:					Any     	=   TFEsmForMaskedLM.from_pretrained("facebook/esm2_t6_8M_UR50D"		)

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

				UpperCAmelCase_	:					List[Any]     	=   tf.constant([[0, 6, 4, 13, 5, 4, 16, 12, 11, 7, 2]]		)
				UpperCAmelCase_	:					Optional[int]     	=   model(lowerCAmelCase_		)[0]
				# compare the actual values for a slice.
				UpperCAmelCase_	:					Optional[Any]     	=   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'''




import os
from shutil import copyfile
from typing import List, Optional, Tuple

from ...tokenization_utils import AddedToken
from ...tokenization_utils_fast import PreTrainedTokenizerFast
from ...utils import is_sentencepiece_available, logging


if is_sentencepiece_available():
 from .tokenization_rembert import RemBertTokenizer
else:
 _SCREAMING_SNAKE_CASE							     =			None

_SCREAMING_SNAKE_CASE							     =			logging.get_logger(__name__)
_SCREAMING_SNAKE_CASE							     =			{"vocab_file": "sentencepiece.model", "tokenizer_file": "tokenizer.json"}

_SCREAMING_SNAKE_CASE							     =			{
    "vocab_file": {
        "google/rembert": "https://huggingface.co/google/rembert/resolve/main/sentencepiece.model",
    },
    "tokenizer_file": {
        "google/rembert": "https://huggingface.co/google/rembert/resolve/main/tokenizer.json",
    },
}

_SCREAMING_SNAKE_CASE							     =			{
    "google/rembert": 256,
}

_SCREAMING_SNAKE_CASE							     =			"▁"
class 			_lowerCAmelCase		(			A__	):






     """simple docstring"""






     snake_case_						      =   VOCAB_FILES_NAMES
     snake_case_						      =   PRETRAINED_VOCAB_FILES_MAP
     snake_case_						      =   PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
     snake_case_						      =   RemBertTokenizer



     def __init__(   self							:			List[Any]   ,       __snake_case							:			int=None   ,       __snake_case							:			Union[str, Any]=None   ,       __snake_case							:			Tuple=True   ,       __snake_case							:			Dict=True   ,       __snake_case							:			str=False   ,       __snake_case							:			Union[str, Any]="[CLS]"   ,       __snake_case							:			Optional[int]="[SEP]"   ,       __snake_case							:			str="<unk>"   ,       __snake_case							:			Dict="[SEP]"   ,       __snake_case							:			Dict="<pad>"   ,       __snake_case							:			Union[str, Any]="[CLS]"   ,       __snake_case							:			int="[MASK]"   ,       **__snake_case							:			Optional[int]   ,       )->	Tuple:
        # Mask token behave like a normal word, i.e. include the space before it
        snake_case 						=   AddedToken(__snake_case   ,       lstrip=__snake_case   ,       rstrip=__snake_case						) if isinstance(__snake_case   ,       __snake_case						) else mask_token

        super().__init__(
            __snake_case   ,       tokenizer_file=__snake_case   ,       do_lower_case=__snake_case   ,       remove_space=__snake_case   ,       keep_accents=__snake_case   ,       bos_token=__snake_case   ,       eos_token=__snake_case   ,       unk_token=__snake_case   ,       sep_token=__snake_case   ,       pad_token=__snake_case   ,       cls_token=__snake_case   ,       mask_token=__snake_case   ,       **__snake_case   ,       )

        snake_case 						=   do_lower_case
        snake_case 						=   remove_space
        snake_case 						=   keep_accents
        snake_case 						=   vocab_file
        snake_case 						=   False if not self.vocab_file else True



     def      lowerCAmelCase			(   self							:			List[str]   ,       __snake_case							:			List[int]   ,       __snake_case							:			Optional[List[int]] = None						)->	List[int]:
        snake_case 						=   [self.sep_token_id]
        snake_case 						=   [self.cls_token_id]
        if token_ids_a is None:
           return cls + token_ids_a + sep
        return cls + token_ids_a + sep + token_ids_a + sep



     def      lowerCAmelCase			(   self							:			List[Any]   ,       __snake_case							:			List[int]   ,       __snake_case							:			Optional[List[int]] = None   ,       __snake_case							:			bool = False						)->	List[int]:

        if already_has_special_tokens:
           if token_ids_a is not None:
              raise ValueError(
                  """You should not supply a second sequence if the provided sequence of """
                  """ids is already formatted with special tokens for the model."""						)
           return [1 if x in [self.sep_token_id, self.cls_token_id] else 0 for x in token_ids_a]

        if token_ids_a is not None:
           return [1] + ([0] * len(__snake_case						)) + [1] + ([0] * len(__snake_case						)) + [1]
        return [1] + ([0] * len(__snake_case						)) + [1]



     def      lowerCAmelCase			(   self							:			int   ,       __snake_case							:			List[int]   ,       __snake_case							:			Optional[List[int]] = None						)->	List[int]:
        snake_case 						=   [self.sep_token_id]
        snake_case 						=   [self.cls_token_id]

        if token_ids_a is None:
           return len(cls + token_ids_a + sep						) * [0]
        return len(cls + token_ids_a + sep						) * [0] + len(token_ids_a + sep						) * [1]



     def      lowerCAmelCase			(   self							:			Tuple   ,       __snake_case							:			str   ,       __snake_case							:			Optional[str] = None						)->	Tuple[str]:
        if not os.path.isdir(__snake_case						):
           logger.error("""Vocabulary path ({}) should be a directory""".format(__snake_case						)						)
           return
        snake_case 						=   os.path.join(
            __snake_case   ,       (filename_prefix + """-""" if filename_prefix else """""") + VOCAB_FILES_NAMES["""vocab_file"""]						)

        if os.path.abspath(self.vocab_file						) != os.path.abspath(__snake_case						):
           copyfile(self.vocab_file   ,       __snake_case						)

        return (out_vocab_file,)

