


import unittest

from transformers import EsmConfig, is_torch_available
from transformers.testing_utils import TestCasePlus, require_torch, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
						import torch

						from transformers import EsmForMaskedLM, EsmForSequenceClassification, EsmForTokenClassification, EsmModel
						from transformers.models.esm.modeling_esm import (
						    ESM_PRETRAINED_MODEL_ARCHIVE_LIST,
						    EsmEmbeddings,
						    create_position_ids_from_input_ids,
						)







class 							lowerCamelCase__						:






			"""simple docstring"""






			def __init__(				self   ,    snake_case   ,    snake_case=13   ,    snake_case=7   ,    snake_case=False   ,    snake_case=True   ,    snake_case=False   ,    snake_case=True   ,    snake_case=33   ,    snake_case=32   ,    snake_case=5   ,    snake_case=4   ,    snake_case=37   ,    snake_case="gelu"   ,    snake_case=0.1   ,    snake_case=0.1   ,    snake_case=512   ,    snake_case=16   ,    snake_case=2   ,    snake_case=0.02   ,    snake_case=3   ,    snake_case=4   ,    snake_case=None   ,    ):

						'''simple docstring'''


						UpperCamelCase__       =							parent
						UpperCamelCase__       =							batch_size
						UpperCamelCase__       =							seq_length
						UpperCamelCase__       =							is_training
						UpperCamelCase__       =							use_input_mask
						UpperCamelCase__       =							use_token_type_ids
						UpperCamelCase__       =							use_labels
						UpperCamelCase__       =							vocab_size
						UpperCamelCase__       =							hidden_size
						UpperCamelCase__       =							num_hidden_layers
						UpperCamelCase__       =							num_attention_heads
						UpperCamelCase__       =							intermediate_size
						UpperCamelCase__       =							hidden_act
						UpperCamelCase__       =							hidden_dropout_prob
						UpperCamelCase__       =							attention_probs_dropout_prob
						UpperCamelCase__       =							max_position_embeddings
						UpperCamelCase__       =							type_vocab_size
						UpperCamelCase__       =							type_sequence_label_size
						UpperCamelCase__       =							initializer_range
						UpperCamelCase__       =							num_labels
						UpperCamelCase__       =							num_choices
						UpperCamelCase__       =							scope




			def 					snake_case__					(				self  ):

						'''simple docstring'''


						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__       =							self.get_config()

						return config, input_ids, input_mask, sequence_labels, token_labels, choice_labels




			def 					snake_case__					(				self  ):

						'''simple docstring'''


						return EsmConfig(
						    vocab_size=self.vocab_size   ,    hidden_size=self.hidden_size   ,    pad_token_id=1   ,    num_hidden_layers=self.num_hidden_layers   ,    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   ,    )




			def 					snake_case__					(				self   ,    snake_case   ,    snake_case   ,    snake_case   ,    snake_case   ,    snake_case   ,    snake_case  ):

						'''simple docstring'''


						UpperCamelCase__       =							EsmModel(config=snake_case  )
						model.to(snake_case  )
						model.eval()
						UpperCamelCase__       =							model(snake_case   ,    attention_mask=snake_case  )
						UpperCamelCase__       =							model(snake_case  )
						UpperCamelCase__       =							model(snake_case  )

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




			def 					snake_case__					(				self   ,    snake_case   ,    snake_case   ,    snake_case   ,    snake_case   ,    snake_case   ,    snake_case  ):

						'''simple docstring'''


						UpperCamelCase__       =							EsmForMaskedLM(config=snake_case  )
						model.to(snake_case  )
						model.eval()
						UpperCamelCase__       =							model(snake_case   ,    attention_mask=snake_case   ,    labels=snake_case  )
						self.parent.assertEqual(result.logits.shape   ,    (self.batch_size, self.seq_length, self.vocab_size)  )




			def 					snake_case__					(				self   ,    snake_case   ,    snake_case   ,    snake_case   ,    snake_case   ,    snake_case   ,    snake_case  ):

						'''simple docstring'''


						UpperCamelCase__       =							self.num_labels
						UpperCamelCase__       =							EsmForTokenClassification(config=snake_case  )
						model.to(snake_case  )
						model.eval()
						UpperCamelCase__       =							model(snake_case   ,    attention_mask=snake_case   ,    labels=snake_case  )
						self.parent.assertEqual(result.logits.shape   ,    (self.batch_size, self.seq_length, self.num_labels)  )







			def 					snake_case__					(				self  ):

						'''simple docstring'''


						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_torch
class 							lowerCamelCase__						(				UpperCAmelCase       ,						UpperCAmelCase       ,						unittest.TestCase   ):






			"""simple docstring"""






			_UpperCamelCase		:       List[str] 			=		False

			_UpperCamelCase		:       Any 			=		(
			    (
			        EsmForMaskedLM,
			        EsmModel,
			        EsmForSequenceClassification,
			        EsmForTokenClassification,
			    )
			    if is_torch_available()
			    else ()
			)
			_UpperCamelCase		:       Optional[int] 			=		()
			_UpperCamelCase		:       Dict 			=		(
			    {
			        'feature-extraction': EsmModel,
			        'fill-mask': EsmForMaskedLM,
			        'text-classification': EsmForSequenceClassification,
			        'token-classification': EsmForTokenClassification,
			        'zero-shot': EsmForSequenceClassification,
			    }
			    if is_torch_available()
			    else {}
			)
			_UpperCamelCase		:       Optional[int] 			=		True




			def 					snake_case__					(				self  ):

						'''simple docstring'''


						UpperCamelCase__       =							EsmModelTester(self  )
						UpperCamelCase__       =							ConfigTester(self   ,    config_class=snake_case   ,    hidden_size=37  )




			def 					snake_case__					(				self  ):

						'''simple docstring'''


						self.config_tester.run_common_tests()




			def 					snake_case__					(				self  ):

						'''simple docstring'''


						UpperCamelCase__       =							self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_model(*snake_case  )




			def 					snake_case__					(				self  ):

						'''simple docstring'''


						UpperCamelCase__       =							self.model_tester.prepare_config_and_inputs()
						for type in ["absolute", "relative_key", "relative_key_query"]:
									UpperCamelCase__       =							type
									self.model_tester.create_and_check_model(*snake_case  )




			def 					snake_case__					(				self  ):

						'''simple docstring'''


						UpperCamelCase__       =							self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_masked_lm(*snake_case  )




			def 					snake_case__					(				self  ):

						'''simple docstring'''


						UpperCamelCase__       =							self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_token_classification(*snake_case  )




			@slow
			def 					snake_case__					(				self  ):

						'''simple docstring'''


						for model_name in ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
									UpperCamelCase__       =							EsmModel.from_pretrained(snake_case  )
									self.assertIsNotNone(snake_case  )




			def 					snake_case__					(				self  ):

						'''simple docstring'''


						UpperCamelCase__       =							self.model_tester.prepare_config_and_inputs()[0]
						UpperCamelCase__       =							EsmEmbeddings(config=snake_case  )

						UpperCamelCase__       =							torch.as_tensor([[12, 31, 13, model.padding_idx]]  )
						UpperCamelCase__       =							torch.as_tensor(
						    [
						        [
						            0 + model.padding_idx + 1,
						            1 + model.padding_idx + 1,
						            2 + model.padding_idx + 1,
						            model.padding_idx,
						        ]
						    ]  )
						UpperCamelCase__       =							create_position_ids_from_input_ids(snake_case   ,    model.padding_idx  )
						self.assertEqual(position_ids.shape   ,    expected_positions.shape  )
						self.assertTrue(torch.all(torch.eq(snake_case   ,    snake_case  )  )  )




			def 					snake_case__					(				self  ):

						'''simple docstring'''


						UpperCamelCase__       =							self.model_tester.prepare_config_and_inputs()[0]
						UpperCamelCase__       =							EsmEmbeddings(config=snake_case  )

						UpperCamelCase__       =							torch.empty(2   ,    4   ,    30  )
						UpperCamelCase__       =							[
						    0 + embeddings.padding_idx + 1,
						    1 + embeddings.padding_idx + 1,
						    2 + embeddings.padding_idx + 1,
						    3 + embeddings.padding_idx + 1,
						]
						UpperCamelCase__       =							torch.as_tensor([expected_single_positions, expected_single_positions]  )
						UpperCamelCase__       =							embeddings.create_position_ids_from_inputs_embeds(snake_case  )
						self.assertEqual(position_ids.shape   ,    expected_positions.shape  )
						self.assertTrue(torch.all(torch.eq(snake_case   ,    snake_case  )  )  )




			@unittest.skip("Esm does not support embedding resizing"  )
			def 					snake_case__					(				self  ):

						'''simple docstring'''


						pass




			@unittest.skip("Esm does not support embedding resizing"  )
			def 					snake_case__					(				self  ):

						'''simple docstring'''


						pass







			@unittest.skip("Will be fixed soon by reducing the size of the model used for common tests."  )
			def 					snake_case__					(				self  ):

						'''simple docstring'''


						pass







@require_torch
class 							lowerCamelCase__						(				UpperCAmelCase   ):






			"""simple docstring"""






			@slow
			def 					snake_case__					(				self  ):

						'''simple docstring'''


						with torch.no_grad():
									UpperCamelCase__       =							EsmForMaskedLM.from_pretrained("facebook/esm2_t6_8M_UR50D"  )
									model.eval()
									UpperCamelCase__       =							torch.tensor([[0, 1, 2, 3, 4, 5]]  )
									UpperCamelCase__       =							model(snake_case  )[0]

									UpperCamelCase__       =							33

									UpperCamelCase__       =							torch.Size((1, 6, vocab_size)  )
									self.assertEqual(output.shape   ,    snake_case  )

									UpperCamelCase__       =							torch.tensor(
									    [[[8.9215, -10.5898, -6.4671], [-6.3967, -13.9114, -1.1212], [-7.7812, -13.9516, -3.7406]]]  )
									self.assertTrue(torch.allclose(output[:, :3, :3]   ,    snake_case   ,    atol=1E-4  )  )







			@slow
			def 					snake_case__					(				self  ):

						'''simple docstring'''


						with torch.no_grad():
									UpperCamelCase__       =							EsmModel.from_pretrained("facebook/esm2_t6_8M_UR50D"  )
									model.eval()

									UpperCamelCase__       =							torch.tensor([[0, 6, 4, 13, 5, 4, 16, 12, 11, 7, 2]]  )
									UpperCamelCase__       =							model(snake_case  )[0]
									# compare the actual values for a slice.
									UpperCamelCase__       =							torch.tensor(
									    [[[0.1444, 0.5413, 0.3248], [0.3034, 0.0053, 0.3108], [0.3228, -0.2499, 0.3415]]]  )
									self.assertTrue(torch.allclose(output[:, :3, :3]   ,    snake_case   ,    atol=1E-4  )  )






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_big_bird import BigBirdTokenizer
else:
						__UpperCamelCase     		=							None

__UpperCamelCase     		=							logging.get_logger(__name__)
__UpperCamelCase     		=							{'vocab_file': 'spiece.model', 'tokenizer_file': 'tokenizer.json'}

__UpperCamelCase     		=							{
    'vocab_file': {
        'google/bigbird-roberta-base': 'https://huggingface.co/google/bigbird-roberta-base/resolve/main/spiece.model',
        'google/bigbird-roberta-large': (
            'https://huggingface.co/google/bigbird-roberta-large/resolve/main/spiece.model'
        ),
        'google/bigbird-base-trivia-itc': (
            'https://huggingface.co/google/bigbird-base-trivia-itc/resolve/main/spiece.model'
        ),
    },
    'tokenizer_file': {
        'google/bigbird-roberta-base': (
            'https://huggingface.co/google/bigbird-roberta-base/resolve/main/tokenizer.json'
        ),
        'google/bigbird-roberta-large': (
            'https://huggingface.co/google/bigbird-roberta-large/resolve/main/tokenizer.json'
        ),
        'google/bigbird-base-trivia-itc': (
            'https://huggingface.co/google/bigbird-base-trivia-itc/resolve/main/tokenizer.json'
        ),
    },
}

__UpperCamelCase     		=							{
    'google/bigbird-roberta-base': 4_0_9_6,
    'google/bigbird-roberta-large': 4_0_9_6,
    'google/bigbird-base-trivia-itc': 4_0_9_6,
}


__UpperCamelCase     		=							'▁'







class 							lowerCamelCase__						(				UpperCAmelCase   ):






			"""simple docstring"""






			_UpperCamelCase		:       Dict 			=		VOCAB_FILES_NAMES
			_UpperCamelCase		:       int 			=		PRETRAINED_VOCAB_FILES_MAP
			_UpperCamelCase		:       List[Any] 			=		PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
			_UpperCamelCase		:       Optional[Any] 			=		BigBirdTokenizer
			_UpperCamelCase		:       Dict 			=		['input_ids', 'attention_mask']
			_UpperCamelCase		:       List[int] 			=		[]




			def __init__(				self   ,    snake_case=None   ,    snake_case=None   ,    snake_case="<unk>"   ,    snake_case="<s>"   ,    snake_case="</s>"   ,    snake_case="<pad>"   ,    snake_case="[SEP]"   ,    snake_case="[MASK]"   ,    snake_case="[CLS]"   ,    **snake_case   ,    ):

						'''simple docstring'''


						UpperCamelCase__       =							AddedToken(snake_case   ,    lstrip=snake_case   ,    rstrip=snake_case  ) if isinstance(snake_case   ,    snake_case  ) else bos_token
						UpperCamelCase__       =							AddedToken(snake_case   ,    lstrip=snake_case   ,    rstrip=snake_case  ) if isinstance(snake_case   ,    snake_case  ) else eos_token
						UpperCamelCase__       =							AddedToken(snake_case   ,    lstrip=snake_case   ,    rstrip=snake_case  ) if isinstance(snake_case   ,    snake_case  ) else unk_token
						UpperCamelCase__       =							AddedToken(snake_case   ,    lstrip=snake_case   ,    rstrip=snake_case  ) if isinstance(snake_case   ,    snake_case  ) else pad_token
						UpperCamelCase__       =							AddedToken(snake_case   ,    lstrip=snake_case   ,    rstrip=snake_case  ) if isinstance(snake_case   ,    snake_case  ) else cls_token
						UpperCamelCase__       =							AddedToken(snake_case   ,    lstrip=snake_case   ,    rstrip=snake_case  ) if isinstance(snake_case   ,    snake_case  ) else sep_token

						# Mask token behave like a normal word, i.e. include the space before it
						UpperCamelCase__       =							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   ,    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   ,    )

						UpperCamelCase__       =							vocab_file
						UpperCamelCase__       =							False if not self.vocab_file else True




			def 					snake_case__					(				self   ,    snake_case   ,    snake_case = None  ):

						'''simple docstring'''


						UpperCamelCase__       =							[self.sep_token_id]
						UpperCamelCase__       =							[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 					snake_case__					(				self   ,    snake_case   ,    snake_case = None   ,    snake_case = False  ):

						'''simple docstring'''


						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 None:
									return [1] + ([0] * len(snake_case  )) + [1]
						return [1] + ([0] * len(snake_case  )) + [1] + ([0] * len(snake_case  )) + [1]




			def 					snake_case__					(				self   ,    snake_case   ,    snake_case = None  ):

						'''simple docstring'''


						UpperCamelCase__       =							[self.sep_token_id]
						UpperCamelCase__       =							[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 					snake_case__					(				self   ,    snake_case   ,    snake_case = None  ):

						'''simple docstring'''


						if not self.can_save_slow_tokenizer:
									raise ValueError(
									    "Your fast tokenizer does not have the necessary information to save the vocabulary for a slow "
									    "tokenizer."  )

						if not os.path.isdir(snake_case  ):
									logger.error(F'''Vocabulary path ({save_directory}) should be a directory'''  )
									return
						UpperCamelCase__       =							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,)


