





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







					'''simple docstring'''

					def __init__(						self    :  List[Any] ,					__lowercase    :  Tuple ,					__lowercase    :  Dict=13 ,					__lowercase    :  Optional[Any]=7 ,					__lowercase    :  Any=True ,					__lowercase    :  str=True ,					__lowercase    :  Any=True ,					__lowercase    :  Optional[int]=True ,					__lowercase    :  Tuple=99 ,					__lowercase    :  Union[str, Any]=16 ,					__lowercase    :  List[str]=36 ,					__lowercase    :  List[str]=6 ,					__lowercase    :  List[str]=6 ,					__lowercase    :  str=6 ,					__lowercase    :  Tuple=37 ,					__lowercase    :  Optional[int]="gelu" ,					__lowercase    :  Dict=0.1 ,					__lowercase    :  Optional[Any]=0.1 ,					__lowercase    :  Optional[int]=5_12 ,					__lowercase    :  str=16 ,					__lowercase    :  Dict=2 ,					__lowercase    :  Any=0.02 ,					__lowercase    :  List[Any]=3 ,					__lowercase    :  Tuple=4 ,					__lowercase    :  Tuple=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_								=     embedding_size
						UpperCAmelCase_								=     hidden_size
						UpperCAmelCase_								=     num_hidden_layers
						UpperCAmelCase_								=     num_hidden_groups
						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        SCREAMING_SNAKE_CASE							(						self    :  List[str]		):






						'''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
						if self.use_token_type_ids:
							UpperCAmelCase_								=     ids_tensor([self.batch_size, self.seq_length] ,					self.type_vocab_size		)

						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, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels





					def        SCREAMING_SNAKE_CASE							(						self    :  List[Any]		):






						'''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        SCREAMING_SNAKE_CASE							(						self    :  List[Any] ,					__lowercase    :  Dict ,					__lowercase    :  int ,					__lowercase    :  Union[str, Any] ,					__lowercase    :  Tuple ,					__lowercase    :  Dict ,					__lowercase    :  Tuple ,					__lowercase    :  List[Any]		):






						'''simple docstring'''







						UpperCAmelCase_								=     AlbertModel(config=__lowercase		)
						model.to(__lowercase		)
						model.eval()
						UpperCAmelCase_								=     model(__lowercase ,					attention_mask=__lowercase ,					token_type_ids=__lowercase		)
						UpperCAmelCase_								=     model(__lowercase ,					token_type_ids=__lowercase		)
						UpperCAmelCase_								=     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        SCREAMING_SNAKE_CASE							(						self    :  Any ,					__lowercase    :  Tuple ,					__lowercase    :  Tuple ,					__lowercase    :  int ,					__lowercase    :  Optional[Any] ,					__lowercase    :  Optional[Any] ,					__lowercase    :  int ,					__lowercase    :  Any		):






						'''simple docstring'''







						UpperCAmelCase_								=     AlbertForPreTraining(config=__lowercase		)
						model.to(__lowercase		)
						model.eval()
						UpperCAmelCase_								=     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        SCREAMING_SNAKE_CASE							(						self    :  Union[str, Any] ,					__lowercase    :  List[Any] ,					__lowercase    :  Optional[int] ,					__lowercase    :  List[str] ,					__lowercase    :  Optional[int] ,					__lowercase    :  List[Any] ,					__lowercase    :  Optional[Any] ,					__lowercase    :  List[Any]		):






						'''simple docstring'''







						UpperCAmelCase_								=     AlbertForMaskedLM(config=__lowercase		)
						model.to(__lowercase		)
						model.eval()
						UpperCAmelCase_								=     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        SCREAMING_SNAKE_CASE							(						self    :  Any ,					__lowercase    :  Optional[int] ,					__lowercase    :  Optional[int] ,					__lowercase    :  str ,					__lowercase    :  Any ,					__lowercase    :  Any ,					__lowercase    :  Any ,					__lowercase    :  Union[str, Any]		):






						'''simple docstring'''







						UpperCAmelCase_								=     AlbertForQuestionAnswering(config=__lowercase		)
						model.to(__lowercase		)
						model.eval()
						UpperCAmelCase_								=     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        SCREAMING_SNAKE_CASE							(						self    :  Dict ,					__lowercase    :  List[Any] ,					__lowercase    :  int ,					__lowercase    :  Optional[int] ,					__lowercase    :  Optional[int] ,					__lowercase    :  List[Any] ,					__lowercase    :  List[Any] ,					__lowercase    :  str		):






						'''simple docstring'''







						UpperCAmelCase_								=     self.num_labels
						UpperCAmelCase_								=     AlbertForSequenceClassification(__lowercase		)
						model.to(__lowercase		)
						model.eval()
						UpperCAmelCase_								=     model(__lowercase ,					attention_mask=__lowercase ,					token_type_ids=__lowercase ,					labels=__lowercase		)
						self.parent.assertEqual(result.logits.shape ,					(self.batch_size, self.num_labels)		)





					def        SCREAMING_SNAKE_CASE							(						self    :  Optional[int] ,					__lowercase    :  Dict ,					__lowercase    :  Union[str, Any] ,					__lowercase    :  Union[str, Any] ,					__lowercase    :  Tuple ,					__lowercase    :  Union[str, Any] ,					__lowercase    :  Any ,					__lowercase    :  List[str]		):






						'''simple docstring'''







						UpperCAmelCase_								=     self.num_labels
						UpperCAmelCase_								=     AlbertForTokenClassification(config=__lowercase		)
						model.to(__lowercase		)
						model.eval()
						UpperCAmelCase_								=     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        SCREAMING_SNAKE_CASE							(						self    :  Tuple ,					__lowercase    :  Dict ,					__lowercase    :  Optional[Any] ,					__lowercase    :  Dict ,					__lowercase    :  int ,					__lowercase    :  str ,					__lowercase    :  Optional[int] ,					__lowercase    :  List[str]		):






						'''simple docstring'''







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





					def        SCREAMING_SNAKE_CASE							(						self    :  Any		):






						'''simple docstring'''







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







@require_torch
class   _UpperCamelCase  (       A_						,       A_						,       unittest.TestCase  ):







					'''simple docstring'''

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





					def        SCREAMING_SNAKE_CASE							(						self    :  Any ,					__lowercase    :  Union[str, Any] ,					__lowercase    :  List[Any] ,					__lowercase    :  Optional[int]=False		):






						'''simple docstring'''







						UpperCAmelCase_								=     super()._prepare_for_class(__lowercase ,					__lowercase ,					return_labels=__lowercase		)

						if return_labels:
							if model_class in get_values(__lowercase		):
								UpperCAmelCase_								=     torch.zeros(
								    (self.model_tester.batch_size, self.model_tester.seq_length) ,					dtype=torch.long ,					device=__lowercase		)
								UpperCAmelCase_								=     torch.zeros(
								    self.model_tester.batch_size ,					dtype=torch.long ,					device=__lowercase		)
						return inputs_dict





					def        SCREAMING_SNAKE_CASE							(						self    :  int		):






						'''simple docstring'''







						UpperCAmelCase_								=     AlbertModelTester(self		)
						UpperCAmelCase_								=     ConfigTester(self ,					config_class=__lowercase ,					hidden_size=37		)





					def        SCREAMING_SNAKE_CASE							(						self    :  Optional[int]		):






						'''simple docstring'''







						self.config_tester.run_common_tests()





					def        SCREAMING_SNAKE_CASE							(						self    :  str		):






						'''simple docstring'''







						UpperCAmelCase_								=     self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_model(*__lowercase		)





					def        SCREAMING_SNAKE_CASE							(						self    :  Optional[int]		):






						'''simple docstring'''







						UpperCAmelCase_								=     self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_pretraining(*__lowercase		)





					def        SCREAMING_SNAKE_CASE							(						self    :  Union[str, Any]		):






						'''simple docstring'''







						UpperCAmelCase_								=     self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_masked_lm(*__lowercase		)





					def        SCREAMING_SNAKE_CASE							(						self    :  Optional[Any]		):






						'''simple docstring'''







						UpperCAmelCase_								=     self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_multiple_choice(*__lowercase		)





					def        SCREAMING_SNAKE_CASE							(						self    :  Union[str, Any]		):






						'''simple docstring'''







						UpperCAmelCase_								=     self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_question_answering(*__lowercase		)





					def        SCREAMING_SNAKE_CASE							(						self    :  Optional[int]		):






						'''simple docstring'''







						UpperCAmelCase_								=     self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_sequence_classification(*__lowercase		)





					def        SCREAMING_SNAKE_CASE							(						self    :  Optional[int]		):






						'''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(*__lowercase		)





					@slow
					def        SCREAMING_SNAKE_CASE							(						self    :  Dict		):






						'''simple docstring'''







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







@require_torch
class   _UpperCamelCase  (       unittest.TestCase  ):







					'''simple docstring'''

					@slow
					def        SCREAMING_SNAKE_CASE							(						self    :  Tuple		):






						'''simple docstring'''







						UpperCAmelCase_								=     AlbertModel.from_pretrained("""albert-base-v2"""		)
						UpperCAmelCase_								=     torch.tensor([[0, 3_45, 2_32, 3_28, 7_40, 1_40, 16_95, 69, 60_78, 15_88, 2]]		)
						UpperCAmelCase_								=     torch.tensor([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]		)
						with torch.no_grad():
							UpperCAmelCase_								=     model(__lowercase ,					attention_mask=__lowercase		)[0]
						UpperCAmelCase_								=     torch.Size((1, 11, 7_68)		)
						self.assertEqual(output.shape ,					__lowercase		)
						UpperCAmelCase_								=     torch.tensor(
						    [[[-0.6513, 1.5035, -0.2766], [-0.6515, 1.5046, -0.2780], [-0.6512, 1.5049, -0.2784]]]		)

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









from torch import nn







class   _UpperCamelCase  (       nn.Module  ):







					'''simple docstring'''

					def __init__(						self    :  Dict ,					__lowercase    :  List[str] ,					__lowercase    :  Dict		):






						'''simple docstring'''







						super().__init__()
						UpperCAmelCase_								=     class_size
						UpperCAmelCase_								=     embed_size
						# self.mlp1 = nn.Linear(embed_size, embed_size)
						# self.mlp2 = (nn.Linear(embed_size, class_size))
						UpperCAmelCase_								=     nn.Linear(__lowercase ,					__lowercase		)





					def        SCREAMING_SNAKE_CASE							(						self    :  Any ,					__lowercase    :  Optional[Any]		):






						'''simple docstring'''







						UpperCAmelCase_								=     self.mlp(__lowercase		)
						return logits


