

'''simple docstring'''





import argparse
import os

import evaluate
import torch
from datasets import load_dataset
from torch.optim import AdamW
from torch.utils.data import DataLoader
from transformers import AutoModelForSequenceClassification, AutoTokenizer, get_linear_schedule_with_warmup, set_seed

from accelerate import Accelerator, DistributedType


########################################################################
# This is a fully working simple example to use Accelerate,
# specifically showcasing the experiment tracking capability,
# and builds off the `nlp_example.py` script.
#
# This example trains a Bert base model on GLUE MRPC
# in any of the following settings (with the same script):
#   - single CPU or single GPU
#   - multi GPUS (using PyTorch distributed mode)
#   - (multi) TPUs
#   - fp16 (mixed-precision) or fp32 (normal precision)
#
# To help focus on the differences in the code, building `DataLoaders`
# was refactored into its own function.
# New additions from the base script can be found quickly by
# looking for the # New Code # tags
#
# To run it in each of these various modes, follow the instructions
# in the readme for examples:
# https://github.com/huggingface/accelerate/tree/main/examples
#
########################################################################

lowerCAmelCase__								=		1_6
lowerCAmelCase__								=		3_2






def 					lowerCamelCase_     (							UpperCAmelCase_							:     Accelerator							,       UpperCAmelCase_							:     int = 1_6	)						->				Tuple:




	'''simple docstring'''




	_UpperCamelCase   :		Tuple			=						AutoTokenizer.from_pretrained('bert-base-cased'	)
	_UpperCamelCase   :		Union[str, Any]			=						load_dataset('glue'							,       'mrpc'	)

	def tokenize_function(UpperCAmelCase_							:     Union[str, Any]	):
		# max_length=None => use the model max length (it's actually the default)
		_UpperCamelCase   :		Optional[int]			=						tokenizer(examples['sentence1']							,       examples['sentence2']							,       truncation=UpperCAmelCase_							,       max_length=UpperCAmelCase_	)
		return outputs

	# Apply the method we just defined to all the examples in all the splits of the dataset
	# starting with the main process first:
	with accelerator.main_process_first():
		_UpperCamelCase   :		Optional[Any]			=						datasets.map(
		    UpperCAmelCase_							,       batched=UpperCAmelCase_							,       remove_columns=['idx', 'sentence1', 'sentence2']							,       )

	# We also rename the 'label' column to 'labels' which is the expected name for labels by the models of the
	# transformers library
	_UpperCamelCase   :		Optional[int]			=						tokenized_datasets.rename_column('label'							,       'labels'	)

	def collate_fn(UpperCAmelCase_							:     Any	):
		# On TPU it's best to pad everything to the same length or training will be very slow.
		_UpperCamelCase   :		int			=						1_2_8 if accelerator.distributed_type == DistributedType.TPU else None
		# When using mixed precision we want round multiples of 8/16
		if accelerator.mixed_precision == "fp8":
			_UpperCamelCase   :		List[str]			=						1_6
		elif accelerator.mixed_precision != "no":
			_UpperCamelCase   :		List[str]			=						8
		else:
			_UpperCamelCase   :		int			=						None

		return tokenizer.pad(
		    UpperCAmelCase_							,       padding='longest'							,       max_length=UpperCAmelCase_							,       pad_to_multiple_of=UpperCAmelCase_							,       return_tensors='pt'							,       )

	# Instantiate dataloaders.
	_UpperCamelCase   :		int			=						DataLoader(
	    tokenized_datasets['train']							,       shuffle=UpperCAmelCase_							,       collate_fn=UpperCAmelCase_							,       batch_size=UpperCAmelCase_	)
	_UpperCamelCase   :		Any			=						DataLoader(
	    tokenized_datasets['validation']							,       shuffle=UpperCAmelCase_							,       collate_fn=UpperCAmelCase_							,       batch_size=UpperCAmelCase_	)

	return train_dataloader, eval_dataloader


# For testing only
if os.environ.get("""TESTING_MOCKED_DATALOADERS""", None) == "1":
							from accelerate.test_utils.training import mocked_dataloaders

							lowerCAmelCase__								=		mocked_dataloaders  # noqa: F811






def 					lowerCamelCase_     (							UpperCAmelCase_							:     List[str]							,       UpperCAmelCase_							:     str	)						->				Any:




	'''simple docstring'''




	if os.environ.get('TESTING_MOCKED_DATALOADERS'							,       UpperCAmelCase_	) == "1":
		_UpperCamelCase   :		Any			=						2
	# Initialize Accelerator

	# New Code #
	# We pass in "all" to `log_with` to grab all available trackers in the environment
	# Note: If using a custom `Tracker` class, should be passed in here such as:
	# >>> log_with = ["all", MyCustomTrackerClassInstance()]
	if args.with_tracking:
		_UpperCamelCase   :		Union[str, Any]			=						Accelerator(
		    cpu=args.cpu							,       mixed_precision=args.mixed_precision							,       log_with='all'							,       project_dir=args.project_dir	)
	else:
		_UpperCamelCase   :		Optional[int]			=						Accelerator(cpu=args.cpu							,       mixed_precision=args.mixed_precision	)
	# Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
	_UpperCamelCase   :		Any			=						config['lr']
	_UpperCamelCase   :		int			=						int(config['num_epochs']	)
	_UpperCamelCase   :		int			=						int(config['seed']	)
	_UpperCamelCase   :		Dict			=						int(config['batch_size']	)
	set_seed(UpperCAmelCase_	)

	_UpperCamelCase   :		Dict			=						get_dataloaders(UpperCAmelCase_							,       UpperCAmelCase_	)
	_UpperCamelCase   :		Dict			=						evaluate.load('glue'							,       'mrpc'	)

	# If the batch size is too big we use gradient accumulation
	_UpperCamelCase   :		Optional[Any]			=						1
	if batch_size > MAX_GPU_BATCH_SIZE and accelerator.distributed_type != DistributedType.TPU:
		_UpperCamelCase   :		Optional[int]			=						batch_size // MAX_GPU_BATCH_SIZE
		_UpperCamelCase   :		Any			=						MAX_GPU_BATCH_SIZE

	# Instantiate the model (we build the model here so that the seed also control new weights initialization)
	_UpperCamelCase   :		int			=						AutoModelForSequenceClassification.from_pretrained('bert-base-cased'							,       return_dict=UpperCAmelCase_	)

	# We could avoid this line since the accelerator is set with `device_placement=True` (default value).
	# Note that if you are placing tensors on devices manually, this line absolutely needs to be before the optimizer
	# creation otherwise training will not work on TPU (`accelerate` will kindly throw an error to make us aware of that).
	_UpperCamelCase   :		Tuple			=						model.to(accelerator.device	)

	# Instantiate optimizer
	_UpperCamelCase   :		int			=						AdamW(params=model.parameters()							,       lr=UpperCAmelCase_	)

	# Instantiate scheduler
	_UpperCamelCase   :		int			=						get_linear_schedule_with_warmup(
	    optimizer=UpperCAmelCase_							,       num_warmup_steps=1_0_0							,       num_training_steps=(len(UpperCAmelCase_	) * num_epochs) // gradient_accumulation_steps							,       )

	# Prepare everything
	# There is no specific order to remember, we just need to unpack the objects in the same order we gave them to the
	# prepare method.
	_UpperCamelCase   :		Optional[Any]			=						accelerator.prepare(
	    UpperCAmelCase_							,       UpperCAmelCase_							,       UpperCAmelCase_							,       UpperCAmelCase_							,       UpperCAmelCase_	)

	# New Code #
	# We need to initialize the trackers we use. Overall configurations can also be stored
	if args.with_tracking:
		_UpperCamelCase   :		Any			=						os.path.split(UpperCAmelCase_	)[-1].split('.'	)[0]
		accelerator.init_trackers(UpperCAmelCase_							,       UpperCAmelCase_	)

	# Now we train the model
	for epoch in range(UpperCAmelCase_	):
		model.train()
		# New Code #
		# For our tracking example, we will log the total loss of each epoch
		if args.with_tracking:
			_UpperCamelCase   :		Optional[int]			=						0
		for step, batch in enumerate(UpperCAmelCase_	):
			# We could avoid this line since we set the accelerator with `device_placement=True`.
			batch.to(accelerator.device	)
			_UpperCamelCase   :		str			=						model(**UpperCAmelCase_	)
			_UpperCamelCase   :		List[str]			=						outputs.loss
			# New Code #
			if args.with_tracking:
				total_loss += loss.detach().float()
			_UpperCamelCase   :		int			=						loss / gradient_accumulation_steps
			accelerator.backward(UpperCAmelCase_	)
			if step % gradient_accumulation_steps == 0:
				optimizer.step()
				lr_scheduler.step()
				optimizer.zero_grad()

		model.eval()
		for step, batch in enumerate(UpperCAmelCase_	):
			# We could avoid this line since we set the accelerator with `device_placement=True` (the default).
			batch.to(accelerator.device	)
			with torch.no_grad():
				_UpperCamelCase   :		List[str]			=						model(**UpperCAmelCase_	)
			_UpperCamelCase   :		List[Any]			=						outputs.logits.argmax(dim=-1	)
			_UpperCamelCase   :		Union[str, Any]			=						accelerator.gather_for_metrics((predictions, batch['labels'])	)
			metric.add_batch(
			    predictions=UpperCAmelCase_							,       references=UpperCAmelCase_							,       )

		_UpperCamelCase   :		Union[str, Any]			=						metric.compute()
		# Use accelerator.print to print only on the main process.
		accelerator.print(F'''epoch {epoch}:'''							,       UpperCAmelCase_	)

		# New Code #
		# To actually log, we call `Accelerator.log`
		# The values passed can be of `str`, `int`, `float` or `dict` of `str` to `float`/`int`
		if args.with_tracking:
			accelerator.log(
			    {
			        'accuracy': eval_metric['accuracy'],
			        'f1': eval_metric['f1'],
			        'train_loss': total_loss.item() / len(UpperCAmelCase_	),
			        'epoch': epoch,
			    }							,       step=UpperCAmelCase_							,       )

    # New Code #
    # When a run is finished, you should call `accelerator.end_training()`
    # to close all of the open trackers
	if args.with_tracking:
		accelerator.end_training()






def 					lowerCamelCase_     (							)						->				List[Any]:




	'''simple docstring'''




	_UpperCamelCase   :		Union[str, Any]			=						argparse.ArgumentParser(description='Simple example of training script.'	)
	parser.add_argument(
	    '--mixed_precision'							,       type=UpperCAmelCase_							,       default=UpperCAmelCase_							,       choices=['no', 'fp16', 'bf16', 'fp8']							,       help='Whether to use mixed precision. Choose'
	    'between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >= 1.10.'
	    'and an Nvidia Ampere GPU.'							,       )
	parser.add_argument('--cpu'							,       action='store_true'							,       help='If passed, will train on the CPU.'	)
	parser.add_argument(
	    '--with_tracking'							,       action='store_true'							,       help='Whether to load in all available experiment trackers from the environment and use them for logging.'							,       )
	parser.add_argument(
	    '--project_dir'							,       type=UpperCAmelCase_							,       default='logs'							,       help='Location on where to store experiment tracking logs` and relevent project information'							,       )
	_UpperCamelCase   :		Dict			=						parser.parse_args()
	_UpperCamelCase   :		Tuple			=						{'lr': 2e-5, 'num_epochs': 3, 'seed': 4_2, 'batch_size': 1_6}
	training_function(UpperCAmelCase_							,       UpperCAmelCase_	)


if __name__ == "__main__":
							main()







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


		"""simple docstring"""




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


		def 		A__					(				self):
			_UpperCamelCase   :		List[str]			=						ids_tensor([self.batch_size, self.seq_length]	,       self.vocab_size)

			_UpperCamelCase   :		Any			=						None
			if self.use_input_mask:
				_UpperCamelCase   :		int			=						random_attention_mask([self.batch_size, self.seq_length])

			_UpperCamelCase   :		Optional[Any]			=						None
			if self.use_token_type_ids:
				_UpperCamelCase   :		Dict			=						ids_tensor([self.batch_size, self.seq_length]	,       self.type_vocab_size)

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

			_UpperCamelCase   :		Union[str, Any]			=						self.get_config()

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


		def 		A__					(				self):
			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):
			_UpperCamelCase   :		List[str]			=						MegatronBertModel(config=__snake_case)
			model.to(__snake_case)
			model.eval()
			_UpperCamelCase   :		Optional[int]			=						model(__snake_case	,       attention_mask=__snake_case	,       token_type_ids=__snake_case)
			_UpperCamelCase   :		Dict			=						model(__snake_case	,       token_type_ids=__snake_case)
			_UpperCamelCase   :		Optional[Any]			=						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):
			_UpperCamelCase   :		int			=						MegatronBertForMaskedLM(config=__snake_case)
			model.to(__snake_case)
			model.eval()
			_UpperCamelCase   :		Dict			=						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):
			_UpperCamelCase   :		str			=						MegatronBertForCausalLM(config=__snake_case)
			model.to(__snake_case)
			model.eval()
			_UpperCamelCase   :		Optional[int]			=						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):
			_UpperCamelCase   :		Tuple			=						MegatronBertForNextSentencePrediction(config=__snake_case)
			model.to(__snake_case)
			model.eval()
			_UpperCamelCase   :		Optional[Any]			=						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):
			_UpperCamelCase   :		Optional[Any]			=						MegatronBertForPreTraining(config=__snake_case)
			model.to(__snake_case)
			model.eval()
			_UpperCamelCase   :		List[str]			=						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):
			_UpperCamelCase   :		int			=						MegatronBertForQuestionAnswering(config=__snake_case)
			model.to(__snake_case)
			model.eval()
			_UpperCamelCase   :		List[Any]			=						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):
			_UpperCamelCase   :		Optional[int]			=						self.num_labels
			_UpperCamelCase   :		Union[str, Any]			=						MegatronBertForSequenceClassification(__snake_case)
			model.to(__snake_case)
			model.eval()
			_UpperCamelCase   :		str			=						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):
			_UpperCamelCase   :		Any			=						self.num_labels
			_UpperCamelCase   :		Optional[int]			=						MegatronBertForTokenClassification(config=__snake_case)
			model.to(__snake_case)
			model.eval()
			_UpperCamelCase   :		Tuple			=						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):
			_UpperCamelCase   :		List[str]			=						self.num_choices
			_UpperCamelCase   :		Optional[int]			=						MegatronBertForMultipleChoice(config=__snake_case)
			model.to(__snake_case)
			model.eval()
			_UpperCamelCase   :		List[Any]			=						input_ids.unsqueeze(1).expand(-1	,       self.num_choices	,       -1).contiguous()
			_UpperCamelCase   :		List[Any]			=						token_type_ids.unsqueeze(1).expand(-1	,       self.num_choices	,       -1).contiguous()
			_UpperCamelCase   :		Optional[Any]			=						input_mask.unsqueeze(1).expand(-1	,       self.num_choices	,       -1).contiguous()
			_UpperCamelCase   :		Union[str, Any]			=						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):
			_UpperCamelCase   :		Union[str, Any]			=						self.prepare_config_and_inputs()
			(
			    (
			    _UpperCamelCase
			)	,		(
			    _UpperCamelCase
			)	,		(
			    _UpperCamelCase
			)	,		(
			    _UpperCamelCase
			)	,		(
			    _UpperCamelCase
			)	,		(
			    _UpperCamelCase
			)	,		(
			    _UpperCamelCase
			)	,		
			)   :		Optional[int]			=						config_and_inputs
			_UpperCamelCase   :		int			=						{'input_ids': input_ids, 'token_type_ids': token_type_ids, 'attention_mask': input_mask}
			return config, inputs_dict







@require_torch
class 	lowercase    (				_lowercase    ,							_lowercase    ,							unittest.TestCase							):


		"""simple docstring"""




		a__					       =  (
		    (
		        MegatronBertModel,
		        MegatronBertForMaskedLM,
		        MegatronBertForCausalLM,
		        MegatronBertForMultipleChoice,
		        MegatronBertForNextSentencePrediction,
		        MegatronBertForPreTraining,
		        MegatronBertForQuestionAnswering,
		        MegatronBertForSequenceClassification,
		        MegatronBertForTokenClassification,
		    )
		    if is_torch_available()
		    else ()
		)
		a__					       =  (
		    {
		        "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 {}
		)
		a__					       =  True
		# test_resize_embeddings = False
		a__					       =  False


		def 		A__					(				self	,       __snake_case	,       __snake_case	,       __snake_case=False):
			_UpperCamelCase   :		str			=						super()._prepare_for_class(__snake_case	,       __snake_case	,       return_labels=__snake_case)

			if return_labels:
				if model_class in get_values(__snake_case):
					_UpperCamelCase   :		Optional[Any]			=						torch.zeros(
					    (self.model_tester.batch_size, self.model_tester.seq_length)	,       dtype=torch.long	,       device=__snake_case)
					_UpperCamelCase   :		str			=						torch.zeros(
					    self.model_tester.batch_size	,       dtype=torch.long	,       device=__snake_case)
			return inputs_dict


		def 		A__					(				self):
			_UpperCamelCase   :		Any			=						MegatronBertModelTester(self)
			_UpperCamelCase   :		int			=						ConfigTester(self	,       config_class=__snake_case	,       hidden_size=37)


		def 		A__					(				self):
			self.config_tester.run_common_tests()


		def 		A__					(				self):
			_UpperCamelCase   :		Optional[int]			=						self.model_tester.prepare_config_and_inputs()
			self.model_tester.create_and_check_megatron_bert_model(*__snake_case)


		def 		A__					(				self):
			_UpperCamelCase   :		Dict			=						self.model_tester.prepare_config_and_inputs()
			self.model_tester.create_and_check_megatron_bert_for_masked_lm(*__snake_case)


		def 		A__					(				self):
			_UpperCamelCase   :		int			=						self.model_tester.prepare_config_and_inputs()
			self.model_tester.create_and_check_megatron_bert_for_multiple_choice(*__snake_case)


		def 		A__					(				self):
			_UpperCamelCase   :		Dict			=						self.model_tester.prepare_config_and_inputs()
			self.model_tester.create_and_check_megatron_bert_for_next_sequence_prediction(*__snake_case)


		def 		A__					(				self):
			_UpperCamelCase   :		str			=						self.model_tester.prepare_config_and_inputs()
			self.model_tester.create_and_check_megatron_bert_for_pretraining(*__snake_case)


		def 		A__					(				self):
			_UpperCamelCase   :		List[Any]			=						self.model_tester.prepare_config_and_inputs()
			self.model_tester.create_and_check_megatron_bert_for_question_answering(*__snake_case)


		def 		A__					(				self):
			_UpperCamelCase   :		Optional[Any]			=						self.model_tester.prepare_config_and_inputs()
			self.model_tester.create_and_check_megatron_bert_for_sequence_classification(*__snake_case)






		def 		A__					(				self):
			_UpperCamelCase   :		Any			=						self.model_tester.prepare_config_and_inputs()
			self.model_tester.create_and_check_megatron_bert_for_token_classification(*__snake_case)






def 					lowerCamelCase_     (							UpperCAmelCase_							:     str	)						->				Optional[Any]:




	'''simple docstring'''




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


lowerCAmelCase__								=		1e-4







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


		"""simple docstring"""




		@slow
		@unittest.skip('Model is not available.')
		def 		A__					(				self):
			_UpperCamelCase   :		int			=						'nvidia/megatron-bert-uncased-345m'
			if "MYDIR" in os.environ:
				_UpperCamelCase   :		int			=						os.path.join(os.environ['MYDIR']	,       __snake_case)
			_UpperCamelCase   :		Optional[int]			=						MegatronBertModel.from_pretrained(__snake_case)
			model.to(__snake_case)
			model.half()
			_UpperCamelCase   :		Optional[Any]			=						_long_tensor([[1_01, 71_10, 10_05, 10_56, 20_23, 1_13_33, 1_74_13, 10_29, 1_02]])
			with torch.no_grad():
				_UpperCamelCase   :		str			=						model(__snake_case)[0]
			_UpperCamelCase   :		Optional[int]			=						torch.Size((1, 9, 10_24))
			self.assertEqual(output.shape	,       __snake_case)

			_UpperCamelCase   :		Union[str, Any]			=						[-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):
					_UpperCamelCase   :		Optional[Any]			=						output[0, ii, jj]
					_UpperCamelCase   :		Dict			=						expected[3 * ii + jj]
					_UpperCamelCase   :		Optional[int]			=						'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)




