


"""simple docstring"""


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







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







						def __a						(							self   :    Any		):
								A					       = ids_tensor([self.batch_size, self.seq_length]			,			self.vocab_size		)

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

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

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

								A					       = self.get_config()

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







						def __a						(							self   :    Union[str, Any]		):
								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=UpperCamelCase__			,			initializer_range=self.initializer_range			,			)







						def __a						(							self   :    Optional[Any]			,			_lowercase   :    List[Any]			,			_lowercase   :    Tuple			,			_lowercase   :    Optional[int]			,			_lowercase   :    Dict			,			_lowercase   :    List[str]			,			_lowercase   :    Dict			,			_lowercase   :    List[Any]		):
								A					       = MegatronBertModel(config=UpperCamelCase__		)
								model.to(UpperCamelCase__		)
								model.eval()
								A					       = model(UpperCamelCase__			,			attention_mask=UpperCamelCase__			,			token_type_ids=UpperCamelCase__		)
								A					       = model(UpperCamelCase__			,			token_type_ids=UpperCamelCase__		)
								A					       = model(UpperCamelCase__		)

								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   :    List[Any]			,			_lowercase   :    List[Any]			,			_lowercase   :    Any			,			_lowercase   :    Tuple			,			_lowercase   :    Any			,			_lowercase   :    Union[str, Any]			,			_lowercase   :    Tuple			,			_lowercase   :    Any		):
								A					       = MegatronBertForMaskedLM(config=UpperCamelCase__		)
								model.to(UpperCamelCase__		)
								model.eval()
								A					       = model(UpperCamelCase__			,			attention_mask=UpperCamelCase__			,			token_type_ids=UpperCamelCase__			,			labels=UpperCamelCase__		)
								self.parent.assertEqual(result.logits.shape			,			(self.batch_size, self.seq_length, self.vocab_size)		)







						def __a						(							self   :    Tuple			,			_lowercase   :    Optional[int]			,			_lowercase   :    str			,			_lowercase   :    str			,			_lowercase   :    List[str]			,			_lowercase   :    List[str]			,			_lowercase   :    Dict			,			_lowercase   :    Dict		):
								A					       = MegatronBertForCausalLM(config=UpperCamelCase__		)
								model.to(UpperCamelCase__		)
								model.eval()
								A					       = model(UpperCamelCase__			,			attention_mask=UpperCamelCase__			,			token_type_ids=UpperCamelCase__			,			labels=UpperCamelCase__		)
								self.parent.assertEqual(result.logits.shape			,			(self.batch_size, self.seq_length, self.vocab_size)		)







						def __a						(							self   :    Optional[Any]			,			_lowercase   :    Dict			,			_lowercase   :    Tuple			,			_lowercase   :    Optional[Any]			,			_lowercase   :    Union[str, Any]			,			_lowercase   :    Dict			,			_lowercase   :    Optional[int]			,			_lowercase   :    List[Any]		):
								A					       = MegatronBertForNextSentencePrediction(config=UpperCamelCase__		)
								model.to(UpperCamelCase__		)
								model.eval()
								A					       = model(
								    UpperCamelCase__			,			attention_mask=UpperCamelCase__			,			token_type_ids=UpperCamelCase__			,			labels=UpperCamelCase__			,			)
								self.parent.assertEqual(result.logits.shape			,			(self.batch_size, 2)		)







						def __a						(							self   :    List[Any]			,			_lowercase   :    Dict			,			_lowercase   :    Optional[int]			,			_lowercase   :    Tuple			,			_lowercase   :    int			,			_lowercase   :    List[Any]			,			_lowercase   :    int			,			_lowercase   :    Dict		):
								A					       = MegatronBertForPreTraining(config=UpperCamelCase__		)
								model.to(UpperCamelCase__		)
								model.eval()
								A					       = model(
								    UpperCamelCase__			,			attention_mask=UpperCamelCase__			,			token_type_ids=UpperCamelCase__			,			labels=UpperCamelCase__			,			next_sentence_label=UpperCamelCase__			,			)
								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   :    Tuple			,			_lowercase   :    int			,			_lowercase   :    Any			,			_lowercase   :    Dict			,			_lowercase   :    Union[str, Any]			,			_lowercase   :    str			,			_lowercase   :    Union[str, Any]			,			_lowercase   :    Tuple		):
								A					       = MegatronBertForQuestionAnswering(config=UpperCamelCase__		)
								model.to(UpperCamelCase__		)
								model.eval()
								A					       = model(
								    UpperCamelCase__			,			attention_mask=UpperCamelCase__			,			token_type_ids=UpperCamelCase__			,			start_positions=UpperCamelCase__			,			end_positions=UpperCamelCase__			,			)
								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   :    Any			,			_lowercase   :    Any			,			_lowercase   :    int			,			_lowercase   :    List[str]			,			_lowercase   :    Dict			,			_lowercase   :    List[Any]			,			_lowercase   :    str			,			_lowercase   :    List[Any]		):
								A					       = self.num_labels
								A					       = MegatronBertForSequenceClassification(UpperCamelCase__		)
								model.to(UpperCamelCase__		)
								model.eval()
								A					       = model(UpperCamelCase__			,			attention_mask=UpperCamelCase__			,			token_type_ids=UpperCamelCase__			,			labels=UpperCamelCase__		)
								self.parent.assertEqual(result.logits.shape			,			(self.batch_size, self.num_labels)		)







						def __a						(							self   :    Union[str, Any]			,			_lowercase   :    Tuple			,			_lowercase   :    str			,			_lowercase   :    Optional[Any]			,			_lowercase   :    Optional[int]			,			_lowercase   :    str			,			_lowercase   :    Any			,			_lowercase   :    Optional[int]		):
								A					       = self.num_labels
								A					       = MegatronBertForTokenClassification(config=UpperCamelCase__		)
								model.to(UpperCamelCase__		)
								model.eval()
								A					       = model(UpperCamelCase__			,			attention_mask=UpperCamelCase__			,			token_type_ids=UpperCamelCase__			,			labels=UpperCamelCase__		)
								self.parent.assertEqual(result.logits.shape			,			(self.batch_size, self.seq_length, self.num_labels)		)







						def __a						(							self   :    Optional[int]			,			_lowercase   :    List[Any]			,			_lowercase   :    Optional[Any]			,			_lowercase   :    Union[str, Any]			,			_lowercase   :    Any			,			_lowercase   :    int			,			_lowercase   :    List[Any]			,			_lowercase   :    int		):
								A					       = self.num_choices
								A					       = MegatronBertForMultipleChoice(config=UpperCamelCase__		)
								model.to(UpperCamelCase__		)
								model.eval()
								A					       = input_ids.unsqueeze(1		).expand(-1			,			self.num_choices			,			-1		).contiguous()
								A					       = token_type_ids.unsqueeze(1		).expand(-1			,			self.num_choices			,			-1		).contiguous()
								A					       = input_mask.unsqueeze(1		).expand(-1			,			self.num_choices			,			-1		).contiguous()
								A					       = model(
								    UpperCamelCase__			,			attention_mask=UpperCamelCase__			,			token_type_ids=UpperCamelCase__			,			labels=UpperCamelCase__			,			)
								self.parent.assertEqual(result.logits.shape			,			(self.batch_size, self.num_choices)		)







						def __a						(							self   :    Union[str, Any]		):
								A					       = self.prepare_config_and_inputs()
								(
								    A       
								)					       = config_and_inputs
								A					       = {'input_ids': input_ids, 'token_type_ids': token_type_ids, 'attention_mask': input_mask}
								return config, inputs_dict



@require_torch
class 			lowerCamelCase__						(   __A							,    __A							,    unittest.TestCase			):
						lowerCAmelCase       =     (
						    (
						        MegatronBertModel,
						        MegatronBertForMaskedLM,
						        MegatronBertForCausalLM,
						        MegatronBertForMultipleChoice,
						        MegatronBertForNextSentencePrediction,
						        MegatronBertForPreTraining,
						        MegatronBertForQuestionAnswering,
						        MegatronBertForSequenceClassification,
						        MegatronBertForTokenClassification,
						    )
						    if is_torch_available()
						    else ()
						)
						lowerCAmelCase       =     (
						    {
						        """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 {}
						)
						lowerCAmelCase       =     True
						# test_resize_embeddings = False
						lowerCAmelCase       =     False







						def __a						(							self   :    Union[str, Any]			,			_lowercase   :    Optional[int]			,			_lowercase   :    Dict			,			_lowercase   :    Dict=False		):
								A					       = super()._prepare_for_class(UpperCamelCase__			,			UpperCamelCase__			,			return_labels=UpperCamelCase__		)

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







						def __a						(							self   :    List[str]		):
								A					       = MegatronBertModelTester(self		)
								A					       = ConfigTester(self			,			config_class=UpperCamelCase__			,			hidden_size=37		)







						def __a						(							self   :    Union[str, Any]		):
								self.config_tester.run_common_tests()







						def __a						(							self   :    Optional[int]		):
								A					       = self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_megatron_bert_model(*UpperCamelCase__		)







						def __a						(							self   :    int		):
								A					       = self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_megatron_bert_for_masked_lm(*UpperCamelCase__		)







						def __a						(							self   :    Union[str, Any]		):
								A					       = self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_megatron_bert_for_multiple_choice(*UpperCamelCase__		)







						def __a						(							self   :    str		):
								A					       = self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_megatron_bert_for_next_sequence_prediction(*UpperCamelCase__		)







						def __a						(							self   :    List[str]		):
								A					       = self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_megatron_bert_for_pretraining(*UpperCamelCase__		)







						def __a						(							self   :    Optional[Any]		):
								A					       = self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_megatron_bert_for_question_answering(*UpperCamelCase__		)







						def __a						(							self   :    Dict		):
								A					       = self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_megatron_bert_for_sequence_classification(*UpperCamelCase__		)







						def __a						(							self   :    List[Any]		):
								A					       = self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_megatron_bert_for_token_classification(*UpperCamelCase__		)







def  __snake_case       (   UpperCamelCase__	)      ->		List[str]:







		"""simple docstring"""





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


UpperCamelCase						:				List[Any]             =      1E-4



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







						@slow
						@unittest.skip('Model is not available.'		)
						def __a						(							self   :    Tuple		):
								A					       = 'nvidia/megatron-bert-uncased-345m'
								if "MYDIR" in os.environ:
										A					       = os.path.join(os.environ['MYDIR']			,			UpperCamelCase__		)
								A					       = MegatronBertModel.from_pretrained(UpperCamelCase__		)
								model.to(UpperCamelCase__		)
								model.half()
								A					       = _long_tensor([[101, 7_110, 1_005, 1_056, 2_023, 11_333, 17_413, 1_029, 102]]		)
								with torch.no_grad():
										A					       = model(UpperCamelCase__		)[0]
								A					       = torch.Size((1, 9, 1_024)		)
								self.assertEqual(output.shape			,			UpperCamelCase__		)

								A					       = [-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		):
												A					       = output[0, ii, jj]
												A					       = expected[3 * ii + jj]
												A					       = 'ii={} jj={} a={} b={}'.format(UpperCamelCase__			,			UpperCamelCase__			,			UpperCamelCase__			,			UpperCamelCase__		)
												self.assertTrue(math.isclose(UpperCamelCase__			,			UpperCamelCase__			,			rel_tol=UpperCamelCase__			,			abs_tol=UpperCamelCase__		)			,			msg=UpperCamelCase__		)









"""simple docstring"""


from statistics import mean, stdev







def  __snake_case       (   UpperCamelCase__ ,			UpperCamelCase__ = 3	)      ->		list:







		"""simple docstring"""





		A					       = min(UpperCamelCase__	)
		A					       = max(UpperCamelCase__	)
		# normalize data
		return [round((x - x_min) / (x_max - x_min) ,			UpperCamelCase__	) for x in data]







def  __snake_case       (   UpperCamelCase__ ,			UpperCamelCase__ = 3	)      ->		list:







		"""simple docstring"""





		A					       = mean(UpperCamelCase__	)
		A					       = stdev(UpperCamelCase__	)
		# standardize data
		return [round((x - mu) / (sigma) ,			UpperCamelCase__	) for x in data]





