"""simple docstring"""




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			,			__lowercase			,			__lowercase=13			,			__lowercase=7			,			__lowercase=False			,			__lowercase=True			,			__lowercase=False			,			__lowercase=True			,			__lowercase=33			,			__lowercase=32			,			__lowercase=5			,			__lowercase=4			,			__lowercase=37			,			__lowercase="gelu"			,			__lowercase=0.1			,			__lowercase=0.1			,			__lowercase=512			,			__lowercase=16			,			__lowercase=2			,			__lowercase=0.02			,			__lowercase=3			,			__lowercase=4			,			__lowercase=None			,			)	->     Any:
        __UpperCamelCase					:List[str]        =		parent
        __UpperCamelCase					:Any        =		batch_size
        __UpperCamelCase					:int        =		seq_length
        __UpperCamelCase					:Tuple        =		is_training
        __UpperCamelCase					:Optional[Any]        =		use_input_mask
        __UpperCamelCase					:Optional[int]        =		use_token_type_ids
        __UpperCamelCase					:Optional[int]        =		use_labels
        __UpperCamelCase					:Optional[Any]        =		vocab_size
        __UpperCamelCase					:Union[str, Any]        =		hidden_size
        __UpperCamelCase					:int        =		num_hidden_layers
        __UpperCamelCase					:Any        =		num_attention_heads
        __UpperCamelCase					:str        =		intermediate_size
        __UpperCamelCase					:Optional[Any]        =		hidden_act
        __UpperCamelCase					:Union[str, Any]        =		hidden_dropout_prob
        __UpperCamelCase					:Any        =		attention_probs_dropout_prob
        __UpperCamelCase					:Dict        =		max_position_embeddings
        __UpperCamelCase					:Union[str, Any]        =		type_vocab_size
        __UpperCamelCase					:str        =		type_sequence_label_size
        __UpperCamelCase					:Tuple        =		initializer_range
        __UpperCamelCase					:Optional[int]        =		num_labels
        __UpperCamelCase					:str        =		num_choices
        __UpperCamelCase					:List[Any]        =		scope







    def       UpperCamelCase__       (     self)	->     Optional[Any]:
        __UpperCamelCase					:Any        =		ids_tensor([self.batch_size, self.seq_length]			,			self.vocab_size)

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

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

        __UpperCamelCase					:Any        =		self.get_config()

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







    def       UpperCamelCase__       (     self)	->     str:
        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       UpperCamelCase__       (     self			,			__lowercase			,			__lowercase			,			__lowercase			,			__lowercase			,			__lowercase			,			__lowercase)	->     int:
        __UpperCamelCase					:Tuple        =		EsmModel(config=_snake_case)
        model.to(_snake_case)
        model.eval()
        __UpperCamelCase					:Dict        =		model(_snake_case			,			attention_mask=_snake_case)
        __UpperCamelCase					:Union[str, Any]        =		model(_snake_case)
        __UpperCamelCase					:int        =		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       UpperCamelCase__       (     self			,			__lowercase			,			__lowercase			,			__lowercase			,			__lowercase			,			__lowercase			,			__lowercase)	->     List[Any]:
        __UpperCamelCase					:int        =		EsmForMaskedLM(config=_snake_case)
        model.to(_snake_case)
        model.eval()
        __UpperCamelCase					:List[str]        =		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       UpperCamelCase__       (     self			,			__lowercase			,			__lowercase			,			__lowercase			,			__lowercase			,			__lowercase			,			__lowercase)	->     Tuple:
        __UpperCamelCase					:List[Any]        =		self.num_labels
        __UpperCamelCase					:int        =		EsmForTokenClassification(config=_snake_case)
        model.to(_snake_case)
        model.eval()
        __UpperCamelCase					:Optional[Any]        =		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       UpperCamelCase__       (     self)	->     Optional[Any]:
        __UpperCamelCase					:Tuple        =		self.prepare_config_and_inputs()
        (
            (
            __UpperCamelCase	
        )  ,	(
            __UpperCamelCase	
        )  ,	(
            __UpperCamelCase	
        )  ,	(
            __UpperCamelCase	
        )  ,	(
            __UpperCamelCase	
        )  ,	(
            __UpperCamelCase	
        )  ,		
        )					:List[str]        =		config_and_inputs
        __UpperCamelCase					:Any        =		{'''input_ids''': input_ids, '''attention_mask''': input_mask}
        return config, inputs_dict


@require_torch
class      lowerCamelCase_					(       UpperCAmelCase_						,     UpperCAmelCase_						,     unittest.TestCase			):






    '''simple docstring'''






    a__     :						List[str]      =				False

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







    def       UpperCamelCase__       (     self)	->     Tuple:
        __UpperCamelCase					:List[Any]        =		EsmModelTester(self)
        __UpperCamelCase					:str        =		ConfigTester(self			,			config_class=_snake_case			,			hidden_size=37)







    def       UpperCamelCase__       (     self)	->     Tuple:
        self.config_tester.run_common_tests()







    def       UpperCamelCase__       (     self)	->     List[str]:
        __UpperCamelCase					:List[Any]        =		self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_model(*_snake_case)







    def       UpperCamelCase__       (     self)	->     Dict:
        __UpperCamelCase					:int        =		self.model_tester.prepare_config_and_inputs()
        for type in ["absolute", "relative_key", "relative_key_query"]:
            __UpperCamelCase					:Optional[int]        =		type
            self.model_tester.create_and_check_model(*_snake_case)







    def       UpperCamelCase__       (     self)	->     str:
        __UpperCamelCase					:Dict        =		self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_masked_lm(*_snake_case)







    def       UpperCamelCase__       (     self)	->     List[str]:
        __UpperCamelCase					:Union[str, Any]        =		self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_token_classification(*_snake_case)







    @slow
    def       UpperCamelCase__       (     self)	->     Union[str, Any]:


        for model_name in ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
            __UpperCamelCase					:List[Any]        =		EsmModel.from_pretrained(_snake_case)
            self.assertIsNotNone(_snake_case)







    def       UpperCamelCase__       (     self)	->     Any:
        __UpperCamelCase					:List[str]        =		self.model_tester.prepare_config_and_inputs()[0]
        __UpperCamelCase					:int        =		EsmEmbeddings(config=_snake_case)

        __UpperCamelCase					:str        =		torch.as_tensor([[12, 31, 13, model.padding_idx]])
        __UpperCamelCase					:List[Any]        =		torch.as_tensor(
            [
                [
                    0 + model.padding_idx + 1,
                    1 + model.padding_idx + 1,
                    2 + model.padding_idx + 1,
                    model.padding_idx,
                ]
            ])
        __UpperCamelCase					:int        =		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       UpperCamelCase__       (     self)	->     List[Any]:
        __UpperCamelCase					:Optional[int]        =		self.model_tester.prepare_config_and_inputs()[0]
        __UpperCamelCase					:Any        =		EsmEmbeddings(config=_snake_case)

        __UpperCamelCase					:Tuple        =		torch.empty(2			,			4			,			30)
        __UpperCamelCase					:int        =		[
            0 + embeddings.padding_idx + 1,
            1 + embeddings.padding_idx + 1,
            2 + embeddings.padding_idx + 1,
            3 + embeddings.padding_idx + 1,
        ]
        __UpperCamelCase					:Union[str, Any]        =		torch.as_tensor([expected_single_positions, expected_single_positions])
        __UpperCamelCase					:Dict        =		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       UpperCamelCase__       (     self)	->     List[str]:
        pass







    @unittest.skip('''Esm does not support embedding resizing''')
    def       UpperCamelCase__       (     self)	->     Dict:
        pass







    @unittest.skip('''Will be fixed soon by reducing the size of the model used for common tests.''')
    def       UpperCamelCase__       (     self)	->     Union[str, Any]:
        pass


@require_torch
class      lowerCamelCase_					(       UpperCAmelCase_			):






    '''simple docstring'''






    @slow
    def       UpperCamelCase__       (     self)	->     Optional[int]:


        with torch.no_grad():
            __UpperCamelCase					:int        =		EsmForMaskedLM.from_pretrained('''facebook/esm2_t6_8M_UR50D''')
            model.eval()
            __UpperCamelCase					:Dict        =		torch.tensor([[0, 1, 2, 3, 4, 5]])
            __UpperCamelCase					:Union[str, Any]        =		model(_snake_case)[0]

            __UpperCamelCase					:List[Any]        =		33

            __UpperCamelCase					:List[Any]        =		torch.Size((1, 6, vocab_size))
            self.assertEqual(output.shape			,			_snake_case)

            __UpperCamelCase					:Union[str, Any]        =		torch.tensor(
                [[[8.92_15, -10.5_898, -6.46_71], [-6.39_67, -13.9_114, -1.12_12], [-7.78_12, -13.9_516, -3.74_06]]])
            self.assertTrue(torch.allclose(output[:, :3, :3]			,			_snake_case			,			atol=1E-4))







    @slow
    def       UpperCamelCase__       (     self)	->     Any:


        with torch.no_grad():
            __UpperCamelCase					:Dict        =		EsmModel.from_pretrained('''facebook/esm2_t6_8M_UR50D''')
            model.eval()

            __UpperCamelCase					:Optional[int]        =		torch.tensor([[0, 6, 4, 13, 5, 4, 16, 12, 11, 7, 2]])
            __UpperCamelCase					:Dict        =		model(_snake_case)[0]
            # compare the actual values for a slice.
            __UpperCamelCase					:List[Any]        =		torch.tensor(
                [[[0.14_44, 0.54_13, 0.32_48], [0.30_34, 0.00_53, 0.31_08], [0.32_28, -0.24_99, 0.34_15]]])
            self.assertTrue(torch.allclose(output[:, :3, :3]			,			_snake_case			,			atol=1E-4))






from typing import List, Optional, Tuple

from ...tokenization_utils_fast import PreTrainedTokenizerFast
from ...utils import logging
from .tokenization_herbert import HerbertTokenizer


__lowercase										= logging.get_logger(__name__)

__lowercase										= {'''vocab_file''': '''vocab.json''', '''merges_file''': '''merges.txt''', '''tokenizer_file''': '''tokenizer.json'''}

__lowercase										= {
    '''vocab_file''': {
        '''allegro/herbert-base-cased''': '''https://huggingface.co/allegro/herbert-base-cased/resolve/main/vocab.json'''
    },
    '''merges_file''': {
        '''allegro/herbert-base-cased''': '''https://huggingface.co/allegro/herbert-base-cased/resolve/main/merges.txt'''
    },
}

__lowercase										= {'''allegro/herbert-base-cased''': 514}
__lowercase										= {}


class      lowerCamelCase_					(       UpperCAmelCase_			):






    '''simple docstring'''






    a__     :						Tuple      =				VOCAB_FILES_NAMES
    a__     :						Dict      =				PRETRAINED_VOCAB_FILES_MAP
    a__     :						Tuple      =				PRETRAINED_INIT_CONFIGURATION
    a__     :						Optional[Any]      =				PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
    a__     :						Union[str, Any]      =				HerbertTokenizer







    def __init__(     self			,			__lowercase=None			,			__lowercase=None			,			__lowercase=None			,			__lowercase="<s>"			,			__lowercase="<unk>"			,			__lowercase="<pad>"			,			__lowercase="<mask>"			,			__lowercase="</s>"			,			**__lowercase			,			)	->     Optional[Any]:
        super().__init__(
            __lowercase			,			__lowercase			,			tokenizer_file=__lowercase			,			cls_token=__lowercase			,			unk_token=__lowercase			,			pad_token=__lowercase			,			mask_token=__lowercase			,			sep_token=__lowercase			,			**__lowercase			,			)







    def       UpperCamelCase__       (     self			,			__lowercase			,			__lowercase = None)	->     List[int]:
        __UpperCamelCase					:List[str]        =		[self.cls_token_id]
        __UpperCamelCase					:Tuple        =		[self.sep_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       UpperCamelCase__       (     self			,			__lowercase			,			__lowercase = None			,			__lowercase = False)	->     List[int]:
        if already_has_special_tokens:
            return super().get_special_tokens_mask(
                token_ids_a=__lowercase			,			token_ids_a=__lowercase			,			already_has_special_tokens=__lowercase)

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







    def       UpperCamelCase__       (     self			,			__lowercase			,			__lowercase = None)	->     List[int]:
        __UpperCamelCase					:Optional[Any]        =		[self.sep_token_id]
        __UpperCamelCase					:int        =		[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       UpperCamelCase__       (     self			,			__lowercase			,			__lowercase = None)	->     Tuple[str]:
        __UpperCamelCase					:Optional[int]        =		self._tokenizer.model.save(__lowercase			,			name=__lowercase)
        return tuple(__lowercase)





