'''simple docstring'''







import json
import os
import unittest

from transformers.models.biogpt.tokenization_biogpt import VOCAB_FILES_NAMES, BioGptTokenizer
from transformers.testing_utils import slow

from ...test_tokenization_common import TokenizerTesterMixin




class _SCREAMING_SNAKE_CASE(							_SCREAMING_SNAKE_CASE						,						unittest.TestCase					):
						A_			:							Dict     		= BioGptTokenizer
						A_			:							Optional[Any]     		= False




						def      __lowerCamelCase       (    self :       Union[str, Any]      )     ->				Union[str, Any]:
													super().setUp()

													# Adapted from Sennrich et al. 2015 and https://github.com/rsennrich/subword-nmt
													SCREAMING_SNAKE_CASE__     :Dict						    = [
													    'l',
													    'o',
													    'w',
													    'e',
													    'r',
													    's',
													    't',
													    'i',
													    'd',
													    'n',
													    'w</w>',
													    'r</w>',
													    't</w>',
													    'lo',
													    'low',
													    'er</w>',
													    'low</w>',
													    'lowest</w>',
													    'newer</w>',
													    'wider</w>',
													    '<unk>',
													]
													SCREAMING_SNAKE_CASE__     :str						    = dict(zip(UpperCamelCase_							,							range(len(UpperCamelCase_      )      )      )      )
													SCREAMING_SNAKE_CASE__     :str						    = ['l o 123', 'lo w 1456', 'e r</w> 1789', '']

													SCREAMING_SNAKE_CASE__     :Tuple						    = os.path.join(self.tmpdirname							,							VOCAB_FILES_NAMES['vocab_file']      )
													SCREAMING_SNAKE_CASE__     :str						    = os.path.join(self.tmpdirname							,							VOCAB_FILES_NAMES['merges_file']      )
													with open(self.vocab_file							,							'w'      ) as fp:
																				fp.write(json.dumps(UpperCamelCase_      )      )
													with open(self.merges_file							,							'w'      ) as fp:
																				fp.write('\n'.join(UpperCamelCase_      )      )




						def      __lowerCamelCase       (    self :       Tuple							,							UpperCamelCase_ :       List[str]      )     ->				List[str]:
													SCREAMING_SNAKE_CASE__     :Union[str, Any]						    = 'lower newer'
													SCREAMING_SNAKE_CASE__     :Union[str, Any]						    = 'lower newer'
													return input_text, output_text




						def      __lowerCamelCase       (    self :       str      )     ->				Optional[int]:
													SCREAMING_SNAKE_CASE__     :List[str]						    = BioGptTokenizer(self.vocab_file							,							self.merges_file      )

													SCREAMING_SNAKE_CASE__     :Any						    = 'lower'
													SCREAMING_SNAKE_CASE__     :Optional[Any]						    = ['low', 'er</w>']
													SCREAMING_SNAKE_CASE__     :Tuple						    = tokenizer.tokenize(UpperCamelCase_      )
													self.assertListEqual(UpperCamelCase_							,							UpperCamelCase_      )

													SCREAMING_SNAKE_CASE__     :Union[str, Any]						    = tokens + ['<unk>']
													SCREAMING_SNAKE_CASE__     :Optional[Any]						    = [14, 15, 20]
													self.assertListEqual(tokenizer.convert_tokens_to_ids(UpperCamelCase_      )							,							UpperCamelCase_      )




						@slow
						def      __lowerCamelCase       (    self :       Union[str, Any]      )     ->				Optional[int]:
													SCREAMING_SNAKE_CASE__     :List[Any]						    = BioGptTokenizer.from_pretrained('microsoft/biogpt'      )

													SCREAMING_SNAKE_CASE__     :Union[str, Any]						    = tokenizer.encode('sequence builders'							,							add_special_tokens=UpperCamelCase_      )
													SCREAMING_SNAKE_CASE__     :List[str]						    = tokenizer.encode('multi-sequence build'							,							add_special_tokens=UpperCamelCase_      )

													SCREAMING_SNAKE_CASE__     :List[str]						    = tokenizer.build_inputs_with_special_tokens(UpperCamelCase_      )
													SCREAMING_SNAKE_CASE__     :str						    = tokenizer.build_inputs_with_special_tokens(UpperCamelCase_							,							UpperCamelCase_      )

													self.assertTrue(encoded_sentence == [2] + text      )
													self.assertTrue(encoded_pair == [2] + text + [2] + text_a      )







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




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




						def      __lowerCamelCase       (    self :       Any      )     ->				Optional[Any]:
													SCREAMING_SNAKE_CASE__     :Optional[int]						    = ids_tensor([self.batch_size, self.seq_length]							,							self.vocab_size      )

													SCREAMING_SNAKE_CASE__     :List[Any]						    = None
													if self.use_input_mask:
																				SCREAMING_SNAKE_CASE__     :str						    = random_attention_mask([self.batch_size, self.seq_length]      )

													SCREAMING_SNAKE_CASE__     :Any						    = None
													if self.use_token_type_ids:
																				SCREAMING_SNAKE_CASE__     :List[Any]						    = ids_tensor([self.batch_size, self.seq_length]							,							self.type_vocab_size      )

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

													SCREAMING_SNAKE_CASE__     :str						    = self.get_config()

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




						def      __lowerCamelCase       (    self :       List[str]      )     ->				Any:
													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=UpperCamelCase_							,							initializer_range=self.initializer_range							,							)




						def      __lowerCamelCase       (    self :       Optional[int]							,							UpperCamelCase_ :       List[str]							,							UpperCamelCase_ :       List[str]							,							UpperCamelCase_ :       str							,							UpperCamelCase_ :       Optional[int]							,							UpperCamelCase_ :       List[str]							,							UpperCamelCase_ :       Any							,							UpperCamelCase_ :       int      )     ->				int:
													SCREAMING_SNAKE_CASE__     :Optional[int]						    = MegatronBertModel(config=UpperCamelCase_      )
													model.to(UpperCamelCase_      )
													model.eval()
													SCREAMING_SNAKE_CASE__     :int						    = model(UpperCamelCase_							,							attention_mask=UpperCamelCase_							,							token_type_ids=UpperCamelCase_      )
													SCREAMING_SNAKE_CASE__     :str						    = model(UpperCamelCase_							,							token_type_ids=UpperCamelCase_      )
													SCREAMING_SNAKE_CASE__     :List[str]						    = model(UpperCamelCase_      )

													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 :       str							,							UpperCamelCase_ :       Dict							,							UpperCamelCase_ :       Optional[int]							,							UpperCamelCase_ :       Optional[Any]							,							UpperCamelCase_ :       str							,							UpperCamelCase_ :       List[str]							,							UpperCamelCase_ :       int							,							UpperCamelCase_ :       Union[str, Any]      )     ->				List[str]:
													SCREAMING_SNAKE_CASE__     :str						    = MegatronBertForMaskedLM(config=UpperCamelCase_      )
													model.to(UpperCamelCase_      )
													model.eval()
													SCREAMING_SNAKE_CASE__     :Optional[int]						    = model(UpperCamelCase_							,							attention_mask=UpperCamelCase_							,							token_type_ids=UpperCamelCase_							,							labels=UpperCamelCase_      )
													self.parent.assertEqual(result.logits.shape							,							(self.batch_size, self.seq_length, self.vocab_size)      )




						def      __lowerCamelCase       (    self :       Any							,							UpperCamelCase_ :       List[Any]							,							UpperCamelCase_ :       Optional[int]							,							UpperCamelCase_ :       Any							,							UpperCamelCase_ :       Optional[int]							,							UpperCamelCase_ :       int							,							UpperCamelCase_ :       int							,							UpperCamelCase_ :       int      )     ->				Dict:
													SCREAMING_SNAKE_CASE__     :Optional[int]						    = MegatronBertForCausalLM(config=UpperCamelCase_      )
													model.to(UpperCamelCase_      )
													model.eval()
													SCREAMING_SNAKE_CASE__     :List[str]						    = model(UpperCamelCase_							,							attention_mask=UpperCamelCase_							,							token_type_ids=UpperCamelCase_							,							labels=UpperCamelCase_      )
													self.parent.assertEqual(result.logits.shape							,							(self.batch_size, self.seq_length, self.vocab_size)      )




						def      __lowerCamelCase       (    self :       List[str]							,							UpperCamelCase_ :       Dict							,							UpperCamelCase_ :       Optional[int]							,							UpperCamelCase_ :       Dict							,							UpperCamelCase_ :       Optional[Any]							,							UpperCamelCase_ :       List[str]							,							UpperCamelCase_ :       Union[str, Any]							,							UpperCamelCase_ :       Optional[int]      )     ->				List[Any]:
													SCREAMING_SNAKE_CASE__     :Optional[int]						    = MegatronBertForNextSentencePrediction(config=UpperCamelCase_      )
													model.to(UpperCamelCase_      )
													model.eval()
													SCREAMING_SNAKE_CASE__     :List[str]						    = model(
													    UpperCamelCase_							,							attention_mask=UpperCamelCase_							,							token_type_ids=UpperCamelCase_							,							labels=UpperCamelCase_							,							)
													self.parent.assertEqual(result.logits.shape							,							(self.batch_size, 2)      )




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




						def      __lowerCamelCase       (    self :       Any							,							UpperCamelCase_ :       int							,							UpperCamelCase_ :       Union[str, Any]							,							UpperCamelCase_ :       List[str]							,							UpperCamelCase_ :       Optional[int]							,							UpperCamelCase_ :       Union[str, Any]							,							UpperCamelCase_ :       Union[str, Any]							,							UpperCamelCase_ :       Optional[Any]      )     ->				Any:
													SCREAMING_SNAKE_CASE__     :Dict						    = self.num_labels
													SCREAMING_SNAKE_CASE__     :Any						    = MegatronBertForTokenClassification(config=UpperCamelCase_      )
													model.to(UpperCamelCase_      )
													model.eval()
													SCREAMING_SNAKE_CASE__     :Optional[Any]						    = model(UpperCamelCase_							,							attention_mask=UpperCamelCase_							,							token_type_ids=UpperCamelCase_							,							labels=UpperCamelCase_      )
													self.parent.assertEqual(result.logits.shape							,							(self.batch_size, self.seq_length, self.num_labels)      )




						def      __lowerCamelCase       (    self :       str							,							UpperCamelCase_ :       Any							,							UpperCamelCase_ :       List[Any]							,							UpperCamelCase_ :       Union[str, Any]							,							UpperCamelCase_ :       int							,							UpperCamelCase_ :       str							,							UpperCamelCase_ :       Any							,							UpperCamelCase_ :       List[Any]      )     ->				Dict:
													SCREAMING_SNAKE_CASE__     :Optional[Any]						    = self.num_choices
													SCREAMING_SNAKE_CASE__     :Tuple						    = MegatronBertForMultipleChoice(config=UpperCamelCase_      )
													model.to(UpperCamelCase_      )
													model.eval()
													SCREAMING_SNAKE_CASE__     :Any						    = input_ids.unsqueeze(1      ).expand(-1							,							self.num_choices							,							-1      ).contiguous()
													SCREAMING_SNAKE_CASE__     :List[Any]						    = token_type_ids.unsqueeze(1      ).expand(-1							,							self.num_choices							,							-1      ).contiguous()
													SCREAMING_SNAKE_CASE__     :Dict						    = input_mask.unsqueeze(1      ).expand(-1							,							self.num_choices							,							-1      ).contiguous()
													SCREAMING_SNAKE_CASE__     :Union[str, Any]						    = model(
													    UpperCamelCase_							,							attention_mask=UpperCamelCase_							,							token_type_ids=UpperCamelCase_							,							labels=UpperCamelCase_							,							)
													self.parent.assertEqual(result.logits.shape							,							(self.batch_size, self.num_choices)      )




						def      __lowerCamelCase       (    self :       Any      )     ->				int:
													SCREAMING_SNAKE_CASE__     :Dict						    = self.prepare_config_and_inputs()
													(
													    (
													    SCREAMING_SNAKE_CASE__
													)    ,      (
													    SCREAMING_SNAKE_CASE__
													)    ,      (
													    SCREAMING_SNAKE_CASE__
													)    ,      (
													    SCREAMING_SNAKE_CASE__
													)    ,      (
													    SCREAMING_SNAKE_CASE__
													)    ,      (
													    SCREAMING_SNAKE_CASE__
													)    ,      (
													    SCREAMING_SNAKE_CASE__
													)    ,      
													)     :Union[str, Any]						    = config_and_inputs
													SCREAMING_SNAKE_CASE__     :Tuple						    = {'input_ids': input_ids, 'token_type_ids': token_type_ids, 'attention_mask': input_mask}
													return config, inputs_dict




@require_torch
class _SCREAMING_SNAKE_CASE(							_SCREAMING_SNAKE_CASE						,						_SCREAMING_SNAKE_CASE						,						unittest.TestCase					):
						A_			:							List[Any]     		= (
						    (
						        MegatronBertModel,
						        MegatronBertForMaskedLM,
						        MegatronBertForCausalLM,
						        MegatronBertForMultipleChoice,
						        MegatronBertForNextSentencePrediction,
						        MegatronBertForPreTraining,
						        MegatronBertForQuestionAnswering,
						        MegatronBertForSequenceClassification,
						        MegatronBertForTokenClassification,
						    )
						    if is_torch_available()
						    else ()
						)
						A_			:							List[Any]     		= (
						    {
						        '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_			:							int     		= True
						# test_resize_embeddings = False
						A_			:							Dict     		= False




						def      __lowerCamelCase       (    self :       Union[str, Any]							,							UpperCamelCase_ :       int							,							UpperCamelCase_ :       int							,							UpperCamelCase_ :       str=False      )     ->				List[Any]:
													SCREAMING_SNAKE_CASE__     :Optional[Any]						    = super()._prepare_for_class(UpperCamelCase_							,							UpperCamelCase_							,							return_labels=UpperCamelCase_      )

													if return_labels:
																				if model_class in get_values(UpperCamelCase_      ):
																											SCREAMING_SNAKE_CASE__     :str						    = torch.zeros(
																											    (self.model_tester.batch_size, self.model_tester.seq_length)							,							dtype=torch.long							,							device=UpperCamelCase_      )
																											SCREAMING_SNAKE_CASE__     :List[str]						    = torch.zeros(
																											    self.model_tester.batch_size							,							dtype=torch.long							,							device=UpperCamelCase_      )
													return inputs_dict




						def      __lowerCamelCase       (    self :       Union[str, Any]      )     ->				Any:
													SCREAMING_SNAKE_CASE__     :Dict						    = MegatronBertModelTester(self      )
													SCREAMING_SNAKE_CASE__     :Union[str, Any]						    = ConfigTester(self							,							config_class=UpperCamelCase_							,							hidden_size=37      )




						def      __lowerCamelCase       (    self :       List[str]      )     ->				Optional[int]:
													self.config_tester.run_common_tests()




						def      __lowerCamelCase       (    self :       List[str]      )     ->				List[str]:
													SCREAMING_SNAKE_CASE__     :Dict						    = self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_megatron_bert_model(*UpperCamelCase_      )




						def      __lowerCamelCase       (    self :       Any      )     ->				Tuple:
													SCREAMING_SNAKE_CASE__     :List[Any]						    = self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_megatron_bert_for_masked_lm(*UpperCamelCase_      )




						def      __lowerCamelCase       (    self :       Union[str, Any]      )     ->				str:
													SCREAMING_SNAKE_CASE__     :Any						    = self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_megatron_bert_for_multiple_choice(*UpperCamelCase_      )




						def      __lowerCamelCase       (    self :       str      )     ->				Dict:
													SCREAMING_SNAKE_CASE__     :Dict						    = self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_megatron_bert_for_next_sequence_prediction(*UpperCamelCase_      )




						def      __lowerCamelCase       (    self :       int      )     ->				Any:
													SCREAMING_SNAKE_CASE__     :Dict						    = self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_megatron_bert_for_pretraining(*UpperCamelCase_      )




						def      __lowerCamelCase       (    self :       Dict      )     ->				Union[str, Any]:
													SCREAMING_SNAKE_CASE__     :Union[str, Any]						    = self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_megatron_bert_for_question_answering(*UpperCamelCase_      )




						def      __lowerCamelCase       (    self :       Any      )     ->				Union[str, Any]:
													SCREAMING_SNAKE_CASE__     :Union[str, Any]						    = self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_megatron_bert_for_sequence_classification(*UpperCamelCase_      )




						def      __lowerCamelCase       (    self :       List[Any]      )     ->				Dict:
													SCREAMING_SNAKE_CASE__     :Dict						    = self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_megatron_bert_for_token_classification(*UpperCamelCase_      )


def    lowerCamelCase       (			UpperCAmelCase__				:  Dict		)			->	Any:


							'''simple docstring'''






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


UpperCamelCase_      	=	1e-4





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




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

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






