



'''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 	UpperCAmelCase	:


	'''simple docstring'''




	def __init__(   self							,  lowerCAmelCase_							,  lowerCAmelCase_=1_3							,  lowerCAmelCase_=7							,  lowerCAmelCase_=True							,  lowerCAmelCase_=True							,  lowerCAmelCase_=True							,  lowerCAmelCase_=True							,  lowerCAmelCase_=9_9							,  lowerCAmelCase_=6_4							,  lowerCAmelCase_=3_2							,  lowerCAmelCase_=5							,  lowerCAmelCase_=4							,  lowerCAmelCase_=3_7							,  lowerCAmelCase_="gelu"							,  lowerCAmelCase_=0.1							,  lowerCAmelCase_=0.1							,  lowerCAmelCase_=5_1_2							,  lowerCAmelCase_=1_6							,  lowerCAmelCase_=2							,  lowerCAmelCase_=0.0_2							,  lowerCAmelCase_=3							,  lowerCAmelCase_=4							,  lowerCAmelCase_=None							,  ) ->      Optional[int]:



						"""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       lowercase_  (   self) ->      Optional[Any]:



						"""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       lowercase_  (   self) ->      List[str]:



						"""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       lowercase_  (   self							,  lowerCAmelCase_							,  lowerCAmelCase_							,  lowerCAmelCase_							,  lowerCAmelCase_							,  lowerCAmelCase_							,  lowerCAmelCase_							,  lowerCAmelCase_) ->      Dict:



						"""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       lowercase_  (   self							,  lowerCAmelCase_							,  lowerCAmelCase_							,  lowerCAmelCase_							,  lowerCAmelCase_							,  lowerCAmelCase_							,  lowerCAmelCase_							,  lowerCAmelCase_) ->      List[Any]:



						"""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       lowercase_  (   self							,  lowerCAmelCase_							,  lowerCAmelCase_							,  lowerCAmelCase_							,  lowerCAmelCase_							,  lowerCAmelCase_							,  lowerCAmelCase_							,  lowerCAmelCase_) ->      Union[str, 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       lowercase_  (   self							,  lowerCAmelCase_							,  lowerCAmelCase_							,  lowerCAmelCase_							,  lowerCAmelCase_							,  lowerCAmelCase_							,  lowerCAmelCase_							,  lowerCAmelCase_) ->      int:



						"""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       lowercase_  (   self							,  lowerCAmelCase_							,  lowerCAmelCase_							,  lowerCAmelCase_							,  lowerCAmelCase_							,  lowerCAmelCase_							,  lowerCAmelCase_							,  lowerCAmelCase_) ->      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       lowercase_  (   self							,  lowerCAmelCase_							,  lowerCAmelCase_							,  lowerCAmelCase_							,  lowerCAmelCase_							,  lowerCAmelCase_							,  lowerCAmelCase_							,  lowerCAmelCase_) ->      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       lowercase_  (   self							,  lowerCAmelCase_							,  lowerCAmelCase_							,  lowerCAmelCase_							,  lowerCAmelCase_							,  lowerCAmelCase_							,  lowerCAmelCase_							,  lowerCAmelCase_) ->      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       lowercase_  (   self							,  lowerCAmelCase_							,  lowerCAmelCase_							,  lowerCAmelCase_							,  lowerCAmelCase_							,  lowerCAmelCase_							,  lowerCAmelCase_							,  lowerCAmelCase_) ->      Any:



						"""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       lowercase_  (   self							,  lowerCAmelCase_							,  lowerCAmelCase_							,  lowerCAmelCase_							,  lowerCAmelCase_							,  lowerCAmelCase_							,  lowerCAmelCase_							,  lowerCAmelCase_) ->      Tuple:



						"""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       lowercase_  (   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 	UpperCAmelCase	(		__a					,	__a					,	unittest.TestCase):


	'''simple docstring'''




	__magic_name__					: int  =      (
	    (
	        MegatronBertModel,
	        MegatronBertForMaskedLM,
	        MegatronBertForCausalLM,
	        MegatronBertForMultipleChoice,
	        MegatronBertForNextSentencePrediction,
	        MegatronBertForPreTraining,
	        MegatronBertForQuestionAnswering,
	        MegatronBertForSequenceClassification,
	        MegatronBertForTokenClassification,
	    )
	    if is_torch_available()
	    else ()
	)
	__magic_name__					: Dict  =      (
	    {
	        "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 {}
	)
	__magic_name__					: Optional[int]  =      True
	# test_resize_embeddings = False
	__magic_name__					: str  =      False




	def       lowercase_  (   self							,  lowerCAmelCase_							,  lowerCAmelCase_							,  lowerCAmelCase_=False) ->      str:



						"""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       lowercase_  (   self) ->      Tuple:



						"""simple docstring"""







						a_       		=MegatronBertModelTester(self)
						a_       		=ConfigTester(self							,  config_class=lowerCAmelCase_							,  hidden_size=3_7)




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



						"""simple docstring"""







						self.config_tester.run_common_tests()




	def       lowercase_  (   self) ->      Optional[int]:



						"""simple docstring"""







						a_       		=self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_megatron_bert_model(*lowerCAmelCase_)




	def       lowercase_  (   self) ->      int:



						"""simple docstring"""







						a_       		=self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_megatron_bert_for_masked_lm(*lowerCAmelCase_)




	def       lowercase_  (   self) ->      int:



						"""simple docstring"""







						a_       		=self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_megatron_bert_for_multiple_choice(*lowerCAmelCase_)




	def       lowercase_  (   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       lowercase_  (   self) ->      List[str]:



						"""simple docstring"""







						a_       		=self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_megatron_bert_for_pretraining(*lowerCAmelCase_)




	def       lowercase_  (   self) ->      Optional[Any]:



						"""simple docstring"""







						a_       		=self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_megatron_bert_for_question_answering(*lowerCAmelCase_)




	def       lowercase_  (   self) ->      str:



						"""simple docstring"""







						a_       		=self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_megatron_bert_for_sequence_classification(*lowerCAmelCase_)




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



						"""simple docstring"""







						a_       		=self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_megatron_bert_for_token_classification(*lowerCAmelCase_)






def  UpperCAmelCase_       (		lowercase__       ):


					'''simple docstring'''







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


lowercase							       =			1e-4







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


	'''simple docstring'''




	@slow
	@unittest.skip("Model is not available.")
	def       lowercase_  (   self) ->      Tuple:



						"""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([[1_0_1, 7_1_1_0, 1_0_0_5, 1_0_5_6, 2_0_2_3, 1_1_3_3_3, 1_7_4_1_3, 1_0_2_9, 1_0_2]])
						with torch.no_grad():
											a_       		=model(lowerCAmelCase_)[0]
						a_       		=torch.Size((1, 9, 1_0_2_4))
						self.assertEqual(output.shape							,  lowerCAmelCase_)

						a_       		=[-0.6_0_4_0, -0.2_5_1_7, -0.1_0_2_5, 0.3_4_2_0, -0.6_7_5_8, -0.0_0_1_7, -0.1_0_8_9, -0.1_9_9_0, 0.5_7_2_8]
						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_)







'''simple docstring'''


import os

# Precomputes a list of the 100 first triangular numbers
lowercase							       =			[int(0.5 * n * (n + 1)) for n in range(1, 101)]






def  UpperCAmelCase_       (		):


     '''simple docstring'''







     a_       		=os.path.dirname(os.path.realpath(lowercase__       )       )
     a_       		=os.path.join(lowercase__       ,	"words.txt"       )

     a_       		=""
     with open(lowercase__       ) as f:
          a_       		=f.readline()

     a_       		=[word.strip("\""       ) for word in words.strip("\r\n"       ).split(","       )]
     a_       		=[
         word
         for word in [sum(ord(lowercase__       ) - 6_4 for x in word       ) for word in words]
         if word in TRIANGULAR_NUMBERS
     ]
     return len(lowercase__       )


if __name__ == "__main__":
  print(solution())


