


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


a			:			Dict												=					os.path.abspath(os.path.dirname(__file__))
with ExtendSysPath(f'''{bindir}/../../examples/pytorch/translation'''):
	from run_translation import main  # noqa


set_seed(42)
a			:			Dict												=					'''sshleifer/student_marian_en_ro_6_1'''
a			:			str												=					'''sshleifer/tiny-mbart'''
@require_torch
class 						lowerCamelCase_							(					lowerCAmelCase__					):



					'''simple docstring'''







					def A   (					self			,							snake_case_=False			,							snake_case_=None			,							snake_case_=True			,							snake_case_=True			,							snake_case_=True			,							snake_case_=True			,							) ->      Optional[int]:







								'''simple docstring'''

								__lowercase			       = self.run_trainer(
								    eval_steps=1			,							max_len=1_2			,							model_name=snake_case_			,							num_train_epochs=1			,							distributed=snake_case_			,							extra_args_str=snake_case_			,							predict_with_generate=snake_case_			,							do_train=snake_case_			,							do_eval=snake_case_			,							do_predict=snake_case_			,							)
								__lowercase			       = TrainerState.load_from_json(os.path.join(snake_case_			,							'''trainer_state.json'''    )    ).log_history

								if not do_eval:
											return

								__lowercase			       = [log for log in logs if '''eval_loss''' in log.keys()]

								__lowercase			       = eval_metrics[0]
								if predict_with_generate:
											assert "eval_bleu" in first_step_stats

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






					@require_torch_non_multi_gpu
					def A   (					self    ) ->      List[str]:







								'''simple docstring'''

								self.run_seqaseq_quick()






					@require_torch_multi_gpu
					def A   (					self    ) ->      Dict:







								'''simple docstring'''

								self.run_seqaseq_quick(distributed=snake_case_    )






					@require_torch_multi_gpu
					def A   (					self    ) ->      int:







								'''simple docstring'''

								self.run_seqaseq_quick(distributed=snake_case_    )






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







								'''simple docstring'''

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






					@unittest.skip('''Requires an update of the env running those tests'''    )
					@require_torch_multi_gpu
					@require_fairscale
					def A   (					self    ) ->      str:







								'''simple docstring'''

								self.run_seqaseq_quick(distributed=snake_case_			,							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 A   (					self    ) ->      Union[str, Any]:







								'''simple docstring'''

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






					@unittest.skip('''Requires an update of the env running those tests'''    )
					@require_torch_multi_gpu
					@require_fairscale
					def A   (					self    ) ->      Dict:







								'''simple docstring'''

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






					@require_apex
					@require_torch_gpu
					def A   (					self    ) ->      Any:







								'''simple docstring'''

								self.run_seqaseq_quick(distributed=snake_case_			,							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=snake_case_			,							extra_args_str='''--fp16 --fp16_backend=apex'''    )






					@parameterized.expand(['''base''', '''low''', '''high''', '''mixed''']    )
					@require_torch_multi_gpu
					def A   (					self			,							snake_case_    ) ->      List[str]:







								'''simple docstring'''

								__lowercase			       = {
								    # 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},
								}

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






					@slow
					def A   (					self    ) ->      Optional[Any]:







								'''simple docstring'''

								__lowercase			       = self.run_trainer(
								    eval_steps=2			,							max_len=1_2_8			,							model_name=snake_case_			,							learning_rate=3e-4			,							num_train_epochs=1_0			,							distributed=snake_case_			,							)

								# Check metrics
								__lowercase			       = TrainerState.load_from_json(os.path.join(snake_case_			,							'''trainer_state.json'''    )    ).log_history
								__lowercase			       = [log for log in logs if '''eval_loss''' in log.keys()]
								__lowercase			       = eval_metrics[0]
								__lowercase			       = eval_metrics[-1]

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

								# test if do_predict saves generations and metrics
								__lowercase			       = os.listdir(snake_case_    )
								__lowercase			       = {os.path.basename(snake_case_    ) for p in contents}
								assert "generated_predictions.txt" in contents
								assert "predict_results.json" in contents






					@slow
					@require_bitsandbytes
					def A   (					self    ) ->      Optional[int]:







								'''simple docstring'''

								from transformers.training_args import OptimizerNames

								def train_and_return_metrics(snake_case_    ) -> Tuple[int, float]:
											__lowercase			       = '''--skip_memory_metrics 0'''

											__lowercase			       = self.run_trainer(
											    max_len=1_2_8			,							model_name=snake_case_			,							learning_rate=3e-4			,							num_train_epochs=1			,							optim=snake_case_			,							distributed=snake_case_			,							extra_args_str=snake_case_			,							do_eval=snake_case_			,							do_predict=snake_case_			,							n_gpus_to_use=1			,							)

											# Check metrics
											__lowercase			       = TrainerState.load_from_json(Path(snake_case_			,							'''trainer_state.json'''    )    ).log_history
											__lowercase			       = int(logs[0]['''train_mem_gpu_peaked_delta'''] / 2**2_0    )
											__lowercase			       = int(logs[0]['''train_mem_gpu_alloc_delta'''] / 2**2_0    )

											__lowercase			       = logs[0]['''train_loss''']
											return gpu_peak_mem_mb, gpu_alloc_mem_mb, loss

								__lowercase       ,			__lowercase       ,			__lowercase			       = train_and_return_metrics(OptimizerNames.ADAMW_TORCH.value    )
								__lowercase       ,			__lowercase       ,			__lowercase			       = train_and_return_metrics(OptimizerNames.ADAMW_BNB.value    )

								__lowercase			       = gpu_alloc_mem_orig - gpu_alloc_mem_bnb

								__lowercase			       = gpu_peak_mem_orig + gpu_alloc_mem_orig
								__lowercase			       = gpu_peak_mem_bnb + gpu_alloc_mem_bnb
								__lowercase			       = 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
								__lowercase			       = 1_2_0

								# 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(
								    snake_case_			,							snake_case_			,							'''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(
								    snake_case_			,							snake_case_			,							'''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(
								    snake_case_			,							snake_case_			,							F'loss should be the same, but got loss_orig={loss_orig}, loss_bnb={loss_bnb}'    )






					def A   (					self			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_ = 3e-3			,							snake_case_ = "adafactor"			,							snake_case_ = False			,							snake_case_ = None			,							snake_case_ = 0			,							snake_case_ = True			,							snake_case_ = True			,							snake_case_ = True			,							snake_case_ = True			,							snake_case_ = None			,							) ->      Tuple:







								'''simple docstring'''

								__lowercase			       = self.test_file_dir / '''../fixtures/tests_samples/wmt_en_ro'''
								__lowercase			       = self.get_auto_remove_tmp_dir()
								__lowercase			       = F'\n            --model_name_or_path {model_name}\n            --train_file {data_dir}/train.json\n            --validation_file {data_dir}/val.json\n            --test_file {data_dir}/test.json\n            --output_dir {output_dir}\n            --overwrite_output_dir\n            --max_train_samples 8\n            --max_source_length {max_len}\n            --max_target_length {max_len}\n            --do_train\n            --num_train_epochs {str(snake_case_    )}\n            --per_device_train_batch_size 4\n            --learning_rate {learning_rate}\n            --warmup_steps 8\n            --logging_steps 0\n            --logging_strategy no\n            --save_steps {str(snake_case_    )}\n            --group_by_length\n            --label_smoothing_factor 0.1\n            --target_lang ro_RO\n            --source_lang en_XX\n        '.split()

								__lowercase			       = F'\n            --do_eval\n            --per_device_eval_batch_size 4\n            --max_eval_samples 8\n            --val_max_target_length {max_len}\n            --evaluation_strategy steps\n            --eval_steps {str(snake_case_    )}\n        '.split()

								__lowercase			       = '''
            --do_predict
        '''.split()

								__lowercase			       = []
								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:
														__lowercase			       = get_gpu_count()
											__lowercase			       = get_torch_dist_unique_port()
											__lowercase			       = F'\n                -m torch.distributed.run\n                --nproc_per_node={n_gpus_to_use}\n                --master_port={master_port}\n                {self.examples_dir_str}/pytorch/translation/run_translation.py\n            '.split()
											__lowercase			       = [sys.executable] + distributed_args + args
											# keep for quick debug
											# print(" ".join([f"\nPYTHONPATH={self.src_dir_str}"] +cmd)); die
											execute_subprocess_async(snake_case_			,							env=self.get_env()    )
								else:
											__lowercase			       = ['''run_translation.py'''] + args
											with patch.object(snake_case_			,							'''argv'''			,							snake_case_    ):
														main()

								return output_dir






import math
import os
import unittest

from transformers import MegatronBertConfig, is_torch_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_sentencepiece, require_tokenizers, 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,
	    MegatronBertForCausalLM,
	    MegatronBertForMaskedLM,
	    MegatronBertForMultipleChoice,
	    MegatronBertForNextSentencePrediction,
	    MegatronBertForPreTraining,
	    MegatronBertForQuestionAnswering,
	    MegatronBertForSequenceClassification,
	    MegatronBertForTokenClassification,
	    MegatronBertModel,
	)
class 						lowerCamelCase_							:



					'''simple docstring'''







					def __init__(					self			,							snake_case_			,							snake_case_=1_3			,							snake_case_=7			,							snake_case_=True			,							snake_case_=True			,							snake_case_=True			,							snake_case_=True			,							snake_case_=9_9			,							snake_case_=6_4			,							snake_case_=3_2			,							snake_case_=5			,							snake_case_=4			,							snake_case_=3_7			,							snake_case_="gelu"			,							snake_case_=0.1			,							snake_case_=0.1			,							snake_case_=5_1_2			,							snake_case_=1_6			,							snake_case_=2			,							snake_case_=0.0_2			,							snake_case_=3			,							snake_case_=4			,							snake_case_=None			,							) ->      Tuple:







								'''simple docstring'''

								__lowercase			       = parent
								__lowercase			       = batch_size
								__lowercase			       = seq_length
								__lowercase			       = is_training
								__lowercase			       = use_input_mask
								__lowercase			       = use_token_type_ids
								__lowercase			       = use_labels
								__lowercase			       = vocab_size
								__lowercase			       = hidden_size
								__lowercase			       = embedding_size
								__lowercase			       = num_hidden_layers
								__lowercase			       = num_attention_heads
								__lowercase			       = intermediate_size
								__lowercase			       = hidden_act
								__lowercase			       = hidden_dropout_prob
								__lowercase			       = attention_probs_dropout_prob
								__lowercase			       = max_position_embeddings
								__lowercase			       = type_vocab_size
								__lowercase			       = type_sequence_label_size
								__lowercase			       = initializer_range
								__lowercase			       = num_labels
								__lowercase			       = num_choices
								__lowercase			       = scope






					def A   (					self    ) ->      List[str]:







								'''simple docstring'''

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

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

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

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

								__lowercase			       = self.get_config()

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






					def A   (					self    ) ->      List[Any]:







								'''simple docstring'''

								return MegatronBertConfig(
								    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			,							embedding_size=self.embedding_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			,							is_decoder=snake_case_			,							initializer_range=self.initializer_range			,							)






					def A   (					self			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_    ) ->      List[str]:







								'''simple docstring'''

								__lowercase			       = MegatronBertModel(config=snake_case_    )
								model.to(snake_case_    )
								model.eval()
								__lowercase			       = model(snake_case_			,							attention_mask=snake_case_			,							token_type_ids=snake_case_    )
								__lowercase			       = model(snake_case_			,							token_type_ids=snake_case_    )
								__lowercase			       = model(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 A   (					self			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_    ) ->      List[str]:







								'''simple docstring'''

								__lowercase			       = MegatronBertForMaskedLM(config=snake_case_    )
								model.to(snake_case_    )
								model.eval()
								__lowercase			       = model(snake_case_			,							attention_mask=snake_case_			,							token_type_ids=snake_case_			,							labels=snake_case_    )
								self.parent.assertEqual(result.logits.shape			,							(self.batch_size, self.seq_length, self.vocab_size)    )






					def A   (					self			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_    ) ->      int:







								'''simple docstring'''

								__lowercase			       = MegatronBertForCausalLM(config=snake_case_    )
								model.to(snake_case_    )
								model.eval()
								__lowercase			       = model(snake_case_			,							attention_mask=snake_case_			,							token_type_ids=snake_case_			,							labels=snake_case_    )
								self.parent.assertEqual(result.logits.shape			,							(self.batch_size, self.seq_length, self.vocab_size)    )






					def A   (					self			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_    ) ->      Tuple:







								'''simple docstring'''

								__lowercase			       = MegatronBertForNextSentencePrediction(config=snake_case_    )
								model.to(snake_case_    )
								model.eval()
								__lowercase			       = model(
								    snake_case_			,							attention_mask=snake_case_			,							token_type_ids=snake_case_			,							labels=snake_case_			,							)
								self.parent.assertEqual(result.logits.shape			,							(self.batch_size, 2)    )






					def A   (					self			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_    ) ->      Optional[int]:







								'''simple docstring'''

								__lowercase			       = MegatronBertForPreTraining(config=snake_case_    )
								model.to(snake_case_    )
								model.eval()
								__lowercase			       = model(
								    snake_case_			,							attention_mask=snake_case_			,							token_type_ids=snake_case_			,							labels=snake_case_			,							next_sentence_label=snake_case_			,							)
								self.parent.assertEqual(result.prediction_logits.shape			,							(self.batch_size, self.seq_length, self.vocab_size)    )
								self.parent.assertEqual(result.seq_relationship_logits.shape			,							(self.batch_size, 2)    )






					def A   (					self			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_    ) ->      Any:







								'''simple docstring'''

								__lowercase			       = MegatronBertForQuestionAnswering(config=snake_case_    )
								model.to(snake_case_    )
								model.eval()
								__lowercase			       = model(
								    snake_case_			,							attention_mask=snake_case_			,							token_type_ids=snake_case_			,							start_positions=snake_case_			,							end_positions=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 A   (					self			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_    ) ->      Optional[int]:







								'''simple docstring'''

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






					def A   (					self			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_    ) ->      Dict:







								'''simple docstring'''

								__lowercase			       = self.num_labels
								__lowercase			       = MegatronBertForTokenClassification(config=snake_case_    )
								model.to(snake_case_    )
								model.eval()
								__lowercase			       = model(snake_case_			,							attention_mask=snake_case_			,							token_type_ids=snake_case_			,							labels=snake_case_    )
								self.parent.assertEqual(result.logits.shape			,							(self.batch_size, self.seq_length, self.num_labels)    )






					def A   (					self			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_    ) ->      Dict:







								'''simple docstring'''

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






					def A   (					self    ) ->      Union[str, Any]:







								'''simple docstring'''

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




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



					'''simple docstring'''







					__UpperCAmelCase							  =					(
					    (
					        MegatronBertModel,
					        MegatronBertForMaskedLM,
					        MegatronBertForCausalLM,
					        MegatronBertForMultipleChoice,
					        MegatronBertForNextSentencePrediction,
					        MegatronBertForPreTraining,
					        MegatronBertForQuestionAnswering,
					        MegatronBertForSequenceClassification,
					        MegatronBertForTokenClassification,
					    )
					    if is_torch_available()
					    else ()
					)
					__UpperCAmelCase							  =					(
					    {
					        "feature-extraction": MegatronBertModel,
					        "fill-mask": MegatronBertForMaskedLM,
					        "question-answering": MegatronBertForQuestionAnswering,
					        "text-classification": MegatronBertForSequenceClassification,
					        "text-generation": MegatronBertForCausalLM,
					        "token-classification": MegatronBertForTokenClassification,
					        "zero-shot": MegatronBertForSequenceClassification,
					    }
					    if is_torch_available()
					    else {}
					)
					__UpperCAmelCase							  =					True
					# test_resize_embeddings = False
					__UpperCAmelCase							  =					False






					def A   (					self			,							snake_case_			,							snake_case_			,							snake_case_=False    ) ->      Tuple:







								'''simple docstring'''

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

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






					def A   (					self    ) ->      List[Any]:







								'''simple docstring'''

								__lowercase			       = MegatronBertModelTester(self    )
								__lowercase			       = ConfigTester(self			,							config_class=snake_case_			,							hidden_size=3_7    )






					def A   (					self    ) ->      int:







								'''simple docstring'''

								self.config_tester.run_common_tests()






					def A   (					self    ) ->      List[str]:







								'''simple docstring'''

								__lowercase			       = self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_megatron_bert_model(*snake_case_    )






					def A   (					self    ) ->      Optional[int]:







								'''simple docstring'''

								__lowercase			       = self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_megatron_bert_for_masked_lm(*snake_case_    )






					def A   (					self    ) ->      str:







								'''simple docstring'''

								__lowercase			       = self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_megatron_bert_for_multiple_choice(*snake_case_    )






					def A   (					self    ) ->      List[Any]:







								'''simple docstring'''

								__lowercase			       = self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_megatron_bert_for_next_sequence_prediction(*snake_case_    )






					def A   (					self    ) ->      Dict:







								'''simple docstring'''

								__lowercase			       = self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_megatron_bert_for_pretraining(*snake_case_    )






					def A   (					self    ) ->      Optional[Any]:







								'''simple docstring'''

								__lowercase			       = self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_megatron_bert_for_question_answering(*snake_case_    )






					def A   (					self    ) ->      List[Any]:







								'''simple docstring'''

								__lowercase			       = self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_megatron_bert_for_sequence_classification(*snake_case_    )






					def A   (					self    ) ->      Dict:







								'''simple docstring'''

								__lowercase			       = self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_megatron_bert_for_token_classification(*snake_case_    )




def 		lowercase_	(  _UpperCamelCase	):


			'''simple docstring'''







			return torch.tensor(
			    _UpperCamelCase							,      dtype=torch.long							,      device=_UpperCamelCase							,      )


a			:			Optional[Any]												=					1E-4





@require_torch
@require_sentencepiece
@require_tokenizers
class 						lowerCamelCase_							(					unittest.TestCase					):



					'''simple docstring'''







					@slow
					@unittest.skip('''Model is not available.'''    )
					def A   (					self    ) ->      Tuple:







								'''simple docstring'''

								__lowercase			       = '''nvidia/megatron-bert-uncased-345m'''
								if "MYDIR" in os.environ:
											__lowercase			       = os.path.join(os.environ['''MYDIR''']			,							snake_case_    )
								__lowercase			       = MegatronBertModel.from_pretrained(snake_case_    )
								model.to(snake_case_    )
								model.half()
								__lowercase			       = _long_tensor([[1_0_1, 7_1_1_0, 1_0_0_5, 1_0_5_6, 2_0_2_3, 1_1_3_3_3, 1_7_4_1_3, 1_0_2_9, 1_0_2]]    )
								with torch.no_grad():
											__lowercase			       = model(snake_case_    )[0]
								__lowercase			       = torch.Size((1, 9, 1_0_2_4)    )
								self.assertEqual(output.shape			,							snake_case_    )

								__lowercase			       = [-0.6_0_4_0, -0.2_5_1_7, -0.1_0_2_5, 0.3_4_2_0, -0.6_7_5_8, -0.0_0_1_7, -0.1_0_8_9, -0.1_9_9_0, 0.5_7_2_8]
								for ii in range(3    ):
											for jj in range(3    ):
														__lowercase			       = output[0, ii, jj]
														__lowercase			       = expected[3 * ii + jj]
														__lowercase			       = '''ii={} jj={} a={} b={}'''.format(snake_case_			,							snake_case_			,							snake_case_			,							snake_case_    )
														self.assertTrue(math.isclose(snake_case_			,							snake_case_			,							rel_tol=snake_case_			,							abs_tol=snake_case_    )			,							msg=snake_case_    )


