



"""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    __lowercase     :







						"""simple docstring"""







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







						def 						__UpperCamelCase							(self       ):
											snake_case_					:		List[str] 					=					ids_tensor([self.batch_size, self.seq_length]   ,		self.vocab_size       )

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

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

											snake_case_					:		Dict 					=					None
											snake_case_					:		Tuple 					=					None
											snake_case_					:		Optional[Any] 					=					None
											if self.use_labels:
																snake_case_					:		str 					=					ids_tensor([self.batch_size]   ,		self.type_sequence_label_size       )
																snake_case_					:		str 					=					ids_tensor([self.batch_size, self.seq_length]   ,		self.num_labels       )
																snake_case_					:		int 					=					ids_tensor([self.batch_size]   ,		self.num_choices       )

											snake_case_					:		Optional[Any] 					=					self.get_config()

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







						def 						__UpperCamelCase							(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 						__UpperCamelCase							(self   ,		lowercase__   ,		lowercase__   ,		lowercase__   ,		lowercase__   ,		lowercase__   ,		lowercase__   ,		lowercase__       ):
											snake_case_					:		Dict 					=					MegatronBertModel(config=__snake_case       )
											model.to(__snake_case       )
											model.eval()
											snake_case_					:		Tuple 					=					model(__snake_case   ,		attention_mask=__snake_case   ,		token_type_ids=__snake_case       )
											snake_case_					:		str 					=					model(__snake_case   ,		token_type_ids=__snake_case       )
											snake_case_					:		Optional[int] 					=					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 						__UpperCamelCase							(self   ,		lowercase__   ,		lowercase__   ,		lowercase__   ,		lowercase__   ,		lowercase__   ,		lowercase__   ,		lowercase__       ):
											snake_case_					:		Dict 					=					MegatronBertForMaskedLM(config=__snake_case       )
											model.to(__snake_case       )
											model.eval()
											snake_case_					:		List[Any] 					=					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 						__UpperCamelCase							(self   ,		lowercase__   ,		lowercase__   ,		lowercase__   ,		lowercase__   ,		lowercase__   ,		lowercase__   ,		lowercase__       ):
											snake_case_					:		Any 					=					MegatronBertForCausalLM(config=__snake_case       )
											model.to(__snake_case       )
											model.eval()
											snake_case_					:		str 					=					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 						__UpperCamelCase							(self   ,		lowercase__   ,		lowercase__   ,		lowercase__   ,		lowercase__   ,		lowercase__   ,		lowercase__   ,		lowercase__       ):
											snake_case_					:		Dict 					=					MegatronBertForNextSentencePrediction(config=__snake_case       )
											model.to(__snake_case       )
											model.eval()
											snake_case_					:		int 					=					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 						__UpperCamelCase							(self   ,		lowercase__   ,		lowercase__   ,		lowercase__   ,		lowercase__   ,		lowercase__   ,		lowercase__   ,		lowercase__       ):
											snake_case_					:		Dict 					=					MegatronBertForPreTraining(config=__snake_case       )
											model.to(__snake_case       )
											model.eval()
											snake_case_					:		Optional[Any] 					=					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 						__UpperCamelCase							(self   ,		lowercase__   ,		lowercase__   ,		lowercase__   ,		lowercase__   ,		lowercase__   ,		lowercase__   ,		lowercase__       ):
											snake_case_					:		Dict 					=					MegatronBertForQuestionAnswering(config=__snake_case       )
											model.to(__snake_case       )
											model.eval()
											snake_case_					:		Tuple 					=					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 						__UpperCamelCase							(self   ,		lowercase__   ,		lowercase__   ,		lowercase__   ,		lowercase__   ,		lowercase__   ,		lowercase__   ,		lowercase__       ):
											snake_case_					:		int 					=					self.num_labels
											snake_case_					:		List[str] 					=					MegatronBertForSequenceClassification(__snake_case       )
											model.to(__snake_case       )
											model.eval()
											snake_case_					:		List[str] 					=					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 						__UpperCamelCase							(self   ,		lowercase__   ,		lowercase__   ,		lowercase__   ,		lowercase__   ,		lowercase__   ,		lowercase__   ,		lowercase__       ):
											snake_case_					:		Optional[Any] 					=					self.num_labels
											snake_case_					:		int 					=					MegatronBertForTokenClassification(config=__snake_case       )
											model.to(__snake_case       )
											model.eval()
											snake_case_					:		Any 					=					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 						__UpperCamelCase							(self   ,		lowercase__   ,		lowercase__   ,		lowercase__   ,		lowercase__   ,		lowercase__   ,		lowercase__   ,		lowercase__       ):
											snake_case_					:		Optional[Any] 					=					self.num_choices
											snake_case_					:		Tuple 					=					MegatronBertForMultipleChoice(config=__snake_case       )
											model.to(__snake_case       )
											model.eval()
											snake_case_					:		List[Any] 					=					input_ids.unsqueeze(1       ).expand(-1   ,		self.num_choices   ,		-1       ).contiguous()
											snake_case_					:		List[Any] 					=					token_type_ids.unsqueeze(1       ).expand(-1   ,		self.num_choices   ,		-1       ).contiguous()
											snake_case_					:		List[str] 					=					input_mask.unsqueeze(1       ).expand(-1   ,		self.num_choices   ,		-1       ).contiguous()
											snake_case_					:		Dict 					=					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 						__UpperCamelCase							(self       ):
											snake_case_					:		List[str] 					=					self.prepare_config_and_inputs()
											(
											    snake_case_    
											)					:		Optional[Any] 					=					config_and_inputs
											snake_case_					:		List[Any] 					=					{'''input_ids''': input_ids, '''token_type_ids''': token_type_ids, '''attention_mask''': input_mask}
											return config, inputs_dict







@require_torch
class    __lowercase     (					__SCREAMING_SNAKE_CASE					,    __SCREAMING_SNAKE_CASE					,    unittest.TestCase):







						"""simple docstring"""







						_A		:							Tuple							       =		(
						    (
						        MegatronBertModel,
						        MegatronBertForMaskedLM,
						        MegatronBertForCausalLM,
						        MegatronBertForMultipleChoice,
						        MegatronBertForNextSentencePrediction,
						        MegatronBertForPreTraining,
						        MegatronBertForQuestionAnswering,
						        MegatronBertForSequenceClassification,
						        MegatronBertForTokenClassification,
						    )
						    if is_torch_available()
						    else ()
						)
						_A		:							Tuple							       =		(
						    {
						        """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		:							List[str]							       =		True
						# test_resize_embeddings = False
						_A		:							Union[str, Any]							       =		False







						def 						__UpperCamelCase							(self   ,		lowercase__   ,		lowercase__   ,		lowercase__=False       ):
											snake_case_					:		Union[str, Any] 					=					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_					:		int 					=					torch.zeros(
																					    (self.model_tester.batch_size, self.model_tester.seq_length)   ,		dtype=torch.long   ,		device=__snake_case       )
																					snake_case_					:		str 					=					torch.zeros(
																					    self.model_tester.batch_size   ,		dtype=torch.long   ,		device=__snake_case       )
											return inputs_dict







						def 						__UpperCamelCase							(self       ):
											snake_case_					:		List[Any] 					=					MegatronBertModelTester(self       )
											snake_case_					:		Union[str, Any] 					=					ConfigTester(self   ,		config_class=__snake_case   ,		hidden_size=37       )







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







						def 						__UpperCamelCase							(self       ):
											snake_case_					:		int 					=					self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_megatron_bert_model(*__snake_case       )







						def 						__UpperCamelCase							(self       ):
											snake_case_					:		Optional[Any] 					=					self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_megatron_bert_for_masked_lm(*__snake_case       )







						def 						__UpperCamelCase							(self       ):
											snake_case_					:		Dict 					=					self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_megatron_bert_for_multiple_choice(*__snake_case       )







						def 						__UpperCamelCase							(self       ):
											snake_case_					:		List[Any] 					=					self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_megatron_bert_for_next_sequence_prediction(*__snake_case       )







						def 						__UpperCamelCase							(self       ):
											snake_case_					:		Tuple 					=					self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_megatron_bert_for_pretraining(*__snake_case       )







						def 						__UpperCamelCase							(self       ):
											snake_case_					:		Dict 					=					self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_megatron_bert_for_question_answering(*__snake_case       )







						def 						__UpperCamelCase							(self       ):
											snake_case_					:		Dict 					=					self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_megatron_bert_for_sequence_classification(*__snake_case       )







						def 						__UpperCamelCase							(self       ):
											snake_case_					:		List[Any] 					=					self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_megatron_bert_for_token_classification(*__snake_case       )







def 							SCREAMING_SNAKE_CASE__      (		SCREAMING_SNAKE_CASE__ :  Optional[Any]			):
					return torch.tensor(
					    a_				,    dtype=torch.long				,    device=a_				,    )


a_			   =					1E-4







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







						"""simple docstring"""







						@slow
						@unittest.skip("""Model is not available."""       )
						def 						__UpperCamelCase							(self       ):
											snake_case_					:		Union[str, Any] 					=					'''nvidia/megatron-bert-uncased-345m'''
											if "MYDIR" in os.environ:
																snake_case_					:		Any 					=					os.path.join(os.environ["""MYDIR"""]   ,		__snake_case       )
											snake_case_					:		Union[str, Any] 					=					MegatronBertModel.from_pretrained(__snake_case       )
											model.to(__snake_case       )
											model.half()
											snake_case_					:		List[str] 					=					_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_					:		Any 					=					model(__snake_case       )[0]
											snake_case_					:		Tuple 					=					torch.Size((1, 9, 10_24)       )
											self.assertEqual(output.shape   ,		__snake_case       )

											snake_case_					:		str 					=					[-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_					:		int 					=					output[0, ii, jj]
																					snake_case_					:		str 					=					expected[3 * ii + jj]
																					snake_case_					:		Union[str, Any] 					=					'''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       )









"""simple docstring"""

from copy import deepcopy







class    __lowercase     :







						"""simple docstring"""







						def __init__(self   ,		lowercase__ = None   ,		lowercase__ = None       ):

											if arr is None and size is not None:
																snake_case_					:		str 					=					size
																snake_case_					:		Optional[Any] 					=					[0] * size
											elif arr is not None:
																self.init(lowercase__       )
											else:
																raise ValueError("""Either arr or size must be specified"""       )







						def 						__UpperCamelCase							(self   ,		lowercase__       ):
											snake_case_					:		Optional[Any] 					=					len(lowercase__       )
											snake_case_					:		int 					=					deepcopy(lowercase__       )
											for i in range(1   ,		self.size       ):
																snake_case_					:		Optional[Any] 					=					self.next_(lowercase__       )
																if j < self.size:
																					self.tree[j] += self.tree[i]







						def 						__UpperCamelCase							(self       ):
											snake_case_					:		Dict 					=					self.tree[:]
											for i in range(self.size - 1   ,		0   ,		-1       ):
																snake_case_					:		Optional[int] 					=					self.next_(lowercase__       )
																if j < self.size:
																					arr[j] -= arr[i]
											return arr







						@staticmethod
						def 						__UpperCamelCase							(lowercase__       ):
											return index + (index & (-index))







						@staticmethod
						def 						__UpperCamelCase							(lowercase__       ):
											return index - (index & (-index))







						def 						__UpperCamelCase							(self   ,		lowercase__   ,		lowercase__       ):
											if index == 0:
																self.tree[0] += value
																return
											while index < self.size:
																self.tree[index] += value
																snake_case_					:		Tuple 					=					self.next_(lowercase__       )







						def 						__UpperCamelCase							(self   ,		lowercase__   ,		lowercase__       ):
											self.add(lowercase__   ,		value - self.get(lowercase__       )       )







						def 						__UpperCamelCase							(self   ,		lowercase__       ):
											if right == 0:
																return 0
											snake_case_					:		List[str] 					=					self.tree[0]
											right -= 1  # make right inclusive
											while right > 0:
																result += self.tree[right]
																snake_case_					:		Optional[int] 					=					self.prev(lowercase__       )
											return result







						def 						__UpperCamelCase							(self   ,		lowercase__   ,		lowercase__       ):
											return self.prefix(lowercase__       ) - self.prefix(lowercase__       )







						def 						__UpperCamelCase							(self   ,		lowercase__       ):
											return self.query(lowercase__   ,		index + 1       )







						def 						__UpperCamelCase							(self   ,		lowercase__       ):
											value -= self.tree[0]
											if value < 0:
																return -1

											snake_case_					:		Tuple 					=					1  # Largest power of 2 <= size
											while j * 2 < self.size:
																j *= 2

											snake_case_					:		Tuple 					=					0

											while j > 0:
																if i + j < self.size and self.tree[i + j] <= value:
																					value -= self.tree[i + j]
																					i += j
																j //= 2
											return i


if __name__ == "__main__":
			import doctest

			doctest.testmod()




