



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





 """simple docstring"""



 def __init__(    self						,     a						,     a=1_3						,     a=7						,     a=True						,     a=True						,     a=True						,     a=True						,     a=9_9						,     a=1_6						,     a=3_6						,     a=6						,     a=6						,     a=6						,     a=3_7						,     a="gelu"						,     a=0.1						,     a=0.1						,     a=5_1_2						,     a=1_6						,     a=2						,     a=0.02						,     a=3						,     a=4						,     a=None						,     )						->				int:




   """simple docstring"""







   _A							     =     parent
   _A							     =     batch_size
   _A							     =     seq_length
   _A							     =     is_training
   _A							     =     use_input_mask
   _A							     =     use_token_type_ids
   _A							     =     use_labels
   _A							     =     vocab_size
   _A							     =     embedding_size
   _A							     =     hidden_size
   _A							     =     num_hidden_layers
   _A							     =     num_hidden_groups
   _A							     =     num_attention_heads
   _A							     =     intermediate_size
   _A							     =     hidden_act
   _A							     =     hidden_dropout_prob
   _A							     =     attention_probs_dropout_prob
   _A							     =     max_position_embeddings
   _A							     =     type_vocab_size
   _A							     =     type_sequence_label_size
   _A							     =     initializer_range
   _A							     =     num_labels
   _A							     =     num_choices
   _A							     =     scope
 def lowercase_		(    self       )						->				Optional[int]:




   """simple docstring"""







   _A							     =     ids_tensor([self.batch_size, self.seq_length]						,     self.vocab_size       )

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

   _A							     =     None
   if self.use_token_type_ids:
     _A							     =     ids_tensor([self.batch_size, self.seq_length]						,     self.type_vocab_size       )

   _A							     =     None
   _A							     =     None
   _A							     =     None
   if self.use_labels:
     _A							     =     ids_tensor([self.batch_size]						,     self.type_sequence_label_size       )
     _A							     =     ids_tensor([self.batch_size, self.seq_length]						,     self.num_labels       )
     _A							     =     ids_tensor([self.batch_size]						,     self.num_choices       )

   _A							     =     self.get_config()

   return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels
 def lowercase_		(    self       )						->				Optional[int]:




   """simple docstring"""







   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 lowercase_		(    self						,     a						,     a						,     a						,     a						,     a						,     a						,     a       )						->				str:




   """simple docstring"""







   _A							     =     AlbertModel(config=_lowercase       )
   model.to(_lowercase       )
   model.eval()
   _A							     =     model(_lowercase						,     attention_mask=_lowercase						,     token_type_ids=_lowercase       )
   _A							     =     model(_lowercase						,     token_type_ids=_lowercase       )
   _A							     =     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 lowercase_		(    self						,     a						,     a						,     a						,     a						,     a						,     a						,     a       )						->				Any:




   """simple docstring"""







   _A							     =     AlbertForPreTraining(config=_lowercase       )
   model.to(_lowercase       )
   model.eval()
   _A							     =     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 lowercase_		(    self						,     a						,     a						,     a						,     a						,     a						,     a						,     a       )						->				Optional[Any]:




   """simple docstring"""







   _A							     =     AlbertForMaskedLM(config=_lowercase       )
   model.to(_lowercase       )
   model.eval()
   _A							     =     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 lowercase_		(    self						,     a						,     a						,     a						,     a						,     a						,     a						,     a       )						->				Optional[int]:




   """simple docstring"""







   _A							     =     AlbertForQuestionAnswering(config=_lowercase       )
   model.to(_lowercase       )
   model.eval()
   _A							     =     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 lowercase_		(    self						,     a						,     a						,     a						,     a						,     a						,     a						,     a       )						->				str:




   """simple docstring"""







   _A							     =     self.num_labels
   _A							     =     AlbertForSequenceClassification(_lowercase       )
   model.to(_lowercase       )
   model.eval()
   _A							     =     model(_lowercase						,     attention_mask=_lowercase						,     token_type_ids=_lowercase						,     labels=_lowercase       )
   self.parent.assertEqual(result.logits.shape						,     (self.batch_size, self.num_labels)       )
 def lowercase_		(    self						,     a						,     a						,     a						,     a						,     a						,     a						,     a       )						->				Optional[Any]:




   """simple docstring"""







   _A							     =     self.num_labels
   _A							     =     AlbertForTokenClassification(config=_lowercase       )
   model.to(_lowercase       )
   model.eval()
   _A							     =     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 lowercase_		(    self						,     a						,     a						,     a						,     a						,     a						,     a						,     a       )						->				int:




   """simple docstring"""







   _A							     =     self.num_choices
   _A							     =     AlbertForMultipleChoice(config=_lowercase       )
   model.to(_lowercase       )
   model.eval()
   _A							     =     input_ids.unsqueeze(1       ).expand(-1						,     self.num_choices						,     -1       ).contiguous()
   _A							     =     token_type_ids.unsqueeze(1       ).expand(-1						,     self.num_choices						,     -1       ).contiguous()
   _A							     =     input_mask.unsqueeze(1       ).expand(-1						,     self.num_choices						,     -1       ).contiguous()
   _A							     =     model(
       _lowercase						,     attention_mask=_lowercase						,     token_type_ids=_lowercase						,     labels=_lowercase						,     )
   self.parent.assertEqual(result.logits.shape						,     (self.batch_size, self.num_choices)       )



 def lowercase_		(    self       )						->				Tuple:




   """simple docstring"""







   _A							     =     self.prepare_config_and_inputs()
   (
       _A				
   )							     =     config_and_inputs
   _A							     =     {'''input_ids''': input_ids, '''token_type_ids''': token_type_ids, '''attention_mask''': input_mask}
   return config, inputs_dict







@require_torch
class 		_snake_case   (      __a       ,__a       ,unittest.TestCase					):





 """simple docstring"""



 lowerCamelCase_        =      (
     (
         AlbertModel,
         AlbertForPreTraining,
         AlbertForMaskedLM,
         AlbertForMultipleChoice,
         AlbertForSequenceClassification,
         AlbertForTokenClassification,
         AlbertForQuestionAnswering,
     )
     if is_torch_available()
     else ()
 )
 lowerCamelCase_        =      (
     {
         '''feature-extraction''': AlbertModel,
         '''fill-mask''': AlbertForMaskedLM,
         '''question-answering''': AlbertForQuestionAnswering,
         '''text-classification''': AlbertForSequenceClassification,
         '''token-classification''': AlbertForTokenClassification,
         '''zero-shot''': AlbertForSequenceClassification,
     }
     if is_torch_available()
     else {}
 )
 lowerCamelCase_        =      True
 def lowercase_		(    self						,     a						,     a						,     a=False       )						->				Union[str, Any]:




   """simple docstring"""







   _A							     =     super()._prepare_for_class(_lowercase						,     _lowercase						,     return_labels=_lowercase       )

   if return_labels:
     if model_class in get_values(_lowercase       ):
       _A							     =     torch.zeros(
           (self.model_tester.batch_size, self.model_tester.seq_length)						,     dtype=torch.long						,     device=_lowercase       )
       _A							     =     torch.zeros(
           self.model_tester.batch_size						,     dtype=torch.long						,     device=_lowercase       )
   return inputs_dict
 def lowercase_		(    self       )						->				Any:




   """simple docstring"""







   _A							     =     AlbertModelTester(self       )
   _A							     =     ConfigTester(self						,     config_class=_lowercase						,     hidden_size=3_7       )
 def lowercase_		(    self       )						->				List[Any]:




   """simple docstring"""







   self.config_tester.run_common_tests()
 def lowercase_		(    self       )						->				Optional[Any]:




   """simple docstring"""







   _A							     =     self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_model(*_lowercase       )
 def lowercase_		(    self       )						->				str:




   """simple docstring"""







   _A							     =     self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_for_pretraining(*_lowercase       )
 def lowercase_		(    self       )						->				int:




   """simple docstring"""







   _A							     =     self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_for_masked_lm(*_lowercase       )
 def lowercase_		(    self       )						->				Optional[Any]:




   """simple docstring"""







   _A							     =     self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_for_multiple_choice(*_lowercase       )
 def lowercase_		(    self       )						->				Dict:




   """simple docstring"""







   _A							     =     self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_for_question_answering(*_lowercase       )
 def lowercase_		(    self       )						->				Tuple:




   """simple docstring"""







   _A							     =     self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_for_sequence_classification(*_lowercase       )
 def lowercase_		(    self       )						->				Any:




   """simple docstring"""







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



 @slow
 def lowercase_		(    self       )						->				List[Any]:




   """simple docstring"""







   for model_name in ALBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
     _A							     =     AlbertModel.from_pretrained(_lowercase       )
     self.assertIsNotNone(_lowercase       )







@require_torch
class 		_snake_case   (      unittest.TestCase					):





 """simple docstring"""



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




   """simple docstring"""







   _A							     =     AlbertModel.from_pretrained('''albert-base-v2'''       )
   _A							     =     torch.tensor([[0, 3_4_5, 2_3_2, 3_2_8, 7_4_0, 1_4_0, 1_6_9_5, 6_9, 6_0_7_8, 1_5_8_8, 2]]       )
   _A							     =     torch.tensor([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]       )
   with torch.no_grad():
     _A							     =     model(_lowercase						,     attention_mask=_lowercase       )[0]
   _A							     =     torch.Size((1, 1_1, 7_6_8)       )
   self.assertEqual(output.shape						,     _lowercase       )
   _A							     =     torch.tensor(
       [[[-0.6_513, 1.5_035, -0.2_766], [-0.6_515, 1.5_046, -0.2_780], [-0.6_512, 1.5_049, -0.2_784]]]       )

   self.assertTrue(torch.allclose(output[:, 1:4, 1:4]						,     _lowercase						,     atol=1e-4       )       )
from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_torch_available,
)


UpperCAmelCase_           =				{
    '''configuration_mega''': ['''MEGA_PRETRAINED_CONFIG_ARCHIVE_MAP''', '''MegaConfig''', '''MegaOnnxConfig'''],
}

try:
  if not is_torch_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  UpperCAmelCase_           =				[
      '''MEGA_PRETRAINED_MODEL_ARCHIVE_LIST''',
      '''MegaForCausalLM''',
      '''MegaForMaskedLM''',
      '''MegaForMultipleChoice''',
      '''MegaForQuestionAnswering''',
      '''MegaForSequenceClassification''',
      '''MegaForTokenClassification''',
      '''MegaModel''',
      '''MegaPreTrainedModel''',
  ]

if TYPE_CHECKING:
  from .configuration_mega import MEGA_PRETRAINED_CONFIG_ARCHIVE_MAP, MegaConfig, MegaOnnxConfig

  try:
    if not is_torch_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .modeling_mega import (
        MEGA_PRETRAINED_MODEL_ARCHIVE_LIST,
        MegaForCausalLM,
        MegaForMaskedLM,
        MegaForMultipleChoice,
        MegaForQuestionAnswering,
        MegaForSequenceClassification,
        MegaForTokenClassification,
        MegaModel,
        MegaPreTrainedModel,
    )

else:
  import sys

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