



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




         """simple docstring"""

         lowercase__  =	parent
         lowercase__  =	batch_size
         lowercase__  =	seq_length
         lowercase__  =	is_training
         lowercase__  =	use_input_mask
         lowercase__  =	use_token_type_ids
         lowercase__  =	use_labels
         lowercase__  =	vocab_size
         lowercase__  =	hidden_size
         lowercase__  =	embedding_size
         lowercase__  =	num_hidden_layers
         lowercase__  =	num_attention_heads
         lowercase__  =	intermediate_size
         lowercase__  =	hidden_act
         lowercase__  =	hidden_dropout_prob
         lowercase__  =	attention_probs_dropout_prob
         lowercase__  =	max_position_embeddings
         lowercase__  =	type_vocab_size
         lowercase__  =	type_sequence_label_size
         lowercase__  =	initializer_range
         lowercase__  =	num_labels
         lowercase__  =	num_choices
         lowercase__  =	scope


   def       UpperCAmelCase		(    self	:      Tuple)						->					Tuple:




         """simple docstring"""

         lowercase__  =	ids_tensor([self.batch_size, self.seq_length]      ,					self.vocab_size)

         lowercase__  =	None
         if self.use_input_mask:
               lowercase__  =	random_attention_mask([self.batch_size, self.seq_length])

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

         lowercase__  =	None
         lowercase__  =	None
         lowercase__  =	None
         if self.use_labels:
               lowercase__  =	ids_tensor([self.batch_size]      ,					self.type_sequence_label_size)
               lowercase__  =	ids_tensor([self.batch_size, self.seq_length]      ,					self.num_labels)
               lowercase__  =	ids_tensor([self.batch_size]      ,					self.num_choices)

         lowercase__  =	self.get_config()

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


   def       UpperCAmelCase		(    self	:      List[Any])						->					Dict:




         """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       UpperCAmelCase		(    self	:      Optional[int]      ,					lowerCAmelCase	:      List[Any]      ,					lowerCAmelCase	:      List[str]      ,					lowerCAmelCase	:      Optional[Any]      ,					lowerCAmelCase	:      Optional[Any]      ,					lowerCAmelCase	:      Tuple      ,					lowerCAmelCase	:      Any      ,					lowerCAmelCase	:      Dict)						->					Optional[int]:




         """simple docstring"""

         lowercase__  =	MegatronBertModel(config=lowerCAmelCase)
         model.to(lowerCAmelCase)
         model.eval()
         lowercase__  =	model(lowerCAmelCase      ,					attention_mask=lowerCAmelCase      ,					token_type_ids=lowerCAmelCase)
         lowercase__  =	model(lowerCAmelCase      ,					token_type_ids=lowerCAmelCase)
         lowercase__  =	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       UpperCAmelCase		(    self	:      List[str]      ,					lowerCAmelCase	:      str      ,					lowerCAmelCase	:      str      ,					lowerCAmelCase	:      Dict      ,					lowerCAmelCase	:      Dict      ,					lowerCAmelCase	:      Tuple      ,					lowerCAmelCase	:      Optional[Any]      ,					lowerCAmelCase	:      Dict)						->					List[str]:




         """simple docstring"""

         lowercase__  =	MegatronBertForMaskedLM(config=lowerCAmelCase)
         model.to(lowerCAmelCase)
         model.eval()
         lowercase__  =	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       UpperCAmelCase		(    self	:      Dict      ,					lowerCAmelCase	:      List[str]      ,					lowerCAmelCase	:      Optional[int]      ,					lowerCAmelCase	:      Optional[int]      ,					lowerCAmelCase	:      int      ,					lowerCAmelCase	:      int      ,					lowerCAmelCase	:      Optional[int]      ,					lowerCAmelCase	:      Any)						->					Tuple:




         """simple docstring"""

         lowercase__  =	MegatronBertForCausalLM(config=lowerCAmelCase)
         model.to(lowerCAmelCase)
         model.eval()
         lowercase__  =	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       UpperCAmelCase		(    self	:      str      ,					lowerCAmelCase	:      Optional[Any]      ,					lowerCAmelCase	:      Dict      ,					lowerCAmelCase	:      Any      ,					lowerCAmelCase	:      Optional[int]      ,					lowerCAmelCase	:      List[str]      ,					lowerCAmelCase	:      List[str]      ,					lowerCAmelCase	:      Optional[Any])						->					Union[str, Any]:




         """simple docstring"""

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


   def       UpperCAmelCase		(    self	:      str      ,					lowerCAmelCase	:      Optional[Any]      ,					lowerCAmelCase	:      Dict      ,					lowerCAmelCase	:      Optional[int]      ,					lowerCAmelCase	:      str      ,					lowerCAmelCase	:      List[Any]      ,					lowerCAmelCase	:      List[str]      ,					lowerCAmelCase	:      Dict)						->					Optional[int]:




         """simple docstring"""

         lowercase__  =	MegatronBertForPreTraining(config=lowerCAmelCase)
         model.to(lowerCAmelCase)
         model.eval()
         lowercase__  =	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       UpperCAmelCase		(    self	:      List[Any]      ,					lowerCAmelCase	:      Union[str, Any]      ,					lowerCAmelCase	:      List[str]      ,					lowerCAmelCase	:      Tuple      ,					lowerCAmelCase	:      List[str]      ,					lowerCAmelCase	:      Union[str, Any]      ,					lowerCAmelCase	:      List[str]      ,					lowerCAmelCase	:      Tuple)						->					Tuple:




         """simple docstring"""

         lowercase__  =	MegatronBertForQuestionAnswering(config=lowerCAmelCase)
         model.to(lowerCAmelCase)
         model.eval()
         lowercase__  =	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       UpperCAmelCase		(    self	:      List[Any]      ,					lowerCAmelCase	:      int      ,					lowerCAmelCase	:      Union[str, Any]      ,					lowerCAmelCase	:      Tuple      ,					lowerCAmelCase	:      List[Any]      ,					lowerCAmelCase	:      Optional[Any]      ,					lowerCAmelCase	:      Optional[Any]      ,					lowerCAmelCase	:      str)						->					Dict:




         """simple docstring"""

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


   def       UpperCAmelCase		(    self	:      Any      ,					lowerCAmelCase	:      str      ,					lowerCAmelCase	:      Union[str, Any]      ,					lowerCAmelCase	:      str      ,					lowerCAmelCase	:      str      ,					lowerCAmelCase	:      Tuple      ,					lowerCAmelCase	:      Dict      ,					lowerCAmelCase	:      Tuple)						->					List[Any]:




         """simple docstring"""

         lowercase__  =	self.num_labels
         lowercase__  =	MegatronBertForTokenClassification(config=lowerCAmelCase)
         model.to(lowerCAmelCase)
         model.eval()
         lowercase__  =	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       UpperCAmelCase		(    self	:      Union[str, Any]      ,					lowerCAmelCase	:      Tuple      ,					lowerCAmelCase	:      List[str]      ,					lowerCAmelCase	:      Optional[Any]      ,					lowerCAmelCase	:      List[Any]      ,					lowerCAmelCase	:      str      ,					lowerCAmelCase	:      Union[str, Any]      ,					lowerCAmelCase	:      int)						->					Optional[int]:




         """simple docstring"""

         lowercase__  =	self.num_choices
         lowercase__  =	MegatronBertForMultipleChoice(config=lowerCAmelCase)
         model.to(lowerCAmelCase)
         model.eval()
         lowercase__  =	input_ids.unsqueeze(1).expand(-1      ,					self.num_choices      ,					-1).contiguous()
         lowercase__  =	token_type_ids.unsqueeze(1).expand(-1      ,					self.num_choices      ,					-1).contiguous()
         lowercase__  =	input_mask.unsqueeze(1).expand(-1      ,					self.num_choices      ,					-1).contiguous()
         lowercase__  =	model(
             lowerCAmelCase      ,					attention_mask=lowerCAmelCase      ,					token_type_ids=lowerCAmelCase      ,					labels=lowerCAmelCase      ,					)
         self.parent.assertEqual(result.logits.shape      ,					(self.batch_size, self.num_choices))



   def       UpperCAmelCase		(    self	:      List[str])						->					Optional[Any]:




         """simple docstring"""

         lowercase__  =	self.prepare_config_and_inputs()
         (
             (
             lowercase__
         ),						(
             lowercase__
         ),						(
             lowercase__
         ),						(
             lowercase__
         ),						(
             lowercase__
         ),						(
             lowercase__
         ),						(
             lowercase__
         ),						
         )  =	config_and_inputs
         lowercase__  =	{'input_ids': input_ids, 'token_type_ids': token_type_ids, 'attention_mask': input_mask}
         return config, inputs_dict



@require_torch
class    UpperCAmelCase__(				lowerCamelCase  ,				lowerCamelCase  ,				unittest.TestCase						):

   '''simple docstring'''




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


   def       UpperCAmelCase		(    self	:      List[Any]      ,					lowerCAmelCase	:      List[Any]      ,					lowerCAmelCase	:      List[str]      ,					lowerCAmelCase	:      int=False)						->					Any:




         """simple docstring"""

         lowercase__  =	super()._prepare_for_class(lowerCAmelCase      ,					lowerCAmelCase      ,					return_labels=lowerCAmelCase)

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


   def       UpperCAmelCase		(    self	:      str)						->					Any:




         """simple docstring"""

         lowercase__  =	MegatronBertModelTester(self)
         lowercase__  =	ConfigTester(self      ,					config_class=lowerCAmelCase      ,					hidden_size=37)


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




         """simple docstring"""

         self.config_tester.run_common_tests()


   def       UpperCAmelCase		(    self	:      Any)						->					List[str]:




         """simple docstring"""

         lowercase__  =	self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_megatron_bert_model(*lowerCAmelCase)


   def       UpperCAmelCase		(    self	:      Any)						->					Union[str, Any]:




         """simple docstring"""

         lowercase__  =	self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_megatron_bert_for_masked_lm(*lowerCAmelCase)


   def       UpperCAmelCase		(    self	:      List[Any])						->					Optional[int]:




         """simple docstring"""

         lowercase__  =	self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_megatron_bert_for_multiple_choice(*lowerCAmelCase)


   def       UpperCAmelCase		(    self	:      int)						->					Any:




         """simple docstring"""

         lowercase__  =	self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_megatron_bert_for_next_sequence_prediction(*lowerCAmelCase)


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




         """simple docstring"""

         lowercase__  =	self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_megatron_bert_for_pretraining(*lowerCAmelCase)


   def       UpperCAmelCase		(    self	:      Dict)						->					Optional[int]:




         """simple docstring"""

         lowercase__  =	self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_megatron_bert_for_question_answering(*lowerCAmelCase)


   def       UpperCAmelCase		(    self	:      List[str])						->					Optional[int]:




         """simple docstring"""

         lowercase__  =	self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_megatron_bert_for_sequence_classification(*lowerCAmelCase)



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




         """simple docstring"""

         lowercase__  =	self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_megatron_bert_for_token_classification(*lowerCAmelCase)

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


a__   :					Dict 						=      1E-4



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

   '''simple docstring'''




   @slow
   @unittest.skip('Model is not available.')
   def       UpperCAmelCase		(    self	:      Union[str, Any])						->					Union[str, Any]:




         """simple docstring"""

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

         lowercase__  =	[-0.60_40, -0.25_17, -0.10_25, 0.34_20, -0.67_58, -0.00_17, -0.10_89, -0.19_90, 0.57_28]
         for ii in range(3):
               for jj in range(3):
                     lowercase__  =	output[0, ii, jj]
                     lowercase__  =	expected[3 * ii + jj]
                     lowercase__  =	'ii={} jj={} a={} b={}'.format(lowerCAmelCase      ,					lowerCAmelCase      ,					lowerCAmelCase      ,					lowerCAmelCase)
                     self.assertTrue(math.isclose(lowerCAmelCase      ,					lowerCAmelCase      ,					rel_tol=lowerCAmelCase      ,					abs_tol=lowerCAmelCase)      ,					msg=lowerCAmelCase)





from argparse import ArgumentParser

from .add_new_model import AddNewModelCommand
from .add_new_model_like import AddNewModelLikeCommand
from .convert import ConvertCommand
from .download import DownloadCommand
from .env import EnvironmentCommand
from .lfs import LfsCommands
from .pt_to_tf import PTtoTFCommand
from .run import RunCommand
from .serving import ServeCommand
from .user import UserCommands

def 	_lowerCAmelCase (    ):
      lowercase__  =	ArgumentParser('Transformers CLI tool'    ,	usage='transformers-cli <command> [<args>]'						)
      lowercase__  =	parser.add_subparsers(help='transformers-cli command helpers'						)

      # Register commands
      ConvertCommand.register_subcommand(A__						)
      DownloadCommand.register_subcommand(A__						)
      EnvironmentCommand.register_subcommand(A__						)
      RunCommand.register_subcommand(A__						)
      ServeCommand.register_subcommand(A__						)
      UserCommands.register_subcommand(A__						)
      AddNewModelCommand.register_subcommand(A__						)
      AddNewModelLikeCommand.register_subcommand(A__						)
      LfsCommands.register_subcommand(A__						)
      PTtoTFCommand.register_subcommand(A__						)

      # Let's go
      lowercase__  =	parser.parse_args()

      if not hasattr(A__    ,	'func'						):
            parser.print_help()
            exit(1						)

      # Run
      lowercase__  =	args.func(A__						)
      service.run()


if __name__ == "__main__":
       main()
