





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




					def __init__(		self				,						lowerCAmelCase_				,						lowerCAmelCase_=13				,						lowerCAmelCase_=7				,						lowerCAmelCase_=True				,						lowerCAmelCase_=True				,						lowerCAmelCase_=True				,						lowerCAmelCase_=True				,						lowerCAmelCase_=99				,						lowerCAmelCase_=64				,						lowerCAmelCase_=32				,						lowerCAmelCase_=5				,						lowerCAmelCase_=4				,						lowerCAmelCase_=37				,						lowerCAmelCase_="gelu"				,						lowerCAmelCase_=0.1				,						lowerCAmelCase_=0.1				,						lowerCAmelCase_=512				,						lowerCAmelCase_=16				,						lowerCAmelCase_=2				,						lowerCAmelCase_=0.02				,						lowerCAmelCase_=3				,						lowerCAmelCase_=4				,						lowerCAmelCase_=None				,						)		->							str:
										_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  )		->							Union[str, Any]:
										_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  )		->							List[str]:
										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=_lowercase				,						initializer_range=self.initializer_range				,						)




					def 	lowerCAmelCase			(		self				,						lowerCAmelCase_				,						lowerCAmelCase_				,						lowerCAmelCase_				,						lowerCAmelCase_				,						lowerCAmelCase_				,						lowerCAmelCase_				,						lowerCAmelCase_  )		->							Optional[Any]:
										_snake_case									=				MegatronBertModel(config=_lowercase  )
										model.to(_lowercase  )
										model.eval()
										_snake_case									=				model(_lowercase				,						attention_mask=_lowercase				,						token_type_ids=_lowercase  )
										_snake_case									=				model(_lowercase				,						token_type_ids=_lowercase  )
										_snake_case									=				model(_lowercase  )

										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				,						lowerCAmelCase_				,						lowerCAmelCase_				,						lowerCAmelCase_				,						lowerCAmelCase_				,						lowerCAmelCase_				,						lowerCAmelCase_				,						lowerCAmelCase_  )		->							Optional[int]:
										_snake_case									=				MegatronBertForMaskedLM(config=_lowercase  )
										model.to(_lowercase  )
										model.eval()
										_snake_case									=				model(_lowercase				,						attention_mask=_lowercase				,						token_type_ids=_lowercase				,						labels=_lowercase  )
										self.parent.assertEqual(result.logits.shape				,						(self.batch_size, self.seq_length, self.vocab_size)  )




					def 	lowerCAmelCase			(		self				,						lowerCAmelCase_				,						lowerCAmelCase_				,						lowerCAmelCase_				,						lowerCAmelCase_				,						lowerCAmelCase_				,						lowerCAmelCase_				,						lowerCAmelCase_  )		->							int:
										_snake_case									=				MegatronBertForCausalLM(config=_lowercase  )
										model.to(_lowercase  )
										model.eval()
										_snake_case									=				model(_lowercase				,						attention_mask=_lowercase				,						token_type_ids=_lowercase				,						labels=_lowercase  )
										self.parent.assertEqual(result.logits.shape				,						(self.batch_size, self.seq_length, self.vocab_size)  )




					def 	lowerCAmelCase			(		self				,						lowerCAmelCase_				,						lowerCAmelCase_				,						lowerCAmelCase_				,						lowerCAmelCase_				,						lowerCAmelCase_				,						lowerCAmelCase_				,						lowerCAmelCase_  )		->							Union[str, Any]:
										_snake_case									=				MegatronBertForNextSentencePrediction(config=_lowercase  )
										model.to(_lowercase  )
										model.eval()
										_snake_case									=				model(
										    _lowercase				,						attention_mask=_lowercase				,						token_type_ids=_lowercase				,						labels=_lowercase				,						)
										self.parent.assertEqual(result.logits.shape				,						(self.batch_size, 2)  )




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




					def 	lowerCAmelCase			(		self				,						lowerCAmelCase_				,						lowerCAmelCase_				,						lowerCAmelCase_				,						lowerCAmelCase_				,						lowerCAmelCase_				,						lowerCAmelCase_				,						lowerCAmelCase_  )		->							Union[str, Any]:
										_snake_case									=				self.num_labels
										_snake_case									=				MegatronBertForTokenClassification(config=_lowercase  )
										model.to(_lowercase  )
										model.eval()
										_snake_case									=				model(_lowercase				,						attention_mask=_lowercase				,						token_type_ids=_lowercase				,						labels=_lowercase  )
										self.parent.assertEqual(result.logits.shape				,						(self.batch_size, self.seq_length, self.num_labels)  )




					def 	lowerCAmelCase			(		self				,						lowerCAmelCase_				,						lowerCAmelCase_				,						lowerCAmelCase_				,						lowerCAmelCase_				,						lowerCAmelCase_				,						lowerCAmelCase_				,						lowerCAmelCase_  )		->							Union[str, Any]:
										_snake_case									=				self.num_choices
										_snake_case									=				MegatronBertForMultipleChoice(config=_lowercase  )
										model.to(_lowercase  )
										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(
										    _lowercase				,						attention_mask=_lowercase				,						token_type_ids=_lowercase				,						labels=_lowercase				,						)
										self.parent.assertEqual(result.logits.shape				,						(self.batch_size, self.num_choices)  )






					def 	lowerCAmelCase			(		self  )		->							List[Any]:
										_snake_case									=				self.prepare_config_and_inputs()
										(
										    _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 			UpperCamelCase_		(  UpperCAmelCase_							,  UpperCAmelCase_							,  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 	lowerCAmelCase			(		self				,						lowerCAmelCase_				,						lowerCAmelCase_				,						lowerCAmelCase_=False  )		->							Dict:
										_snake_case									=				super()._prepare_for_class(_lowercase				,						_lowercase				,						return_labels=_lowercase  )

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




					def 	lowerCAmelCase			(		self  )		->							Optional[int]:
										_snake_case									=				MegatronBertModelTester(self  )
										_snake_case									=				ConfigTester(self				,						config_class=_lowercase				,						hidden_size=37  )




					def 	lowerCAmelCase			(		self  )		->							List[Any]:
										self.config_tester.run_common_tests()




					def 	lowerCAmelCase			(		self  )		->							Dict:
										_snake_case									=				self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_megatron_bert_model(*_lowercase  )




					def 	lowerCAmelCase			(		self  )		->							List[str]:
										_snake_case									=				self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_megatron_bert_for_masked_lm(*_lowercase  )




					def 	lowerCAmelCase			(		self  )		->							Tuple:
										_snake_case									=				self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_megatron_bert_for_multiple_choice(*_lowercase  )




					def 	lowerCAmelCase			(		self  )		->							Any:
										_snake_case									=				self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_megatron_bert_for_next_sequence_prediction(*_lowercase  )




					def 	lowerCAmelCase			(		self  )		->							Tuple:
										_snake_case									=				self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_megatron_bert_for_pretraining(*_lowercase  )




					def 	lowerCAmelCase			(		self  )		->							Optional[int]:
										_snake_case									=				self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_megatron_bert_for_question_answering(*_lowercase  )




					def 	lowerCAmelCase			(		self  )		->							Dict:
										_snake_case									=				self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_megatron_bert_for_sequence_classification(*_lowercase  )






					def 	lowerCAmelCase			(		self  )		->							List[Any]:
										_snake_case									=				self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_megatron_bert_for_token_classification(*_lowercase  )





def 						lowerCamelCase__     (     UpperCamelCase__			:       List[Any]						)			-> List[Any]:



					'''simple docstring'''




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


UpperCAmelCase_      							=     1E-4




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




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

										_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(_lowercase				,						_lowercase				,						_lowercase				,						_lowercase  )
																				self.assertTrue(math.isclose(_lowercase				,						_lowercase				,						rel_tol=_lowercase				,						abs_tol=_lowercase  )				,						msg=_lowercase  )










import unittest

from accelerate import debug_launcher
from accelerate.test_utils import require_cpu, test_ops, test_script
@require_cpu
class 			UpperCamelCase_		(  unittest.TestCase					):




					def 	lowerCAmelCase			(		self  )		->							List[Any]:
										debug_launcher(test_script.main  )






					def 	lowerCAmelCase			(		self  )		->							Optional[Any]:
										debug_launcher(test_ops.main  )



