


'''simple docstring'''
import math
import os
import re
import sys
import unittest
from pathlib import Path
from typing import Tuple
from unittest.mock import patch

from parameterized import parameterized

from transformers.testing_utils import (
    CaptureStderr,
    ExtendSysPath,
    TestCasePlus,
    execute_subprocess_async,
    get_gpu_count,
    get_torch_dist_unique_port,
    require_apex,
    require_bitsandbytes,
    require_fairscale,
    require_torch,
    require_torch_gpu,
    require_torch_multi_gpu,
    require_torch_non_multi_gpu,
    slow,
)
from transformers.trainer_callback import TrainerState
from transformers.trainer_utils import set_seed


UpperCamelCase__          = os.path.abspath(os.path.dirname(__file__))
with ExtendSysPath(F"""{bindir}/../../examples/pytorch/translation"""):
  from run_translation import main  # noqa


set_seed(4_2)
UpperCamelCase__          = '''sshleifer/student_marian_en_ro_6_1'''
UpperCamelCase__          = '''sshleifer/tiny-mbart'''




@require_torch
class      lowerCamelCase_  (      __a ):






      def lowercase_				(   self							:       Dict		,		_A							:       int=False		,		_A							:       Optional[Any]=None		,		_A							:       Optional[Any]=True		,		_A							:       Dict=True		,		_A							:       Dict=True		,		_A							:       List[Any]=True		,		):


          '''simple docstring'''


          UpperCAmelCase__     :						Union[str, Any]						=							self.run_trainer(
              eval_steps=1		,		max_len=12		,		model_name=_A		,		num_train_epochs=1		,		distributed=_A		,		extra_args_str=_A		,		predict_with_generate=_A		,		do_train=_A		,		do_eval=_A		,		do_predict=_A		,		)
          UpperCAmelCase__     :						Optional[Any]						=							TrainerState.load_from_json(os.path.join(_A		,		'''trainer_state.json'''					)					).log_history

          if not do_eval:
              return

          UpperCAmelCase__     :						Tuple						=							[log for log in logs if '''eval_loss''' in log.keys()]

          UpperCAmelCase__     :						List[str]						=							eval_metrics[0]
          if predict_with_generate:
              assert "eval_bleu" in first_step_stats

              UpperCAmelCase__     :						Dict						=							eval_metrics[-1]
              assert isinstance(last_step_stats['''eval_bleu''']		,		_A					)
              assert not math.isnan(float(last_step_stats['''eval_loss''']					)					), "eval_loss must not be `nan`"






      @require_torch_non_multi_gpu
      def lowercase_				(   self							:       Optional[Any]					):


          '''simple docstring'''


          self.run_seqaseq_quick()






      @require_torch_multi_gpu
      def lowercase_				(   self							:       Optional[int]					):


          '''simple docstring'''


          self.run_seqaseq_quick(distributed=_A					)






      @require_torch_multi_gpu
      def lowercase_				(   self							:       str					):


          '''simple docstring'''


          self.run_seqaseq_quick(distributed=_A					)






      @unittest.skip('''Requires an update of the env running those tests'''					)
      @require_torch_multi_gpu
      @require_fairscale
      def lowercase_				(   self							:       List[Any]					):


          '''simple docstring'''


          self.run_seqaseq_quick(distributed=_A		,		extra_args_str='''--sharded_ddp simple'''					)






      @unittest.skip('''Requires an update of the env running those tests'''					)
      @require_torch_multi_gpu
      @require_fairscale
      def lowercase_				(   self							:       List[Any]					):


          '''simple docstring'''


          self.run_seqaseq_quick(distributed=_A		,		extra_args_str='''--sharded_ddp simple --fp16'''					)






      @unittest.skip('''Requires an update of the env running those tests'''					)
      @require_torch_multi_gpu
      @require_fairscale
      def lowercase_				(   self							:       List[Any]					):


          '''simple docstring'''


          self.run_seqaseq_quick(distributed=_A		,		extra_args_str='''--sharded_ddp zero_dp_2'''		,		predict_with_generate=_A					)






      @unittest.skip('''Requires an update of the env running those tests'''					)
      @require_torch_multi_gpu
      @require_fairscale
      def lowercase_				(   self							:       Any					):


          '''simple docstring'''


          self.run_seqaseq_quick(
              distributed=_A		,		extra_args_str='''--sharded_ddp zero_dp_2 --fp16'''		,		predict_with_generate=_A					)






      @require_apex
      @require_torch_gpu
      def lowercase_				(   self							:       List[str]					):


          '''simple docstring'''


          self.run_seqaseq_quick(distributed=_A		,		extra_args_str='''--fp16 --fp16_backend=apex'''					)
          # test 2nd time - was getting eval_loss': nan'
          # to reproduce the problem set distributed=False
          self.run_seqaseq_quick(distributed=_A		,		extra_args_str='''--fp16 --fp16_backend=apex'''					)






      @parameterized.expand(['''base''', '''low''', '''high''', '''mixed''']					)
      @require_torch_multi_gpu
      def lowercase_				(   self							:       str		,		_A							:       Dict					):


          '''simple docstring'''


          UpperCAmelCase__     :						Dict						=							{
              # test with the default log_level - should be info and thus log info once
              '''base''': {'''extra_args_str''': '''''', '''n_matches''': 1},
              # test with low log_level and log_level_replica - should be noisy on all processes
              # now the info string should appear twice on 2 processes
              '''low''': {'''extra_args_str''': '''--log_level debug --log_level_replica debug''', '''n_matches''': 2},
              # test with high log_level and low log_level_replica
              # now the info string should appear once only on the replica
              '''high''': {'''extra_args_str''': '''--log_level error --log_level_replica debug''', '''n_matches''': 1},
              # test with high log_level and log_level_replica - should be quiet on all processes
              '''mixed''': {'''extra_args_str''': '''--log_level error --log_level_replica error''', '''n_matches''': 0},
          }

          UpperCAmelCase__     :						int						=							experiments[experiment_id]
          UpperCAmelCase__     :						int						=							{'''distributed''': True, '''predict_with_generate''': False, '''do_eval''': False, '''do_predict''': False}
          UpperCAmelCase__     :						Optional[int]						=							'''Running training'''
          with CaptureStderr() as cl:
              self.run_seqaseq_quick(**_A		,		extra_args_str=data['''extra_args_str''']					)
          UpperCAmelCase__     :						List[Any]						=							len(re.findall(_A		,		cl.err					)					)
          self.assertEqual(_A		,		data['''n_matches''']					)






      @slow
      def lowercase_				(   self							:       int					):


          '''simple docstring'''


          UpperCAmelCase__     :						List[Any]						=							self.run_trainer(
              eval_steps=2		,		max_len=128		,		model_name=_A		,		learning_rate=3e-4		,		num_train_epochs=10		,		distributed=_A		,		)

          # Check metrics
          UpperCAmelCase__     :						List[Any]						=							TrainerState.load_from_json(os.path.join(_A		,		'''trainer_state.json'''					)					).log_history
          UpperCAmelCase__     :						int						=							[log for log in logs if '''eval_loss''' in log.keys()]
          UpperCAmelCase__     :						Dict						=							eval_metrics[0]
          UpperCAmelCase__     :						int						=							eval_metrics[-1]

          assert first_step_stats["eval_loss"] > last_step_stats["eval_loss"], "model learned nothing"
          assert isinstance(last_step_stats['''eval_bleu''']		,		_A					)

          # test if do_predict saves generations and metrics
          UpperCAmelCase__     :						Tuple						=							os.listdir(_A					)
          UpperCAmelCase__     :						Optional[Any]						=							{os.path.basename(_A					) for p in contents}
          assert "generated_predictions.txt" in contents
          assert "predict_results.json" in contents






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


          '''simple docstring'''


          from transformers.training_args import OptimizerNames

          def train_and_return_metrics(_A							:       str					) -> Tuple[int, float]:
              UpperCAmelCase__     :						Any						=							'''--skip_memory_metrics 0'''

              UpperCAmelCase__     :						List[Any]						=							self.run_trainer(
                  max_len=128		,		model_name=_A		,		learning_rate=3e-4		,		num_train_epochs=1		,		optim=_A		,		distributed=_A		,		extra_args_str=_A		,		do_eval=_A		,		do_predict=_A		,		n_gpus_to_use=1		,		)

              # Check metrics
              UpperCAmelCase__     :						Optional[Any]						=							TrainerState.load_from_json(Path(_A		,		'''trainer_state.json'''					)					).log_history
              UpperCAmelCase__     :						Optional[int]						=							int(logs[0]['''train_mem_gpu_peaked_delta'''] / 2**20					)
              UpperCAmelCase__     :						str						=							int(logs[0]['''train_mem_gpu_alloc_delta'''] / 2**20					)

              UpperCAmelCase__     :						Any						=							logs[0]['''train_loss''']
              return gpu_peak_mem_mb, gpu_alloc_mem_mb, loss

          UpperCAmelCase__		,		UpperCAmelCase__		,		UpperCAmelCase__     :						List[Any]						=							train_and_return_metrics(OptimizerNames.ADAMW_TORCH.value					)
          UpperCAmelCase__		,		UpperCAmelCase__		,		UpperCAmelCase__     :						Optional[int]						=							train_and_return_metrics(OptimizerNames.ADAMW_BNB.value					)

          UpperCAmelCase__     :						Optional[Any]						=							gpu_alloc_mem_orig - gpu_alloc_mem_bnb

          UpperCAmelCase__     :						Union[str, Any]						=							gpu_peak_mem_orig + gpu_alloc_mem_orig
          UpperCAmelCase__     :						Optional[Any]						=							gpu_peak_mem_bnb + gpu_alloc_mem_bnb
          UpperCAmelCase__     :						List[str]						=							gpu_total_mem_orig - gpu_total_mem_bnb

          # sshleifer/student_marian_en_ro_6_1 has 54M parameter, 29M of which is `nn.Embedding` which
          # doesn't get quantized and remains in fp32. Therefore we only have 25M parameters quantized
          # in 2 bytes and the diff in optim memory usage is derived as so:
          #
          # - normal 25*8=~200MB (8 bytes per param)
          # - bnb    25*2= ~50MB (2 bytes per param)
          #
          # Thus we should expect ~150MB total memory saved.
          #
          # Peak memory should be the same - the total should be different by about that same margin
          #
          # After leaving a small margin to accommodate for differences between gpus let's check
          # that we have at least 120MB in savings
          UpperCAmelCase__     :						Union[str, Any]						=							120

          # uncomment the following if this test starts failing - requires py38 for a new print feature
          # gpu_peak_mem_diff = gpu_peak_mem_orig - gpu_peak_mem_bnb
          # print(f"{gpu_alloc_mem_orig=}MB {gpu_peak_mem_orig=}MB {gpu_alloc_mem_orig+gpu_peak_mem_orig=}MB")
          # print(f" {gpu_alloc_mem_bnb=}MB  {gpu_peak_mem_bnb=}MB  {gpu_alloc_mem_bnb+gpu_peak_mem_bnb=}MB")
          # print(f"{gpu_alloc_mem_diff=}MB")
          # print(f"{gpu_peak_mem_diff=}MB")
          # print(f"{gpu_total_mem_orig=}MB, {gpu_total_mem_bnb=}MB")
          # print(f"{gpu_total_mem_diff=}MB, {gpu_total_mem_diff=}MB")

          self.assertGreater(
              _A		,		_A		,		'''should use ~150MB less alloc gpu memory with BNB, compared to without it for this model but got'''
              f""" a difference of {gpu_alloc_mem_diff}MB, with gpu_alloc_mem_orig={gpu_alloc_mem_orig}MB and"""
              f""" gpu_alloc_mem_bnb={gpu_alloc_mem_bnb}MB"""		,		)

          self.assertGreater(
              _A		,		_A		,		'''should use ~150MB less total gpu memory with BNB, compared to without it for this model but got'''
              f""" a difference of {gpu_total_mem_diff}MB, with gpu_total_mem_orig={gpu_total_mem_orig}MB and"""
              f""" gpu_total_mem_bnb={gpu_total_mem_bnb}MB"""		,		)

          self.assertEqual(
              _A		,		_A		,		f"""loss should be the same, but got loss_orig={loss_orig}, loss_bnb={loss_bnb}"""					)






      def lowercase_				(   self							:       Optional[int]		,		_A							:       int		,		_A							:       str		,		_A							:       int		,		_A							:       float = 3e-3		,		_A							:       str = "adafactor"		,		_A							:       bool = False		,		_A							:       str = None		,		_A							:       int = 0		,		_A							:       bool = True		,		_A							:       bool = True		,		_A							:       bool = True		,		_A							:       bool = True		,		_A							:       int = None		,		):


          '''simple docstring'''


          UpperCAmelCase__     :						str						=							self.test_file_dir / '''../fixtures/tests_samples/wmt_en_ro'''
          UpperCAmelCase__     :						List[Any]						=							self.get_auto_remove_tmp_dir()
          UpperCAmelCase__     :						Optional[int]						=							f"""
            --model_name_or_path {model_name}
            --train_file {data_dir}/train.json
            --validation_file {data_dir}/val.json
            --test_file {data_dir}/test.json
            --output_dir {output_dir}
            --overwrite_output_dir
            --max_train_samples 8
            --max_source_length {max_len}
            --max_target_length {max_len}
            --do_train
            --num_train_epochs {str(_A					)}
            --per_device_train_batch_size 4
            --learning_rate {learning_rate}
            --warmup_steps 8
            --logging_steps 0
            --logging_strategy no
            --save_steps {str(_A					)}
            --group_by_length
            --label_smoothing_factor 0.1
            --target_lang ro_RO
            --source_lang en_XX
        """.split()

          UpperCAmelCase__     :						List[Any]						=							f"""
            --do_eval
            --per_device_eval_batch_size 4
            --max_eval_samples 8
            --val_max_target_length {max_len}
            --evaluation_strategy steps
            --eval_steps {str(_A					)}
        """.split()

          UpperCAmelCase__     :						Dict						=							'''
            --do_predict
        '''.split()

          UpperCAmelCase__     :						Dict						=							[]
          if do_train:
              args += args_train

          if do_eval:
              args += args_eval

          if do_predict:
              args += args_predict

          if predict_with_generate:
              args += "--predict_with_generate".split()

          if do_train:
              if optim == "adafactor":
                  args += "--adafactor".split()
              else:
                  args += f"""--optim {optim}""".split()

          if extra_args_str is not None:
              args += extra_args_str.split()

          if distributed:
              if n_gpus_to_use is None:
                  UpperCAmelCase__     :						Union[str, Any]						=							get_gpu_count()
              UpperCAmelCase__     :						Any						=							get_torch_dist_unique_port()
              UpperCAmelCase__     :						Tuple						=							f"""
                -m torch.distributed.run
                --nproc_per_node={n_gpus_to_use}
                --master_port={master_port}
                {self.examples_dir_str}/pytorch/translation/run_translation.py
            """.split()
              UpperCAmelCase__     :						List[Any]						=							[sys.executable] + distributed_args + args
              # keep for quick debug
              # print(" ".join([f"\nPYTHONPATH={self.src_dir_str}"] +cmd)); die
              execute_subprocess_async(_A		,		env=self.get_env()					)
          else:
              UpperCAmelCase__     :						Optional[Any]						=							['''run_translation.py'''] + args
              with patch.object(_A		,		'''argv'''		,		_A					):
                  main()

          return output_dir





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

			def __init__(self,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_=13,   SCREAMING_SNAKE_CASE_=7,   SCREAMING_SNAKE_CASE_=True,   SCREAMING_SNAKE_CASE_=True,   SCREAMING_SNAKE_CASE_=True,   SCREAMING_SNAKE_CASE_=True,   SCREAMING_SNAKE_CASE_=99,   SCREAMING_SNAKE_CASE_=16,   SCREAMING_SNAKE_CASE_=36,   SCREAMING_SNAKE_CASE_=6,   SCREAMING_SNAKE_CASE_=6,   SCREAMING_SNAKE_CASE_=6,   SCREAMING_SNAKE_CASE_=37,   SCREAMING_SNAKE_CASE_="gelu",   SCREAMING_SNAKE_CASE_=0.1,   SCREAMING_SNAKE_CASE_=0.1,   SCREAMING_SNAKE_CASE_=512,   SCREAMING_SNAKE_CASE_=16,   SCREAMING_SNAKE_CASE_=2,   SCREAMING_SNAKE_CASE_=0.0_2,   SCREAMING_SNAKE_CASE_=3,   SCREAMING_SNAKE_CASE_=4,   SCREAMING_SNAKE_CASE_=None,   )       ->   Tuple:
						UpperCAmelCase_:		Any							      =     parent
						UpperCAmelCase_:		Union[str, Any]							      =     batch_size
						UpperCAmelCase_:		int							      =     seq_length
						UpperCAmelCase_:		Any							      =     is_training
						UpperCAmelCase_:		str							      =     use_input_mask
						UpperCAmelCase_:		Dict							      =     use_token_type_ids
						UpperCAmelCase_:		Any							      =     use_labels
						UpperCAmelCase_:		Any							      =     vocab_size
						UpperCAmelCase_:		Dict							      =     embedding_size
						UpperCAmelCase_:		Dict							      =     hidden_size
						UpperCAmelCase_:		Dict							      =     num_hidden_layers
						UpperCAmelCase_:		Dict							      =     num_hidden_groups
						UpperCAmelCase_:		int							      =     num_attention_heads
						UpperCAmelCase_:		List[str]							      =     intermediate_size
						UpperCAmelCase_:		Any							      =     hidden_act
						UpperCAmelCase_:		Dict							      =     hidden_dropout_prob
						UpperCAmelCase_:		Any							      =     attention_probs_dropout_prob
						UpperCAmelCase_:		Tuple							      =     max_position_embeddings
						UpperCAmelCase_:		Tuple							      =     type_vocab_size
						UpperCAmelCase_:		List[Any]							      =     type_sequence_label_size
						UpperCAmelCase_:		Any							      =     initializer_range
						UpperCAmelCase_:		int							      =     num_labels
						UpperCAmelCase_:		Optional[int]							      =     num_choices
						UpperCAmelCase_:		Tuple							      =     scope

			def      __snake_case       (self					)       ->   List[Any]:
						UpperCAmelCase_:		Union[str, Any]							      =     ids_tensor([self.batch_size, self.seq_length],   self.vocab_size					)

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

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

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

						UpperCAmelCase_:		Optional[int]							      =     self.get_config()

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

			def      __snake_case       (self					)       ->   Optional[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      __snake_case       (self,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_					)       ->   Optional[Any]:
						UpperCAmelCase_:		int							      =     AlbertModel(config=SCREAMING_SNAKE_CASE_					)
						model.to(SCREAMING_SNAKE_CASE_					)
						model.eval()
						UpperCAmelCase_:		Optional[int]							      =     model(SCREAMING_SNAKE_CASE_,   attention_mask=SCREAMING_SNAKE_CASE_,   token_type_ids=SCREAMING_SNAKE_CASE_					)
						UpperCAmelCase_:		List[Any]							      =     model(SCREAMING_SNAKE_CASE_,   token_type_ids=SCREAMING_SNAKE_CASE_					)
						UpperCAmelCase_:		Union[str, Any]							      =     model(SCREAMING_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      __snake_case       (self,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_					)       ->   str:
						UpperCAmelCase_:		Any							      =     AlbertForPreTraining(config=SCREAMING_SNAKE_CASE_					)
						model.to(SCREAMING_SNAKE_CASE_					)
						model.eval()
						UpperCAmelCase_:		List[Any]							      =     model(
						    SCREAMING_SNAKE_CASE_,   attention_mask=SCREAMING_SNAKE_CASE_,   token_type_ids=SCREAMING_SNAKE_CASE_,   labels=SCREAMING_SNAKE_CASE_,   sentence_order_label=SCREAMING_SNAKE_CASE_,   )
						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      __snake_case       (self,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_					)       ->   List[str]:
						UpperCAmelCase_:		Optional[int]							      =     AlbertForMaskedLM(config=SCREAMING_SNAKE_CASE_					)
						model.to(SCREAMING_SNAKE_CASE_					)
						model.eval()
						UpperCAmelCase_:		Optional[Any]							      =     model(SCREAMING_SNAKE_CASE_,   attention_mask=SCREAMING_SNAKE_CASE_,   token_type_ids=SCREAMING_SNAKE_CASE_,   labels=SCREAMING_SNAKE_CASE_					)
						self.parent.assertEqual(result.logits.shape,   (self.batch_size, self.seq_length, self.vocab_size)					)

			def      __snake_case       (self,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_					)       ->   Tuple:
						UpperCAmelCase_:		Optional[Any]							      =     AlbertForQuestionAnswering(config=SCREAMING_SNAKE_CASE_					)
						model.to(SCREAMING_SNAKE_CASE_					)
						model.eval()
						UpperCAmelCase_:		Any							      =     model(
						    SCREAMING_SNAKE_CASE_,   attention_mask=SCREAMING_SNAKE_CASE_,   token_type_ids=SCREAMING_SNAKE_CASE_,   start_positions=SCREAMING_SNAKE_CASE_,   end_positions=SCREAMING_SNAKE_CASE_,   )
						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      __snake_case       (self,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_					)       ->   Any:
						UpperCAmelCase_:		Optional[Any]							      =     self.num_labels
						UpperCAmelCase_:		List[Any]							      =     AlbertForSequenceClassification(SCREAMING_SNAKE_CASE_					)
						model.to(SCREAMING_SNAKE_CASE_					)
						model.eval()
						UpperCAmelCase_:		str							      =     model(SCREAMING_SNAKE_CASE_,   attention_mask=SCREAMING_SNAKE_CASE_,   token_type_ids=SCREAMING_SNAKE_CASE_,   labels=SCREAMING_SNAKE_CASE_					)
						self.parent.assertEqual(result.logits.shape,   (self.batch_size, self.num_labels)					)

			def      __snake_case       (self,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_					)       ->   List[Any]:
						UpperCAmelCase_:		List[str]							      =     self.num_labels
						UpperCAmelCase_:		Union[str, Any]							      =     AlbertForTokenClassification(config=SCREAMING_SNAKE_CASE_					)
						model.to(SCREAMING_SNAKE_CASE_					)
						model.eval()
						UpperCAmelCase_:		List[str]							      =     model(SCREAMING_SNAKE_CASE_,   attention_mask=SCREAMING_SNAKE_CASE_,   token_type_ids=SCREAMING_SNAKE_CASE_,   labels=SCREAMING_SNAKE_CASE_					)
						self.parent.assertEqual(result.logits.shape,   (self.batch_size, self.seq_length, self.num_labels)					)

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

			def      __snake_case       (self					)       ->   Union[str, Any]:
						UpperCAmelCase_:		Optional[Any]							      =     self.prepare_config_and_inputs()
						(
						    (
						    UpperCAmelCase_
						)      ,							(
						    UpperCAmelCase_
						)      ,							(
						    UpperCAmelCase_
						)      ,							(
						    UpperCAmelCase_
						)      ,							(
						    UpperCAmelCase_
						)      ,							(
						    UpperCAmelCase_
						)      ,							(
						    UpperCAmelCase_
						)      ,							
						):		List[str]							      =     config_and_inputs
						UpperCAmelCase_:		Tuple							      =     {"""input_ids""": input_ids, """token_type_ids""": token_type_ids, """attention_mask""": input_mask}
						return config, inputs_dict







@require_torch
class 				_a				(   _lowerCAmelCase ,		_lowerCAmelCase ,		unittest.TestCase					):
			A      		=	(
			    (
			        AlbertModel,
			        AlbertForPreTraining,
			        AlbertForMaskedLM,
			        AlbertForMultipleChoice,
			        AlbertForSequenceClassification,
			        AlbertForTokenClassification,
			        AlbertForQuestionAnswering,
			    )
			    if is_torch_available()
			    else ()
			)
			A      		=	(
			    {
			        '''feature-extraction''': AlbertModel,
			        '''fill-mask''': AlbertForMaskedLM,
			        '''question-answering''': AlbertForQuestionAnswering,
			        '''text-classification''': AlbertForSequenceClassification,
			        '''token-classification''': AlbertForTokenClassification,
			        '''zero-shot''': AlbertForSequenceClassification,
			    }
			    if is_torch_available()
			    else {}
			)
			A      		=	True

			def      __snake_case       (self,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_=False					)       ->   Any:
						UpperCAmelCase_:		Tuple							      =     super()._prepare_for_class(SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   return_labels=SCREAMING_SNAKE_CASE_					)

						if return_labels:
									if model_class in get_values(SCREAMING_SNAKE_CASE_					):
												UpperCAmelCase_:		Dict							      =     torch.zeros(
												    (self.model_tester.batch_size, self.model_tester.seq_length),   dtype=torch.long,   device=SCREAMING_SNAKE_CASE_					)
												UpperCAmelCase_:		Optional[int]							      =     torch.zeros(
												    self.model_tester.batch_size,   dtype=torch.long,   device=SCREAMING_SNAKE_CASE_					)
						return inputs_dict

			def      __snake_case       (self					)       ->   Optional[Any]:
						UpperCAmelCase_:		Any							      =     AlbertModelTester(self					)
						UpperCAmelCase_:		Union[str, Any]							      =     ConfigTester(self,   config_class=SCREAMING_SNAKE_CASE_,   hidden_size=37					)

			def      __snake_case       (self					)       ->   List[Any]:
						self.config_tester.run_common_tests()

			def      __snake_case       (self					)       ->   List[Any]:
						UpperCAmelCase_:		List[Any]							      =     self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_model(*SCREAMING_SNAKE_CASE_					)

			def      __snake_case       (self					)       ->   Any:
						UpperCAmelCase_:		int							      =     self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_pretraining(*SCREAMING_SNAKE_CASE_					)

			def      __snake_case       (self					)       ->   Optional[Any]:
						UpperCAmelCase_:		Tuple							      =     self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_masked_lm(*SCREAMING_SNAKE_CASE_					)

			def      __snake_case       (self					)       ->   str:
						UpperCAmelCase_:		Optional[Any]							      =     self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_multiple_choice(*SCREAMING_SNAKE_CASE_					)

			def      __snake_case       (self					)       ->   Any:
						UpperCAmelCase_:		List[str]							      =     self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_question_answering(*SCREAMING_SNAKE_CASE_					)

			def      __snake_case       (self					)       ->   str:
						UpperCAmelCase_:		Optional[int]							      =     self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_sequence_classification(*SCREAMING_SNAKE_CASE_					)

			def      __snake_case       (self					)       ->   Dict:
						UpperCAmelCase_:		Optional[Any]							      =     self.model_tester.prepare_config_and_inputs()
						for type in ["absolute", "relative_key", "relative_key_query"]:
									UpperCAmelCase_:		Union[str, Any]							      =     type
									self.model_tester.create_and_check_model(*SCREAMING_SNAKE_CASE_					)

			@slow
			def      __snake_case       (self					)       ->   List[Any]:
						for model_name in ALBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
									UpperCAmelCase_:		Dict							      =     AlbertModel.from_pretrained(SCREAMING_SNAKE_CASE_					)
									self.assertIsNotNone(SCREAMING_SNAKE_CASE_					)







@require_torch
class 				_a				(   unittest.TestCase					):

			@slow
			def      __snake_case       (self					)       ->   Dict:
						UpperCAmelCase_:		str							      =     AlbertModel.from_pretrained("""albert-base-v2"""					)
						UpperCAmelCase_:		List[str]							      =     torch.tensor([[0, 345, 232, 328, 740, 140, 1695, 69, 6078, 1588, 2]]					)
						UpperCAmelCase_:		List[str]							      =     torch.tensor([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]					)
						with torch.no_grad():
									UpperCAmelCase_:		Union[str, Any]							      =     model(SCREAMING_SNAKE_CASE_,   attention_mask=SCREAMING_SNAKE_CASE_					)[0]
						UpperCAmelCase_:		List[str]							      =     torch.Size((1, 11, 768)					)
						self.assertEqual(output.shape,   SCREAMING_SNAKE_CASE_					)
						UpperCAmelCase_:		Dict							      =     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],   SCREAMING_SNAKE_CASE_,   atol=1E-4					)					)






