

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       :




				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   ,     ):
											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 					a							(      self ):
											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 					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 ):
											snake_case_           =					MegatronBertModel(config=snake_case )
											model.to(snake_case )
											model.eval()
											snake_case_           =					model(snake_case   ,     attention_mask=snake_case   ,     token_type_ids=snake_case )
											snake_case_           =					model(snake_case   ,     token_type_ids=snake_case )
											snake_case_           =					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 ):
											snake_case_           =					MegatronBertForMaskedLM(config=snake_case )
											model.to(snake_case )
											model.eval()
											snake_case_           =					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 ):
											snake_case_           =					MegatronBertForCausalLM(config=snake_case )
											model.to(snake_case )
											model.eval()
											snake_case_           =					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 ):
											snake_case_           =					MegatronBertForNextSentencePrediction(config=snake_case )
											model.to(snake_case )
											model.eval()
											snake_case_           =					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 ):
											snake_case_           =					MegatronBertForPreTraining(config=snake_case )
											model.to(snake_case )
											model.eval()
											snake_case_           =					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 ):
											snake_case_           =					MegatronBertForQuestionAnswering(config=snake_case )
											model.to(snake_case )
											model.eval()
											snake_case_           =					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 ):
											snake_case_           =					self.num_labels
											snake_case_           =					MegatronBertForSequenceClassification(snake_case )
											model.to(snake_case )
											model.eval()
											snake_case_           =					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 ):
											snake_case_           =					self.num_labels
											snake_case_           =					MegatronBertForTokenClassification(config=snake_case )
											model.to(snake_case )
											model.eval()
											snake_case_           =					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 ):
											snake_case_           =					self.num_choices
											snake_case_           =					MegatronBertForMultipleChoice(config=snake_case )
											model.to(snake_case )
											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(
											    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 ):
											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 				lowercase       ( lowercase_			,						lowercase_			,						unittest.TestCase    ):
				__SCREAMING_SNAKE_CASE      :				int     =    (
				    (
				        MegatronBertModel,
				        MegatronBertForMaskedLM,
				        MegatronBertForCausalLM,
				        MegatronBertForMultipleChoice,
				        MegatronBertForNextSentencePrediction,
				        MegatronBertForPreTraining,
				        MegatronBertForQuestionAnswering,
				        MegatronBertForSequenceClassification,
				        MegatronBertForTokenClassification,
				    )
				    if is_torch_available()
				    else ()
				)
				__SCREAMING_SNAKE_CASE      :				List[str]     =    (
				    {
				        '''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      :				Tuple     =    True
				# test_resize_embeddings = False
				__SCREAMING_SNAKE_CASE      :				Any     =    False




				def 					a							(      self   ,     snake_case   ,     snake_case   ,     snake_case=False ):
											snake_case_           =					super()._prepare_for_class(snake_case   ,     snake_case   ,     return_labels=snake_case )

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




				def 					a							(      self ):
											snake_case_           =					MegatronBertModelTester(self )
											snake_case_           =					ConfigTester(self   ,     config_class=snake_case   ,     hidden_size=37 )




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




				def 					a							(      self ):
											snake_case_           =					self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_megatron_bert_model(*snake_case )




				def 					a							(      self ):
											snake_case_           =					self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_megatron_bert_for_masked_lm(*snake_case )




				def 					a							(      self ):
											snake_case_           =					self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_megatron_bert_for_multiple_choice(*snake_case )




				def 					a							(      self ):
											snake_case_           =					self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_megatron_bert_for_next_sequence_prediction(*snake_case )




				def 					a							(      self ):
											snake_case_           =					self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_megatron_bert_for_pretraining(*snake_case )




				def 					a							(      self ):
											snake_case_           =					self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_megatron_bert_for_question_answering(*snake_case )




				def 					a							(      self ):
											snake_case_           =					self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_megatron_bert_for_sequence_classification(*snake_case )




				def 					a							(      self ):
											snake_case_           =					self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_megatron_bert_for_token_classification(*snake_case )





def       __lowerCamelCase			(				UpperCamelCase__   ):







							'''simple docstring'''






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


_UpperCAmelCase	:     Any		       = 1e-4







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




				@slow
				@unittest.skip('Model is not available.' )
				def 					a							(      self ):
											snake_case_           =					'nvidia/megatron-bert-uncased-345m'
											if "MYDIR" in os.environ:
																		snake_case_           =					os.path.join(os.environ['MYDIR']   ,     snake_case )
											snake_case_           =					MegatronBertModel.from_pretrained(snake_case )
											model.to(snake_case )
											model.half()
											snake_case_           =					_long_tensor([[101, 7110, 1005, 1056, 2023, 1_1333, 1_7413, 1029, 102]] )
											with torch.no_grad():
																		snake_case_           =					model(snake_case )[0]
											snake_case_           =					torch.Size((1, 9, 1024) )
											self.assertEqual(output.shape   ,     snake_case )

											snake_case_           =					[-0.60_40, -0.25_17, -0.10_25, 0.34_20, -0.67_58, -0.00_17, -0.10_89, -0.19_90, 0.57_28]
											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(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 )







_UpperCAmelCase	:     str		       = [0, 2, 4, 6, 8]
_UpperCAmelCase	:     Any		       = [1, 3, 5, 7, 9]





def       __lowerCamelCase			(				UpperCamelCase__     ,							UpperCamelCase__     ,							UpperCamelCase__     ,							UpperCamelCase__   ):







							'''simple docstring'''






							if remaining_length == 0:
														if digits[0] == 0 or digits[-1] == 0:
																					return 0

														for i in range(length // 2 - 1     ,							-1     ,							-1   ):
																					remainder += digits[i] + digits[length - i - 1]

																					if remainder % 2 == 0:
																												return 0

																					remainder //= 10

														return 1

							if remaining_length == 1:
														if remainder % 2 == 0:
																					return 0

														snake_case_           =					0
														for digit in range(10   ):
																					snake_case_           =					digit
																					result += reversible_numbers(
																					    0     ,							(remainder + 2 * digit) // 10     ,							UpperCamelCase__     ,							UpperCamelCase__   )
														return result

							snake_case_           =					0
							for digita in range(10   ):
														snake_case_           =					digita

														if (remainder + digita) % 2 == 0:
																					snake_case_           =					ODD_DIGITS
														else:
																					snake_case_           =					EVEN_DIGITS

														for digita in other_parity_digits:
																					snake_case_           =					digita
																					result += reversible_numbers(
																					    remaining_length - 2     ,							(remainder + digita + digita) // 10     ,							UpperCamelCase__     ,							UpperCamelCase__     ,							)
							return result





def       __lowerCamelCase			(				UpperCamelCase__ = 9   ):







							'''simple docstring'''






							snake_case_           =					0
							for length in range(1     ,							max_power + 1   ):
														result += reversible_numbers(UpperCamelCase__     ,							0     ,							[0] * length     ,							UpperCamelCase__   )
							return result


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




