




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							:     Tuple       ,	__lowercase							:     Dict       ,	__lowercase							:     Optional[Any]=13       ,	__lowercase							:     List[Any]=7       ,	__lowercase							:     List[str]=True       ,	__lowercase							:     Optional[int]=True       ,	__lowercase							:     Any=True       ,	__lowercase							:     Dict=True       ,	__lowercase							:     Optional[Any]=99       ,	__lowercase							:     List[Any]=64       ,	__lowercase							:     List[Any]=32       ,	__lowercase							:     Optional[int]=5       ,	__lowercase							:     int=4       ,	__lowercase							:     Optional[int]=37       ,	__lowercase							:     int="gelu"       ,	__lowercase							:     Tuple=0.1       ,	__lowercase							:     Tuple=0.1       ,	__lowercase							:     List[Any]=5_12       ,	__lowercase							:     Union[str, Any]=16       ,	__lowercase							:     str=2       ,	__lowercase							:     List[Any]=0.02       ,	__lowercase							:     Union[str, Any]=3       ,	__lowercase							:     Union[str, Any]=4       ,	__lowercase							:     Tuple=None       ,	):






							"""simple docstring"""






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






							"""simple docstring"""






							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        snake_case__						(						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=_a       ,	initializer_range=self.initializer_range       ,	)






						def        snake_case__						(						self							:     Any       ,	__lowercase							:     Dict       ,	__lowercase							:     Optional[int]       ,	__lowercase							:     Optional[int]       ,	__lowercase							:     Optional[int]       ,	__lowercase							:     str       ,	__lowercase							:     Dict       ,	__lowercase							:     Optional[int]						):






							"""simple docstring"""






							snake_case_		    =     MegatronBertModel(config=_a						)
							model.to(_a						)
							model.eval()
							snake_case_		    =     model(_a       ,	attention_mask=_a       ,	token_type_ids=_a						)
							snake_case_		    =     model(_a       ,	token_type_ids=_a						)
							snake_case_		    =     model(_a						)

							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        snake_case__						(						self							:     Union[str, Any]       ,	__lowercase							:     List[str]       ,	__lowercase							:     str       ,	__lowercase							:     List[str]       ,	__lowercase							:     int       ,	__lowercase							:     List[str]       ,	__lowercase							:     Union[str, Any]       ,	__lowercase							:     Dict						):






							"""simple docstring"""






							snake_case_		    =     MegatronBertForMaskedLM(config=_a						)
							model.to(_a						)
							model.eval()
							snake_case_		    =     model(_a       ,	attention_mask=_a       ,	token_type_ids=_a       ,	labels=_a						)
							self.parent.assertEqual(result.logits.shape       ,	(self.batch_size, self.seq_length, self.vocab_size)						)






						def        snake_case__						(						self							:     Any       ,	__lowercase							:     List[str]       ,	__lowercase							:     Dict       ,	__lowercase							:     Dict       ,	__lowercase							:     int       ,	__lowercase							:     str       ,	__lowercase							:     int       ,	__lowercase							:     int						):






							"""simple docstring"""






							snake_case_		    =     MegatronBertForCausalLM(config=_a						)
							model.to(_a						)
							model.eval()
							snake_case_		    =     model(_a       ,	attention_mask=_a       ,	token_type_ids=_a       ,	labels=_a						)
							self.parent.assertEqual(result.logits.shape       ,	(self.batch_size, self.seq_length, self.vocab_size)						)






						def        snake_case__						(						self							:     Optional[int]       ,	__lowercase							:     Optional[int]       ,	__lowercase							:     Optional[int]       ,	__lowercase							:     Any       ,	__lowercase							:     Union[str, Any]       ,	__lowercase							:     Tuple       ,	__lowercase							:     str       ,	__lowercase							:     Optional[int]						):






							"""simple docstring"""






							snake_case_		    =     MegatronBertForNextSentencePrediction(config=_a						)
							model.to(_a						)
							model.eval()
							snake_case_		    =     model(
							    _a       ,	attention_mask=_a       ,	token_type_ids=_a       ,	labels=_a       ,	)
							self.parent.assertEqual(result.logits.shape       ,	(self.batch_size, 2)						)






						def        snake_case__						(						self							:     List[str]       ,	__lowercase							:     str       ,	__lowercase							:     Union[str, Any]       ,	__lowercase							:     Optional[int]       ,	__lowercase							:     Optional[int]       ,	__lowercase							:     str       ,	__lowercase							:     Tuple       ,	__lowercase							:     Union[str, Any]						):






							"""simple docstring"""






							snake_case_		    =     MegatronBertForPreTraining(config=_a						)
							model.to(_a						)
							model.eval()
							snake_case_		    =     model(
							    _a       ,	attention_mask=_a       ,	token_type_ids=_a       ,	labels=_a       ,	next_sentence_label=_a       ,	)
							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        snake_case__						(						self							:     int       ,	__lowercase							:     List[Any]       ,	__lowercase							:     Tuple       ,	__lowercase							:     Optional[int]       ,	__lowercase							:     List[str]       ,	__lowercase							:     Union[str, Any]       ,	__lowercase							:     Any       ,	__lowercase							:     Any						):






							"""simple docstring"""






							snake_case_		    =     MegatronBertForQuestionAnswering(config=_a						)
							model.to(_a						)
							model.eval()
							snake_case_		    =     model(
							    _a       ,	attention_mask=_a       ,	token_type_ids=_a       ,	start_positions=_a       ,	end_positions=_a       ,	)
							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        snake_case__						(						self							:     Optional[int]       ,	__lowercase							:     int       ,	__lowercase							:     List[Any]       ,	__lowercase							:     Union[str, Any]       ,	__lowercase							:     Union[str, Any]       ,	__lowercase							:     Tuple       ,	__lowercase							:     int       ,	__lowercase							:     Dict						):






							"""simple docstring"""






							snake_case_		    =     self.num_labels
							snake_case_		    =     MegatronBertForSequenceClassification(_a						)
							model.to(_a						)
							model.eval()
							snake_case_		    =     model(_a       ,	attention_mask=_a       ,	token_type_ids=_a       ,	labels=_a						)
							self.parent.assertEqual(result.logits.shape       ,	(self.batch_size, self.num_labels)						)






						def        snake_case__						(						self							:     str       ,	__lowercase							:     List[str]       ,	__lowercase							:     List[str]       ,	__lowercase							:     Dict       ,	__lowercase							:     List[Any]       ,	__lowercase							:     Optional[int]       ,	__lowercase							:     int       ,	__lowercase							:     Optional[Any]						):






							"""simple docstring"""






							snake_case_		    =     self.num_labels
							snake_case_		    =     MegatronBertForTokenClassification(config=_a						)
							model.to(_a						)
							model.eval()
							snake_case_		    =     model(_a       ,	attention_mask=_a       ,	token_type_ids=_a       ,	labels=_a						)
							self.parent.assertEqual(result.logits.shape       ,	(self.batch_size, self.seq_length, self.num_labels)						)






						def        snake_case__						(						self							:     Dict       ,	__lowercase							:     Optional[int]       ,	__lowercase							:     Union[str, Any]       ,	__lowercase							:     Optional[int]       ,	__lowercase							:     Tuple       ,	__lowercase							:     Any       ,	__lowercase							:     Tuple       ,	__lowercase							:     str						):






							"""simple docstring"""






							snake_case_		    =     self.num_choices
							snake_case_		    =     MegatronBertForMultipleChoice(config=_a						)
							model.to(_a						)
							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(
							    _a       ,	attention_mask=_a       ,	token_type_ids=_a       ,	labels=_a       ,	)
							self.parent.assertEqual(result.logits.shape       ,	(self.batch_size, self.num_choices)						)






						def        snake_case__						(						self							:     List[Any]						):






							"""simple docstring"""






							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   (				SCREAMING_SNAKE_CASE_       ,   SCREAMING_SNAKE_CASE_       ,   unittest.TestCase			):







						'''simple docstring'''





						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        snake_case__						(						self							:     List[str]       ,	__lowercase							:     int       ,	__lowercase							:     Union[str, Any]       ,	__lowercase							:     Optional[Any]=False						):






							"""simple docstring"""






							snake_case_		    =     super()._prepare_for_class(_a       ,	_a       ,	return_labels=_a						)

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






						def        snake_case__						(						self							:     List[str]						):






							"""simple docstring"""






							snake_case_		    =     MegatronBertModelTester(self						)
							snake_case_		    =     ConfigTester(self       ,	config_class=_a       ,	hidden_size=37						)






						def        snake_case__						(						self							:     Union[str, Any]						):






							"""simple docstring"""






							self.config_tester.run_common_tests()






						def        snake_case__						(						self							:     Dict						):






							"""simple docstring"""






							snake_case_		    =     self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_megatron_bert_model(*_a						)






						def        snake_case__						(						self							:     int						):






							"""simple docstring"""






							snake_case_		    =     self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_megatron_bert_for_masked_lm(*_a						)






						def        snake_case__						(						self							:     Optional[Any]						):






							"""simple docstring"""






							snake_case_		    =     self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_megatron_bert_for_multiple_choice(*_a						)






						def        snake_case__						(						self							:     Optional[Any]						):






							"""simple docstring"""






							snake_case_		    =     self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_megatron_bert_for_next_sequence_prediction(*_a						)






						def        snake_case__						(						self							:     Optional[Any]						):






							"""simple docstring"""






							snake_case_		    =     self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_megatron_bert_for_pretraining(*_a						)






						def        snake_case__						(						self							:     Dict						):






							"""simple docstring"""






							snake_case_		    =     self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_megatron_bert_for_question_answering(*_a						)






						def        snake_case__						(						self							:     Union[str, Any]						):






							"""simple docstring"""






							snake_case_		    =     self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_megatron_bert_for_sequence_classification(*_a						)






						def        snake_case__						(						self							:     Tuple						):






							"""simple docstring"""






							snake_case_		    =     self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_megatron_bert_for_token_classification(*_a						)







def 						lowerCamelCase__    (		_A							):







	'''simple docstring'''





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


lowercase__				:     int 		=						1e-4






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







						'''simple docstring'''





						@slow
						@unittest.skip("Model is not available."						)
						def        snake_case__						(						self							:     Union[str, Any]						):






							"""simple docstring"""






							snake_case_		    =     'nvidia/megatron-bert-uncased-345m'
							if "MYDIR" in os.environ:
								snake_case_		    =     os.path.join(os.environ["MYDIR"]       ,	_a						)
							snake_case_		    =     MegatronBertModel.from_pretrained(_a						)
							model.to(_a						)
							model.half()
							snake_case_		    =     _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_		    =     model(_a						)[0]
							snake_case_		    =     torch.Size((1, 9, 10_24)						)
							self.assertEqual(output.shape       ,	_a						)

							snake_case_		    =     [-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_		    =     output[0, ii, jj]
									snake_case_		    =     expected[3 * ii + jj]
									snake_case_		    =     'ii={} jj={} a={} b={}'.format(_a       ,	_a       ,	_a       ,	_a						)
									self.assertTrue(math.isclose(_a       ,	_a       ,	rel_tol=_a       ,	abs_tol=_a						)       ,	msg=_a						)





'''simple docstring'''





from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_tokenizers_available, is_torch_available


_UpperCamelCase											=						{
    'configuration_biogpt': ['BIOGPT_PRETRAINED_CONFIG_ARCHIVE_MAP', 'BioGptConfig'],
    'tokenization_biogpt': ['BioGptTokenizer'],
}

try:
   if not is_torch_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   _UpperCamelCase											=						[
       'BIOGPT_PRETRAINED_MODEL_ARCHIVE_LIST',
       'BioGptForCausalLM',
       'BioGptForTokenClassification',
       'BioGptForSequenceClassification',
       'BioGptModel',
       'BioGptPreTrainedModel',
   ]


if TYPE_CHECKING:
   from .configuration_biogpt import BIOGPT_PRETRAINED_CONFIG_ARCHIVE_MAP, BioGptConfig
   from .tokenization_biogpt import BioGptTokenizer

   try:
      if not is_torch_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .modeling_biogpt import (
          BIOGPT_PRETRAINED_MODEL_ARCHIVE_LIST,
          BioGptForCausalLM,
          BioGptForSequenceClassification,
          BioGptForTokenClassification,
          BioGptModel,
          BioGptPreTrainedModel,
      )


else:
   import sys

   _UpperCamelCase											=						_LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)






