import unittest

from transformers import AlbertConfig, is_torch_available
from transformers.models.auto import get_values
from transformers.testing_utils import 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 (
        MODEL_FOR_PRETRAINING_MAPPING,
        AlbertForMaskedLM,
        AlbertForMultipleChoice,
        AlbertForPreTraining,
        AlbertForQuestionAnswering,
        AlbertForSequenceClassification,
        AlbertForTokenClassification,
        AlbertModel,
    )
    from transformers.models.albert.modeling_albert import ALBERT_PRETRAINED_MODEL_ARCHIVE_LIST


class 	lowerCamelCase_     :






    '''simple docstring'''






    def __init__(	self , __lowercase , __lowercase=13 , __lowercase=7 , __lowercase=True , __lowercase=True , __lowercase=True , __lowercase=True , __lowercase=99 , __lowercase=16 , __lowercase=36 , __lowercase=6 , __lowercase=6 , __lowercase=6 , __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 , )       ->	Optional[Any]:
        __UpperCamelCase     :str			  =		parent
        __UpperCamelCase     :List[str]			  =		batch_size
        __UpperCamelCase     :List[Any]			  =		seq_length
        __UpperCamelCase     :Optional[Any]			  =		is_training
        __UpperCamelCase     :List[Any]			  =		use_input_mask
        __UpperCamelCase     :str			  =		use_token_type_ids
        __UpperCamelCase     :Dict			  =		use_labels
        __UpperCamelCase     :List[str]			  =		vocab_size
        __UpperCamelCase     :Any			  =		embedding_size
        __UpperCamelCase     :Union[str, Any]			  =		hidden_size
        __UpperCamelCase     :Tuple			  =		num_hidden_layers
        __UpperCamelCase     :List[Any]			  =		num_hidden_groups
        __UpperCamelCase     :Optional[Any]			  =		num_attention_heads
        __UpperCamelCase     :Tuple			  =		intermediate_size
        __UpperCamelCase     :Any			  =		hidden_act
        __UpperCamelCase     :Optional[Any]			  =		hidden_dropout_prob
        __UpperCamelCase     :Any			  =		attention_probs_dropout_prob
        __UpperCamelCase     :List[Any]			  =		max_position_embeddings
        __UpperCamelCase     :str			  =		type_vocab_size
        __UpperCamelCase     :Any			  =		type_sequence_label_size
        __UpperCamelCase     :int			  =		initializer_range
        __UpperCamelCase     :int			  =		num_labels
        __UpperCamelCase     :Tuple			  =		num_choices
        __UpperCamelCase     :str			  =		scope







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

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

        __UpperCamelCase     :Dict			  =		None
        if self.use_token_type_ids:
            __UpperCamelCase     :Optional[Any]			  =		ids_tensor([self.batch_size, self.seq_length] , self.type_vocab_size)

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

        __UpperCamelCase     :Tuple			  =		self.get_config()

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







    def 			UpperCamelCase__   (	self)       ->	List[Any]:
        return AlbertConfig(
            vocab_size=self.vocab_size , hidden_size=self.hidden_size , 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 , num_hidden_groups=self.num_hidden_groups , )







    def 			UpperCamelCase__   (	self , __lowercase , __lowercase , __lowercase , __lowercase , __lowercase , __lowercase , __lowercase)       ->	Union[str, Any]:
        __UpperCamelCase     :Optional[int]			  =		AlbertModel(config=__lowercase)
        model.to(__lowercase)
        model.eval()
        __UpperCamelCase     :Optional[Any]			  =		model(__lowercase , attention_mask=__lowercase , token_type_ids=__lowercase)
        __UpperCamelCase     :List[Any]			  =		model(__lowercase , token_type_ids=__lowercase)
        __UpperCamelCase     :int			  =		model(__lowercase)
        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 , __lowercase)       ->	Optional[Any]:
        __UpperCamelCase     :int			  =		AlbertForPreTraining(config=__lowercase)
        model.to(__lowercase)
        model.eval()
        __UpperCamelCase     :List[str]			  =		model(
            __lowercase , attention_mask=__lowercase , token_type_ids=__lowercase , labels=__lowercase , sentence_order_label=__lowercase , )
        self.parent.assertEqual(result.prediction_logits.shape , (self.batch_size, self.seq_length, self.vocab_size))
        self.parent.assertEqual(result.sop_logits.shape , (self.batch_size, config.num_labels))







    def 			UpperCamelCase__   (	self , __lowercase , __lowercase , __lowercase , __lowercase , __lowercase , __lowercase , __lowercase)       ->	Union[str, Any]:
        __UpperCamelCase     :Dict			  =		AlbertForMaskedLM(config=__lowercase)
        model.to(__lowercase)
        model.eval()
        __UpperCamelCase     :List[Any]			  =		model(__lowercase , attention_mask=__lowercase , token_type_ids=__lowercase , labels=__lowercase)
        self.parent.assertEqual(result.logits.shape , (self.batch_size, self.seq_length, self.vocab_size))







    def 			UpperCamelCase__   (	self , __lowercase , __lowercase , __lowercase , __lowercase , __lowercase , __lowercase , __lowercase)       ->	List[str]:
        __UpperCamelCase     :List[str]			  =		AlbertForQuestionAnswering(config=__lowercase)
        model.to(__lowercase)
        model.eval()
        __UpperCamelCase     :List[str]			  =		model(
            __lowercase , attention_mask=__lowercase , token_type_ids=__lowercase , start_positions=__lowercase , end_positions=__lowercase , )
        self.parent.assertEqual(result.start_logits.shape , (self.batch_size, self.seq_length))
        self.parent.assertEqual(result.end_logits.shape , (self.batch_size, self.seq_length))







    def 			UpperCamelCase__   (	self , __lowercase , __lowercase , __lowercase , __lowercase , __lowercase , __lowercase , __lowercase)       ->	List[str]:
        __UpperCamelCase     :List[str]			  =		self.num_labels
        __UpperCamelCase     :Optional[Any]			  =		AlbertForSequenceClassification(__lowercase)
        model.to(__lowercase)
        model.eval()
        __UpperCamelCase     :List[str]			  =		model(__lowercase , attention_mask=__lowercase , token_type_ids=__lowercase , labels=__lowercase)
        self.parent.assertEqual(result.logits.shape , (self.batch_size, self.num_labels))







    def 			UpperCamelCase__   (	self , __lowercase , __lowercase , __lowercase , __lowercase , __lowercase , __lowercase , __lowercase)       ->	List[str]:
        __UpperCamelCase     :List[Any]			  =		self.num_labels
        __UpperCamelCase     :Optional[Any]			  =		AlbertForTokenClassification(config=__lowercase)
        model.to(__lowercase)
        model.eval()
        __UpperCamelCase     :int			  =		model(__lowercase , attention_mask=__lowercase , token_type_ids=__lowercase , labels=__lowercase)
        self.parent.assertEqual(result.logits.shape , (self.batch_size, self.seq_length, self.num_labels))







    def 			UpperCamelCase__   (	self , __lowercase , __lowercase , __lowercase , __lowercase , __lowercase , __lowercase , __lowercase)       ->	Union[str, Any]:
        __UpperCamelCase     :Optional[int]			  =		self.num_choices
        __UpperCamelCase     :Optional[Any]			  =		AlbertForMultipleChoice(config=__lowercase)
        model.to(__lowercase)
        model.eval()
        __UpperCamelCase     :List[Any]			  =		input_ids.unsqueeze(1).expand(-1 , self.num_choices , -1).contiguous()
        __UpperCamelCase     :List[str]			  =		token_type_ids.unsqueeze(1).expand(-1 , self.num_choices , -1).contiguous()
        __UpperCamelCase     :str			  =		input_mask.unsqueeze(1).expand(-1 , self.num_choices , -1).contiguous()
        __UpperCamelCase     :Optional[int]			  =		model(
            __lowercase , attention_mask=__lowercase , token_type_ids=__lowercase , labels=__lowercase , )
        self.parent.assertEqual(result.logits.shape , (self.batch_size, self.num_choices))







    def 			UpperCamelCase__   (	self)       ->	str:
        __UpperCamelCase     :Optional[Any]			  =		self.prepare_config_and_inputs()
        (
            (
            __UpperCamelCase
        )  ,       (
            __UpperCamelCase
        )  ,       (
            __UpperCamelCase
        )  ,       (
            __UpperCamelCase
        )  ,       (
            __UpperCamelCase
        )  ,       (
            __UpperCamelCase
        )  ,       (
            __UpperCamelCase
        )  ,       
        )     :Optional[Any]			  =		config_and_inputs
        __UpperCamelCase     :Optional[int]			  =		{'''input_ids''': input_ids, '''token_type_ids''': token_type_ids, '''attention_mask''': input_mask}
        return config, inputs_dict


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






    '''simple docstring'''






    a__	:						Union[str, Any]						=      (
        (
            AlbertModel,
            AlbertForPreTraining,
            AlbertForMaskedLM,
            AlbertForMultipleChoice,
            AlbertForSequenceClassification,
            AlbertForTokenClassification,
            AlbertForQuestionAnswering,
        )
        if is_torch_available()
        else ()
    )
    a__	:						int						=      (
        {
            """feature-extraction""": AlbertModel,
            """fill-mask""": AlbertForMaskedLM,
            """question-answering""": AlbertForQuestionAnswering,
            """text-classification""": AlbertForSequenceClassification,
            """token-classification""": AlbertForTokenClassification,
            """zero-shot""": AlbertForSequenceClassification,
        }
        if is_torch_available()
        else {}
    )
    a__	:						List[Any]						=      True







    def 			UpperCamelCase__   (	self , __lowercase , __lowercase , __lowercase=False)       ->	Any:
        __UpperCamelCase     :List[Any]			  =		super()._prepare_for_class(__lowercase , __lowercase , return_labels=__lowercase)

        if return_labels:
            if model_class in get_values(__lowercase):
                __UpperCamelCase     :Union[str, Any]			  =		torch.zeros(
                    (self.model_tester.batch_size, self.model_tester.seq_length) , dtype=torch.long , device=__lowercase)
                __UpperCamelCase     :Union[str, Any]			  =		torch.zeros(
                    self.model_tester.batch_size , dtype=torch.long , device=__lowercase)
        return inputs_dict







    def 			UpperCamelCase__   (	self)       ->	Optional[int]:
        __UpperCamelCase     :Dict			  =		AlbertModelTester(self)
        __UpperCamelCase     :Dict			  =		ConfigTester(self , config_class=__lowercase , hidden_size=37)







    def 			UpperCamelCase__   (	self)       ->	Optional[int]:
        self.config_tester.run_common_tests()







    def 			UpperCamelCase__   (	self)       ->	Any:
        __UpperCamelCase     :Optional[int]			  =		self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_model(*__lowercase)







    def 			UpperCamelCase__   (	self)       ->	Union[str, Any]:
        __UpperCamelCase     :Dict			  =		self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_pretraining(*__lowercase)







    def 			UpperCamelCase__   (	self)       ->	Union[str, Any]:
        __UpperCamelCase     :Tuple			  =		self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_masked_lm(*__lowercase)







    def 			UpperCamelCase__   (	self)       ->	Union[str, Any]:
        __UpperCamelCase     :Optional[Any]			  =		self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_multiple_choice(*__lowercase)







    def 			UpperCamelCase__   (	self)       ->	Optional[int]:
        __UpperCamelCase     :Optional[Any]			  =		self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_question_answering(*__lowercase)







    def 			UpperCamelCase__   (	self)       ->	Optional[int]:
        __UpperCamelCase     :Union[str, Any]			  =		self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_sequence_classification(*__lowercase)







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







    @slow
    def 			UpperCamelCase__   (	self)       ->	str:
        for model_name in ALBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
            __UpperCamelCase     :List[Any]			  =		AlbertModel.from_pretrained(__lowercase)
            self.assertIsNotNone(__lowercase)


@require_torch
class 	lowerCamelCase_     (   unittest.TestCase ):






    '''simple docstring'''






    @slow
    def 			UpperCamelCase__   (	self)       ->	Optional[Any]:
        __UpperCamelCase     :int			  =		AlbertModel.from_pretrained('''albert-base-v2''')
        __UpperCamelCase     :Union[str, Any]			  =		torch.tensor([[0, 345, 232, 328, 740, 140, 1_695, 69, 6_078, 1_588, 2]])
        __UpperCamelCase     :str			  =		torch.tensor([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]])
        with torch.no_grad():
            __UpperCamelCase     :str			  =		model(__lowercase , attention_mask=__lowercase)[0]
        __UpperCamelCase     :Optional[Any]			  =		torch.Size((1, 11, 768))
        self.assertEqual(output.shape , __lowercase)
        __UpperCamelCase     :Tuple			  =		torch.tensor(
            [[[-0.65_13, 1.50_35, -0.27_66], [-0.65_15, 1.50_46, -0.27_80], [-0.65_12, 1.50_49, -0.27_84]]])

        self.assertTrue(torch.allclose(output[:, 1:4, 1:4] , __lowercase , atol=1E-4))






from ...configuration_utils import PretrainedConfig
from ...utils import logging


__lowercase							 =					logging.get_logger(__name__)

__lowercase							 =					{'''ctrl''': '''https://huggingface.co/ctrl/resolve/main/config.json'''}


class 	lowerCamelCase_     (   UpperCAmelCase_ ):






    '''simple docstring'''






    a__	:						str						=      """ctrl"""
    a__	:						Dict						=      ["""past_key_values"""]
    a__	:						Tuple						=      {
        """max_position_embeddings""": """n_positions""",
        """hidden_size""": """n_embd""",
        """num_attention_heads""": """n_head""",
        """num_hidden_layers""": """n_layer""",
    }







    def __init__(	self , __lowercase=246_534 , __lowercase=256 , __lowercase=1_280 , __lowercase=8_192 , __lowercase=48 , __lowercase=16 , __lowercase=0.1 , __lowercase=0.1 , __lowercase=1E-6 , __lowercase=0.02 , __lowercase=True , **__lowercase , )       ->	List[Any]:
        __UpperCamelCase     :List[str]			  =		vocab_size
        __UpperCamelCase     :Optional[Any]			  =		n_positions
        __UpperCamelCase     :Dict			  =		n_embd
        __UpperCamelCase     :Dict			  =		n_layer
        __UpperCamelCase     :List[Any]			  =		n_head
        __UpperCamelCase     :int			  =		dff
        __UpperCamelCase     :Union[str, Any]			  =		resid_pdrop
        __UpperCamelCase     :Optional[int]			  =		embd_pdrop
        __UpperCamelCase     :List[Any]			  =		layer_norm_epsilon
        __UpperCamelCase     :Dict			  =		initializer_range

        __UpperCamelCase     :Any			  =		use_cache

        super().__init__(**__lowercase)





