

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





			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=512					,  __snake_case=16					,  __snake_case=2					,  __snake_case=0.02					,  __snake_case=3					,  __snake_case=4					,  __snake_case=None					,  )   -> str:




						'''simple docstring'''
						__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 		__magic_name__    (		self					)   -> Tuple:




						'''simple docstring'''
						__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 		__magic_name__    (		self					)   -> Optional[int]:




						'''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=lowerCamelCase__					,  initializer_range=self.initializer_range					,  )





			def 		__magic_name__    (		self					,  __snake_case					,  __snake_case					,  __snake_case					,  __snake_case					,  __snake_case					,  __snake_case					,  __snake_case					)   -> List[Any]:




						'''simple docstring'''
						__a        =MegatronBertModel(config=lowerCamelCase__					)
						model.to(lowerCamelCase__					)
						model.eval()
						__a        =model(lowerCamelCase__					,  attention_mask=lowerCamelCase__					,  token_type_ids=lowerCamelCase__					)
						__a        =model(lowerCamelCase__					,  token_type_ids=lowerCamelCase__					)
						__a        =model(lowerCamelCase__					)

						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 		__magic_name__    (		self					,  __snake_case					,  __snake_case					,  __snake_case					,  __snake_case					,  __snake_case					,  __snake_case					,  __snake_case					)   -> List[str]:




						'''simple docstring'''
						__a        =MegatronBertForMaskedLM(config=lowerCamelCase__					)
						model.to(lowerCamelCase__					)
						model.eval()
						__a        =model(lowerCamelCase__					,  attention_mask=lowerCamelCase__					,  token_type_ids=lowerCamelCase__					,  labels=lowerCamelCase__					)
						self.parent.assertEqual(result.logits.shape					,  (self.batch_size, self.seq_length, self.vocab_size)					)





			def 		__magic_name__    (		self					,  __snake_case					,  __snake_case					,  __snake_case					,  __snake_case					,  __snake_case					,  __snake_case					,  __snake_case					)   -> Optional[Any]:




						'''simple docstring'''
						__a        =MegatronBertForCausalLM(config=lowerCamelCase__					)
						model.to(lowerCamelCase__					)
						model.eval()
						__a        =model(lowerCamelCase__					,  attention_mask=lowerCamelCase__					,  token_type_ids=lowerCamelCase__					,  labels=lowerCamelCase__					)
						self.parent.assertEqual(result.logits.shape					,  (self.batch_size, self.seq_length, self.vocab_size)					)





			def 		__magic_name__    (		self					,  __snake_case					,  __snake_case					,  __snake_case					,  __snake_case					,  __snake_case					,  __snake_case					,  __snake_case					)   -> List[str]:




						'''simple docstring'''
						__a        =MegatronBertForNextSentencePrediction(config=lowerCamelCase__					)
						model.to(lowerCamelCase__					)
						model.eval()
						__a        =model(
						    lowerCamelCase__					,  attention_mask=lowerCamelCase__					,  token_type_ids=lowerCamelCase__					,  labels=lowerCamelCase__					,  )
						self.parent.assertEqual(result.logits.shape					,  (self.batch_size, 2)					)





			def 		__magic_name__    (		self					,  __snake_case					,  __snake_case					,  __snake_case					,  __snake_case					,  __snake_case					,  __snake_case					,  __snake_case					)   -> Optional[int]:




						'''simple docstring'''
						__a        =MegatronBertForPreTraining(config=lowerCamelCase__					)
						model.to(lowerCamelCase__					)
						model.eval()
						__a        =model(
						    lowerCamelCase__					,  attention_mask=lowerCamelCase__					,  token_type_ids=lowerCamelCase__					,  labels=lowerCamelCase__					,  next_sentence_label=lowerCamelCase__					,  )
						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 		__magic_name__    (		self					,  __snake_case					,  __snake_case					,  __snake_case					,  __snake_case					,  __snake_case					,  __snake_case					,  __snake_case					)   -> Optional[int]:




						'''simple docstring'''
						__a        =MegatronBertForQuestionAnswering(config=lowerCamelCase__					)
						model.to(lowerCamelCase__					)
						model.eval()
						__a        =model(
						    lowerCamelCase__					,  attention_mask=lowerCamelCase__					,  token_type_ids=lowerCamelCase__					,  start_positions=lowerCamelCase__					,  end_positions=lowerCamelCase__					,  )
						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 		__magic_name__    (		self					,  __snake_case					,  __snake_case					,  __snake_case					,  __snake_case					,  __snake_case					,  __snake_case					,  __snake_case					)   -> str:




						'''simple docstring'''
						__a        =self.num_labels
						__a        =MegatronBertForSequenceClassification(lowerCamelCase__					)
						model.to(lowerCamelCase__					)
						model.eval()
						__a        =model(lowerCamelCase__					,  attention_mask=lowerCamelCase__					,  token_type_ids=lowerCamelCase__					,  labels=lowerCamelCase__					)
						self.parent.assertEqual(result.logits.shape					,  (self.batch_size, self.num_labels)					)





			def 		__magic_name__    (		self					,  __snake_case					,  __snake_case					,  __snake_case					,  __snake_case					,  __snake_case					,  __snake_case					,  __snake_case					)   -> str:




						'''simple docstring'''
						__a        =self.num_labels
						__a        =MegatronBertForTokenClassification(config=lowerCamelCase__					)
						model.to(lowerCamelCase__					)
						model.eval()
						__a        =model(lowerCamelCase__					,  attention_mask=lowerCamelCase__					,  token_type_ids=lowerCamelCase__					,  labels=lowerCamelCase__					)
						self.parent.assertEqual(result.logits.shape					,  (self.batch_size, self.seq_length, self.num_labels)					)





			def 		__magic_name__    (		self					,  __snake_case					,  __snake_case					,  __snake_case					,  __snake_case					,  __snake_case					,  __snake_case					,  __snake_case					)   -> str:




						'''simple docstring'''
						__a        =self.num_choices
						__a        =MegatronBertForMultipleChoice(config=lowerCamelCase__					)
						model.to(lowerCamelCase__					)
						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(
						    lowerCamelCase__					,  attention_mask=lowerCamelCase__					,  token_type_ids=lowerCamelCase__					,  labels=lowerCamelCase__					,  )
						self.parent.assertEqual(result.logits.shape					,  (self.batch_size, self.num_choices)					)





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




						'''simple docstring'''
						__a        =self.prepare_config_and_inputs()
						(
						    (
						    __a						
						)      ,    (
						    __a						
						)      ,    (
						    __a						
						)      ,    (
						    __a						
						)      ,    (
						    __a						
						)      ,    (
						    __a						
						)      ,    (
						    __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 	__magic_name__    (			lowerCAmelCase_  ,    lowerCAmelCase_  ,    unittest.TestCase      ):
			SCREAMING_SNAKE_CASE							   =						(
			    (
			        MegatronBertModel,
			        MegatronBertForMaskedLM,
			        MegatronBertForCausalLM,
			        MegatronBertForMultipleChoice,
			        MegatronBertForNextSentencePrediction,
			        MegatronBertForPreTraining,
			        MegatronBertForQuestionAnswering,
			        MegatronBertForSequenceClassification,
			        MegatronBertForTokenClassification,
			    )
			    if is_torch_available()
			    else ()
			)
			SCREAMING_SNAKE_CASE							   =						(
			    {
			        '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 {}
			)
			SCREAMING_SNAKE_CASE							   =						True
			# test_resize_embeddings = False
			SCREAMING_SNAKE_CASE							   =						False





			def 		__magic_name__    (		self					,  __snake_case					,  __snake_case					,  __snake_case=False					)   -> Dict:




						'''simple docstring'''
						__a        =super()._prepare_for_class(lowerCamelCase__					,  lowerCamelCase__					,  return_labels=lowerCamelCase__					)

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





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




						'''simple docstring'''
						__a        =MegatronBertModelTester(self					)
						__a        =ConfigTester(self					,  config_class=lowerCamelCase__					,  hidden_size=37					)





			def 		__magic_name__    (		self					)   -> Dict:




						'''simple docstring'''
						self.config_tester.run_common_tests()





			def 		__magic_name__    (		self					)   -> Tuple:




						'''simple docstring'''
						__a        =self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_megatron_bert_model(*lowerCamelCase__					)





			def 		__magic_name__    (		self					)   -> Tuple:




						'''simple docstring'''
						__a        =self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_megatron_bert_for_masked_lm(*lowerCamelCase__					)





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




						'''simple docstring'''
						__a        =self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_megatron_bert_for_multiple_choice(*lowerCamelCase__					)





			def 		__magic_name__    (		self					)   -> int:




						'''simple docstring'''
						__a        =self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_megatron_bert_for_next_sequence_prediction(*lowerCamelCase__					)





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




						'''simple docstring'''
						__a        =self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_megatron_bert_for_pretraining(*lowerCamelCase__					)





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




						'''simple docstring'''
						__a        =self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_megatron_bert_for_question_answering(*lowerCamelCase__					)





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




						'''simple docstring'''
						__a        =self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_megatron_bert_for_sequence_classification(*lowerCamelCase__					)





			def 		__magic_name__    (		self					)   -> Dict:




						'''simple docstring'''
						__a        =self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_megatron_bert_for_token_classification(*lowerCamelCase__					)






def        UpperCamelCase_(							_snake_case   :  int	):




			"""simple docstring"""






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


_lowerCAmelCase					:			Dict  				=			1E-4






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





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




						'''simple docstring'''
						__a        ='nvidia/megatron-bert-uncased-345m'
						if "MYDIR" in os.environ:
									__a        =os.path.join(os.environ['MYDIR']					,  lowerCamelCase__					)
						__a        =MegatronBertModel.from_pretrained(lowerCamelCase__					)
						model.to(lowerCamelCase__					)
						model.half()
						__a        =_long_tensor([[101, 7110, 1005, 1056, 2023, 1_1333, 1_7413, 1029, 102]]					)
						with torch.no_grad():
									__a        =model(lowerCamelCase__					)[0]
						__a        =torch.Size((1, 9, 1024)					)
						self.assertEqual(output.shape					,  lowerCamelCase__					)

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












from math import sqrt







def 			lowerCamelCase_		(       UpperCamelCase__				:      int      )     ->       bool:





				"""simple docstring"""


				if 1 < number < 4:
								# 2 and 3 are primes
								return True
				elif number < 2 or number % 2 == 0 or number % 3 == 0:
								# Negatives, 0, 1, all even numbers, all multiples of 3 are not primes
								return False

				# All primes number are in format of 6k +/- 1
				for i in range(5 ,     int(sqrt(UpperCamelCase__      ) + 1      ) ,     6      ):
								if number % i == 0 or number % (i + 2) == 0:
												return False
				return True







def 			lowerCamelCase_		(       UpperCamelCase__				:      int = 1_0001      )     ->       int:





				"""simple docstring"""


				__lowerCamelCase		       =    0
				__lowerCamelCase		       =    1
				while count != nth and number < 3:
								number += 1
								if is_prime(UpperCamelCase__      ):
												count += 1
				while count != nth:
								number += 2
								if is_prime(UpperCamelCase__      ):
												count += 1
				return number


if __name__ == "__main__":
					print(f'''{solution() = }''')


