

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






			"""simple docstring"""




			def __init__(    self ,				__UpperCamelCase ,				__UpperCamelCase=13 ,				__UpperCamelCase=7 ,				__UpperCamelCase=True ,				__UpperCamelCase=True ,				__UpperCamelCase=True ,				__UpperCamelCase=True ,				__UpperCamelCase=99 ,				__UpperCamelCase=64 ,				__UpperCamelCase=32 ,				__UpperCamelCase=5 ,				__UpperCamelCase=4 ,				__UpperCamelCase=37 ,				__UpperCamelCase="gelu" ,				__UpperCamelCase=0.1 ,				__UpperCamelCase=0.1 ,				__UpperCamelCase=5_12 ,				__UpperCamelCase=16 ,				__UpperCamelCase=2 ,				__UpperCamelCase=0.02 ,				__UpperCamelCase=3 ,				__UpperCamelCase=4 ,				__UpperCamelCase=None ,				):

				"""simple docstring"""


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


			def 			__lowerCAmelCase (    self							):

				"""simple docstring"""


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

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

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

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

				snake_case_      		= self.get_config()

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


			def 			__lowerCAmelCase (    self							):

				"""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=__UpperCamelCase ,				initializer_range=self.initializer_range ,				)


			def 			__lowerCAmelCase (    self ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase							):

				"""simple docstring"""


				snake_case_      		= MegatronBertModel(config=__UpperCamelCase							)
				model.to(__UpperCamelCase							)
				model.eval()
				snake_case_      		= model(__UpperCamelCase ,				attention_mask=__UpperCamelCase ,				token_type_ids=__UpperCamelCase							)
				snake_case_      		= model(__UpperCamelCase ,				token_type_ids=__UpperCamelCase							)
				snake_case_      		= 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 			__lowerCAmelCase (    self ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase							):

				"""simple docstring"""


				snake_case_      		= MegatronBertForMaskedLM(config=__UpperCamelCase							)
				model.to(__UpperCamelCase							)
				model.eval()
				snake_case_      		= 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 			__lowerCAmelCase (    self ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase							):

				"""simple docstring"""


				snake_case_      		= MegatronBertForCausalLM(config=__UpperCamelCase							)
				model.to(__UpperCamelCase							)
				model.eval()
				snake_case_      		= 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 			__lowerCAmelCase (    self ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase							):

				"""simple docstring"""


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


			def 			__lowerCAmelCase (    self ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase							):

				"""simple docstring"""


				snake_case_      		= MegatronBertForPreTraining(config=__UpperCamelCase							)
				model.to(__UpperCamelCase							)
				model.eval()
				snake_case_      		= 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 			__lowerCAmelCase (    self ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase							):

				"""simple docstring"""


				snake_case_      		= MegatronBertForQuestionAnswering(config=__UpperCamelCase							)
				model.to(__UpperCamelCase							)
				model.eval()
				snake_case_      		= 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 			__lowerCAmelCase (    self ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase							):

				"""simple docstring"""


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


			def 			__lowerCAmelCase (    self ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase							):

				"""simple docstring"""


				snake_case_      		= self.num_labels
				snake_case_      		= MegatronBertForTokenClassification(config=__UpperCamelCase							)
				model.to(__UpperCamelCase							)
				model.eval()
				snake_case_      		= 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 			__lowerCAmelCase (    self ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase							):

				"""simple docstring"""


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







			def 			__lowerCAmelCase (    self							):

				"""simple docstring"""


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







@require_torch
class SCREAMING_SNAKE_CASE  (		__snake_case							,						__snake_case							,						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 			__lowerCAmelCase (    self ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase=False							):

				"""simple docstring"""


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

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


			def 			__lowerCAmelCase (    self							):

				"""simple docstring"""


				snake_case_      		= MegatronBertModelTester(self							)
				snake_case_      		= ConfigTester(self ,				config_class=__UpperCamelCase ,				hidden_size=37							)


			def 			__lowerCAmelCase (    self							):

				"""simple docstring"""


				self.config_tester.run_common_tests()


			def 			__lowerCAmelCase (    self							):

				"""simple docstring"""


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


			def 			__lowerCAmelCase (    self							):

				"""simple docstring"""


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


			def 			__lowerCAmelCase (    self							):

				"""simple docstring"""


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


			def 			__lowerCAmelCase (    self							):

				"""simple docstring"""


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


			def 			__lowerCAmelCase (    self							):

				"""simple docstring"""


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


			def 			__lowerCAmelCase (    self							):

				"""simple docstring"""


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


			def 			__lowerCAmelCase (    self							):

				"""simple docstring"""


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







			def 			__lowerCAmelCase (    self							):

				"""simple docstring"""


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







def a(lowercase__	):





	'''simple docstring'''
	return torch.tensor(
	    lowercase__			, dtype=torch.long			, device=lowercase__			, )


A											=     1e-4






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






			"""simple docstring"""




			@slow
			@unittest.skip('Model is not available.'							)
			def 			__lowerCAmelCase (    self							):

				"""simple docstring"""


				snake_case_      		= 'nvidia/megatron-bert-uncased-345m'
				if "MYDIR" in os.environ:
					snake_case_      		= os.path.join(os.environ['MYDIR'] ,				__UpperCamelCase							)
				snake_case_      		= MegatronBertModel.from_pretrained(__UpperCamelCase							)
				model.to(__UpperCamelCase							)
				model.half()
				snake_case_      		= _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():
					snake_case_      		= model(__UpperCamelCase							)[0]
				snake_case_      		= torch.Size((1, 9, 10_24)							)
				self.assertEqual(output.shape ,				__UpperCamelCase							)

				snake_case_      		= [-0.6040, -0.2517, -0.1025, 0.3420, -0.6758, -0.0017, -0.1089, -0.1990, 0.5728]
				for ii in range(3							):
					for jj in range(3							):
						snake_case_      		= output[0, ii, jj]
						snake_case_      		= expected[3 * ii + jj]
						snake_case_      		= 'ii={} jj={} a={} b={}'.format(__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase							)
						self.assertTrue(math.isclose(__UpperCamelCase ,				__UpperCamelCase ,				rel_tol=__UpperCamelCase ,				abs_tol=__UpperCamelCase							) ,				msg=__UpperCamelCase							)



def a(lowercase__			, lowercase__	):





	'''simple docstring'''
	if not isinstance(lowercase__			, lowercase__	):
		raise ValueError('iterations must be defined as integers'	)
	if not isinstance(lowercase__			, lowercase__	) or not number >= 1:
		raise ValueError(
		    'starting number must be\n                         and integer and be more than 0'	)
	if not iterations >= 1:
		raise ValueError('Iterations must be done more than 0 times to play FizzBuzz'	)

	snake_case_      		= ''
	while number <= iterations:
		if number % 3 == 0:
			out += "Fizz"
		if number % 5 == 0:
			out += "Buzz"
		if 0 not in (number % 3, number % 5):
			out += str(lowercase__	)

		# print(out)
		number += 1
		out += " "
	return out


if __name__ == "__main__":
		import doctest

		doctest.testmod()
