





'''simple docstring'''

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      :
 def __init__(							self :Dict		,			lowercase_ :Optional[int]		,			lowercase_ :List[Any]=13		,			lowercase_ :List[Any]=7		,			lowercase_ :List[Any]=True		,			lowercase_ :Any=True		,			lowercase_ :int=True		,			lowercase_ :Optional[Any]=True		,			lowercase_ :Tuple=99		,			lowercase_ :int=16		,			lowercase_ :Optional[Any]=36		,			lowercase_ :List[str]=6		,			lowercase_ :Optional[Any]=6		,			lowercase_ :str=6		,			lowercase_ :int=37		,			lowercase_ :Dict="gelu"		,			lowercase_ :Union[str, Any]=0.1		,			lowercase_ :List[str]=0.1		,			lowercase_ :Dict=5_12		,			lowercase_ :Optional[int]=16		,			lowercase_ :int=2		,			lowercase_ :Any=0.0_2		,			lowercase_ :Dict=3		,			lowercase_ :List[Any]=4		,			lowercase_ :Any=None		,			)->				Any:
     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 	UpperCAmelCase_				(							self :Optional[Any]				)->				str:
     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 	UpperCAmelCase_				(							self :List[str]				)->				List[str]:
     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 :Tuple		,			lowercase_ :Optional[int]		,			lowercase_ :Dict		,			lowercase_ :Tuple		,			lowercase_ :List[str]		,			lowercase_ :Tuple		,			lowercase_ :Tuple		,			lowercase_ :List[Any]				)->				str:
     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 	UpperCAmelCase_				(							self :int		,			lowercase_ :str		,			lowercase_ :int		,			lowercase_ :List[str]		,			lowercase_ :int		,			lowercase_ :Tuple		,			lowercase_ :Optional[int]		,			lowercase_ :Any				)->				Dict:
     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 	UpperCAmelCase_				(							self :List[Any]		,			lowercase_ :Optional[Any]		,			lowercase_ :Optional[int]		,			lowercase_ :str		,			lowercase_ :Union[str, Any]		,			lowercase_ :List[str]		,			lowercase_ :Dict		,			lowercase_ :Any				)->				int:
     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 	UpperCAmelCase_				(							self :List[Any]		,			lowercase_ :Optional[Any]		,			lowercase_ :Dict		,			lowercase_ :Optional[Any]		,			lowercase_ :Optional[Any]		,			lowercase_ :Union[str, Any]		,			lowercase_ :int		,			lowercase_ :List[Any]				)->				Tuple:
     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 	UpperCAmelCase_				(							self :Optional[int]		,			lowercase_ :List[Any]		,			lowercase_ :Optional[Any]		,			lowercase_ :str		,			lowercase_ :Optional[int]		,			lowercase_ :Tuple		,			lowercase_ :Any		,			lowercase_ :Optional[Any]				)->				Dict:
     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 	UpperCAmelCase_				(							self :List[Any]		,			lowercase_ :str		,			lowercase_ :int		,			lowercase_ :Optional[Any]		,			lowercase_ :List[str]		,			lowercase_ :Optional[int]		,			lowercase_ :Optional[Any]		,			lowercase_ :List[str]				)->				Dict:
     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 	UpperCAmelCase_				(							self :Union[str, Any]		,			lowercase_ :List[Any]		,			lowercase_ :Dict		,			lowercase_ :Tuple		,			lowercase_ :Tuple		,			lowercase_ :List[Any]		,			lowercase_ :List[Any]		,			lowercase_ :Union[str, Any]				)->				Union[str, Any]:
     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 	UpperCAmelCase_				(							self :Tuple				)->				List[Any]:
     A__							=     self.prepare_config_and_inputs()
     (
         (
         A__
     ),      (
         A__
     ),      (
         A__
     ),      (
         A__
     ),      (
         A__
     ),      (
         A__
     ),      (
         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   UpperCAmelCase      (			UpperCamelCase__     ,			UpperCamelCase__     ,			unittest.TestCase							):
 __lowercase  						=							(
     (
         AlbertModel,
         AlbertForPreTraining,
         AlbertForMaskedLM,
         AlbertForMultipleChoice,
         AlbertForSequenceClassification,
         AlbertForTokenClassification,
         AlbertForQuestionAnswering,
     )
     if is_torch_available()
     else ()
 )
 __lowercase  						=							(
     {
         """feature-extraction""": AlbertModel,
         """fill-mask""": AlbertForMaskedLM,
         """question-answering""": AlbertForQuestionAnswering,
         """text-classification""": AlbertForSequenceClassification,
         """token-classification""": AlbertForTokenClassification,
         """zero-shot""": AlbertForSequenceClassification,
     }
     if is_torch_available()
     else {}
 )
 __lowercase  						=							True
 def 	UpperCAmelCase_				(							self :int		,			lowercase_ :Optional[int]		,			lowercase_ :Any		,			lowercase_ :Optional[Any]=False				)->				Optional[int]:
     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 	UpperCAmelCase_				(							self :Optional[int]				)->				Optional[Any]:
     A__							=     AlbertModelTester(self				)
     A__							=     ConfigTester(self		,			config_class=lowercase_		,			hidden_size=37				)
 def 	UpperCAmelCase_				(							self :List[str]				)->				Dict:
     self.config_tester.run_common_tests()
 def 	UpperCAmelCase_				(							self :Dict				)->				Optional[int]:
     A__							=     self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_model(*lowercase_				)
 def 	UpperCAmelCase_				(							self :Optional[Any]				)->				Union[str, Any]:
     A__							=     self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_pretraining(*lowercase_				)
 def 	UpperCAmelCase_				(							self :str				)->				Optional[int]:
     A__							=     self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_masked_lm(*lowercase_				)
 def 	UpperCAmelCase_				(							self :List[str]				)->				Any:
     A__							=     self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_multiple_choice(*lowercase_				)
 def 	UpperCAmelCase_				(							self :str				)->				Any:
     A__							=     self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_question_answering(*lowercase_				)
 def 	UpperCAmelCase_				(							self :Optional[Any]				)->				List[str]:
     A__							=     self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_sequence_classification(*lowercase_				)
 def 	UpperCAmelCase_				(							self :List[Any]				)->				Union[str, Any]:
     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 	UpperCAmelCase_				(							self :Dict				)->				Tuple:
     for model_name in ALBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
         A__							=     AlbertModel.from_pretrained(lowercase_				)
         self.assertIsNotNone(lowercase_				)






@require_torch
class   UpperCAmelCase      (			unittest.TestCase							):
 @slow
 def 	UpperCAmelCase_				(							self :Optional[int]				)->				Any:
     A__							=     AlbertModel.from_pretrained("albert-base-v2"				)
     A__							=     torch.tensor([[0, 3_45, 2_32, 3_28, 7_40, 1_40, 16_95, 69, 60_78, 15_88, 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, 11, 7_68)				)
     self.assertEqual(output.shape		,			lowercase_				)
     A__							=     torch.tensor(
         [[[-0.6_5_1_3, 1.5_0_3_5, -0.2_7_6_6], [-0.6_5_1_5, 1.5_0_4_6, -0.2_7_8_0], [-0.6_5_1_2, 1.5_0_4_9, -0.2_7_8_4]]]				)

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










'''simple docstring'''

import os
import tempfile
import unittest

from transformers import is_torch_available
from transformers.testing_utils import require_torch


if is_torch_available():
  import torch
  from torch import nn

  from transformers import (
      Adafactor,
      AdamW,
      get_constant_schedule,
      get_constant_schedule_with_warmup,
      get_cosine_schedule_with_warmup,
      get_cosine_with_hard_restarts_schedule_with_warmup,
      get_inverse_sqrt_schedule,
      get_linear_schedule_with_warmup,
      get_polynomial_decay_schedule_with_warmup,
  )





def 		UpperCamelCase							(				_lowerCamelCase   :   Any		,			_lowerCamelCase   :   Optional[Any]=10			):
    A__							=     []
    for _ in range(_lowerCamelCase			):
        lrs.append(scheduler.get_lr()[0]			)
        scheduler.step()
    return lrs





def 		UpperCamelCase							(				_lowerCamelCase   :   Tuple		,			_lowerCamelCase   :   List[str]=10			):
    A__							=     []
    for step in range(_lowerCamelCase			):
        lrs.append(scheduler.get_lr()[0]			)
        scheduler.step()
        if step == num_steps // 2:
            with tempfile.TemporaryDirectory() as tmpdirname:
                A__							=     os.path.join(_lowerCamelCase		,			"schedule.bin"			)
                torch.save(scheduler.state_dict()		,			_lowerCamelCase			)

                A__							=     torch.load(_lowerCamelCase			)
                scheduler.load_state_dict(_lowerCamelCase			)
    return lrs






@require_torch
class   UpperCAmelCase      (			unittest.TestCase							):
 def 	UpperCAmelCase_				(							self :int		,			lowercase_ :Any		,			lowercase_ :List[Any]		,			lowercase_ :Optional[Any]				)->				int:
     self.assertEqual(len(lowercase_				)		,			len(lowercase_				)				)
     for a, b in zip(lowercase_		,			lowercase_				):
         self.assertAlmostEqual(lowercase_		,			lowercase_		,			delta=lowercase_				)
 def 	UpperCAmelCase_				(							self :Union[str, Any]				)->				Dict:
     A__							=     torch.tensor([0.1, -0.2, -0.1]		,			requires_grad=lowercase_				)
     A__							=     torch.tensor([0.4, 0.2, -0.5]				)
     A__							=     nn.MSELoss()
     # No warmup, constant schedule, no gradient clipping
     A__							=     AdamW(params=[w]		,			lr=2E-1		,			weight_decay=0.0				)
     for _ in range(1_00				):
         A__							=     criterion(lowercase_		,			lowercase_				)
         loss.backward()
         optimizer.step()
         w.grad.detach_()  # No zero_grad() function on simple tensors. we do it ourselves.
         w.grad.zero_()
     self.assertListAlmostEqual(w.tolist()		,			[0.4, 0.2, -0.5]		,			tol=1E-2				)



 def 	UpperCAmelCase_				(							self :Tuple				)->				List[str]:
     A__							=     torch.tensor([0.1, -0.2, -0.1]		,			requires_grad=lowercase_				)
     A__							=     torch.tensor([0.4, 0.2, -0.5]				)
     A__							=     nn.MSELoss()
     # No warmup, constant schedule, no gradient clipping
     A__							=     Adafactor(
         params=[w]		,			lr=1E-2		,			eps=(1E-30, 1E-3)		,			clip_threshold=1.0		,			decay_rate=-0.8		,			betaa=lowercase_		,			weight_decay=0.0		,			relative_step=lowercase_		,			scale_parameter=lowercase_		,			warmup_init=lowercase_		,			)
     for _ in range(10_00				):
         A__							=     criterion(lowercase_		,			lowercase_				)
         loss.backward()
         optimizer.step()
         w.grad.detach_()  # No zero_grad() function on simple tensors. we do it ourselves.
         w.grad.zero_()
     self.assertListAlmostEqual(w.tolist()		,			[0.4, 0.2, -0.5]		,			tol=1E-2				)






@require_torch
class   UpperCAmelCase      (			unittest.TestCase							):
 __lowercase  						=							nn.Linear(50     ,			50							) if is_torch_available() else None
 __lowercase  						=							AdamW(m.parameters()     ,			lr=1_0.0							) if is_torch_available() else None
 __lowercase  						=							10
 def 	UpperCAmelCase_				(							self :Tuple		,			lowercase_ :Any		,			lowercase_ :List[Any]		,			lowercase_ :List[Any]		,			lowercase_ :str=None				)->				Optional[int]:
     self.assertEqual(len(lowercase_				)		,			len(lowercase_				)				)
     for a, b in zip(lowercase_		,			lowercase_				):
         self.assertAlmostEqual(lowercase_		,			lowercase_		,			delta=lowercase_		,			msg=lowercase_				)



 def 	UpperCAmelCase_				(							self :Optional[Any]				)->				Any:
     A__							=     {"num_warmup_steps": 2, "num_training_steps": 10}
     # schedulers doct format
     # function: (sched_args_dict, expected_learning_rates)
     A__							=     {
         get_constant_schedule: ({}, [1_0.0] * self.num_steps),
         get_constant_schedule_with_warmup: (
             {"num_warmup_steps": 4},
             [0.0, 2.5, 5.0, 7.5, 1_0.0, 1_0.0, 1_0.0, 1_0.0, 1_0.0, 1_0.0],
         ),
         get_linear_schedule_with_warmup: (
             {**common_kwargs},
             [0.0, 5.0, 1_0.0, 8.7_5, 7.5, 6.2_5, 5.0, 3.7_5, 2.5, 1.2_5],
         ),
         get_cosine_schedule_with_warmup: (
             {**common_kwargs},
             [0.0, 5.0, 1_0.0, 9.6_1, 8.5_3, 6.9_1, 5.0, 3.0_8, 1.4_6, 0.3_8],
         ),
         get_cosine_with_hard_restarts_schedule_with_warmup: (
             {**common_kwargs, "num_cycles": 2},
             [0.0, 5.0, 1_0.0, 8.5_3, 5.0, 1.4_6, 1_0.0, 8.5_3, 5.0, 1.4_6],
         ),
         get_polynomial_decay_schedule_with_warmup: (
             {**common_kwargs, "power": 2.0, "lr_end": 1E-7},
             [0.0, 5.0, 1_0.0, 7.6_5_6, 5.6_2_5, 3.9_0_6, 2.5, 1.4_0_6, 0.6_2_5, 0.1_5_6],
         ),
         get_inverse_sqrt_schedule: (
             {"num_warmup_steps": 2},
             [0.0, 5.0, 1_0.0, 8.1_6_5, 7.0_7_1, 6.3_2_5, 5.7_7_4, 5.3_4_5, 5.0, 4.7_1_4],
         ),
     }

     for scheduler_func, data in scheds.items():
         A__,      A__							=     data

         A__							=     scheduler_func(self.optimizer		,			**lowercase_				)
         self.assertEqual(len([scheduler.get_lr()[0]]				)		,			1				)
         A__							=     unwrap_schedule(lowercase_		,			self.num_steps				)
         self.assertListAlmostEqual(
             lowercase_		,			lowercase_		,			tol=1E-2		,			msg=F"failed for {scheduler_func} in normal scheduler"		,			)

         A__							=     scheduler_func(self.optimizer		,			**lowercase_				)
         if scheduler_func.__name__ != "get_constant_schedule":
             LambdaScheduleWrapper.wrap_scheduler(lowercase_				)  # wrap to test picklability of the schedule
         A__							=     unwrap_and_save_reload_schedule(lowercase_		,			self.num_steps				)
         self.assertListEqual(lowercase_		,			lowercase_		,			msg=F"failed for {scheduler_func} in save and reload"				)






class   UpperCAmelCase      :
 def __init__(							self :str		,			lowercase_ :List[str]				)->				Tuple:
     A__							=     fn
 def __call__(							self :List[Any]		,			*lowercase_ :Dict		,			**lowercase_ :Dict				)->				Tuple:
     return self.fn(*lowercase_		,			**lowercase_				)



 @classmethod
 def 	UpperCAmelCase_				(							self :Any		,			lowercase_ :Tuple				)->				List[Any]:
     A__							=     list(map(self		,			scheduler.lr_lambdas				)				)



