





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  lowerCAmelCase							:







  def __init__(      self  :int   ,							_lowercase  :Any   ,							_lowercase  :Tuple=13   ,							_lowercase  :int=7   ,							_lowercase  :Optional[int]=True   ,							_lowercase  :Tuple=True   ,							_lowercase  :List[str]=True   ,							_lowercase  :str=True   ,							_lowercase  :Optional[Any]=99   ,							_lowercase  :Union[str, Any]=64   ,							_lowercase  :str=32   ,							_lowercase  :Any=5   ,							_lowercase  :Tuple=4   ,							_lowercase  :List[Any]=37   ,							_lowercase  :Optional[Any]="gelu"   ,							_lowercase  :Optional[int]=0.1   ,							_lowercase  :int=0.1   ,							_lowercase  :Optional[int]=5_12   ,							_lowercase  :List[Any]=16   ,							_lowercase  :Optional[int]=2   ,							_lowercase  :int=0.02   ,							_lowercase  :Tuple=3   ,							_lowercase  :Tuple=4   ,							_lowercase  :Any=None   ,							):






     '''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  :Any ):






     '''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  :int ):






     '''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=_lowercase   ,							initializer_range=self.initializer_range   ,							)







  def   UpperCAmelCase					(      self  :Optional[Any]   ,							_lowercase  :int   ,							_lowercase  :List[Any]   ,							_lowercase  :Optional[Any]   ,							_lowercase  :Optional[Any]   ,							_lowercase  :str   ,							_lowercase  :int   ,							_lowercase  :Dict ):






     '''simple docstring'''





     lowercase__			       =  MegatronBertModel(config=_lowercase )
     model.to(_lowercase )
     model.eval()
     lowercase__			       =  model(_lowercase   ,							attention_mask=_lowercase   ,							token_type_ids=_lowercase )
     lowercase__			       =  model(_lowercase   ,							token_type_ids=_lowercase )
     lowercase__			       =  model(_lowercase )

     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  :Any   ,							_lowercase  :Optional[int]   ,							_lowercase  :Dict   ,							_lowercase  :Tuple   ,							_lowercase  :str   ,							_lowercase  :Union[str, Any]   ,							_lowercase  :List[str]   ,							_lowercase  :Dict ):






     '''simple docstring'''





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







  def   UpperCAmelCase					(      self  :List[Any]   ,							_lowercase  :List[Any]   ,							_lowercase  :int   ,							_lowercase  :Union[str, Any]   ,							_lowercase  :Tuple   ,							_lowercase  :Optional[int]   ,							_lowercase  :Any   ,							_lowercase  :Any ):






     '''simple docstring'''





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







  def   UpperCAmelCase					(      self  :int   ,							_lowercase  :Union[str, Any]   ,							_lowercase  :Union[str, Any]   ,							_lowercase  :Dict   ,							_lowercase  :Optional[int]   ,							_lowercase  :Tuple   ,							_lowercase  :int   ,							_lowercase  :str ):






     '''simple docstring'''





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







  def   UpperCAmelCase					(      self  :Dict   ,							_lowercase  :Any   ,							_lowercase  :Any   ,							_lowercase  :int   ,							_lowercase  :Any   ,							_lowercase  :Optional[int]   ,							_lowercase  :Dict   ,							_lowercase  :Union[str, Any] ):






     '''simple docstring'''





     lowercase__			       =  MegatronBertForPreTraining(config=_lowercase )
     model.to(_lowercase )
     model.eval()
     lowercase__			       =  model(
         _lowercase   ,							attention_mask=_lowercase   ,							token_type_ids=_lowercase   ,							labels=_lowercase   ,							next_sentence_label=_lowercase   ,							)
     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  :Union[str, Any]   ,							_lowercase  :Any   ,							_lowercase  :Optional[int]   ,							_lowercase  :Tuple   ,							_lowercase  :Union[str, Any]   ,							_lowercase  :Tuple   ,							_lowercase  :Optional[int]   ,							_lowercase  :str ):






     '''simple docstring'''





     lowercase__			       =  MegatronBertForQuestionAnswering(config=_lowercase )
     model.to(_lowercase )
     model.eval()
     lowercase__			       =  model(
         _lowercase   ,							attention_mask=_lowercase   ,							token_type_ids=_lowercase   ,							start_positions=_lowercase   ,							end_positions=_lowercase   ,							)
     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  :Any   ,							_lowercase  :List[Any]   ,							_lowercase  :List[str]   ,							_lowercase  :List[str]   ,							_lowercase  :List[str]   ,							_lowercase  :Optional[Any]   ,							_lowercase  :List[Any]   ,							_lowercase  :List[Any] ):






     '''simple docstring'''





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







  def   UpperCAmelCase					(      self  :str   ,							_lowercase  :Dict   ,							_lowercase  :Any   ,							_lowercase  :List[Any]   ,							_lowercase  :str   ,							_lowercase  :List[str]   ,							_lowercase  :int   ,							_lowercase  :Tuple ):






     '''simple docstring'''





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







  def   UpperCAmelCase					(      self  :int   ,							_lowercase  :List[str]   ,							_lowercase  :List[str]   ,							_lowercase  :Optional[Any]   ,							_lowercase  :Dict   ,							_lowercase  :Optional[Any]   ,							_lowercase  :Dict   ,							_lowercase  :Any ):






     '''simple docstring'''





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







  def   UpperCAmelCase					(      self  :List[str] ):






     '''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  lowerCAmelCase							(       lowercase_							,							lowercase_							,							unittest.TestCase ):
  __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   UpperCAmelCase					(      self  :Dict   ,							_lowercase  :Optional[int]   ,							_lowercase  :List[Any]   ,							_lowercase  :Tuple=False ):






     '''simple docstring'''





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

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







  def   UpperCAmelCase					(      self  :Any ):






     '''simple docstring'''





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







  def   UpperCAmelCase					(      self  :List[str] ):






     '''simple docstring'''





     self.config_tester.run_common_tests()







  def   UpperCAmelCase					(      self  :Tuple ):






     '''simple docstring'''





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







  def   UpperCAmelCase					(      self  :Optional[Any] ):






     '''simple docstring'''





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







  def   UpperCAmelCase					(      self  :Any ):






     '''simple docstring'''





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







  def   UpperCAmelCase					(      self  :int ):






     '''simple docstring'''





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







  def   UpperCAmelCase					(      self  :Optional[Any] ):






     '''simple docstring'''





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







  def   UpperCAmelCase					(      self  :Dict ):






     '''simple docstring'''





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







  def   UpperCAmelCase					(      self  :Any ):






     '''simple docstring'''





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







  def   UpperCAmelCase					(      self  :List[str] ):






     '''simple docstring'''





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


def _A       (   __magic_name__	):
   return torch.tensor(
       __magic_name__  , dtype=torch.long  , device=__magic_name__  , )


_snake_case           =			1E-4







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







  @slow
  @unittest.skip("Model is not available." )
  def   UpperCAmelCase					(      self  :Optional[Any] ):






     '''simple docstring'''





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

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













from __future__ import annotations


def _A       (   __magic_name__  , __magic_name__	):
   lowercase__			       =  []
   create_all_state(1  , __magic_name__  , __magic_name__  , []  , __magic_name__	)
   return result


def _A       (   __magic_name__  , __magic_name__  , __magic_name__  , __magic_name__  , __magic_name__  , ):
   if level == 0:
      total_list.append(current_list[:]	)
      return

   for i in range(__magic_name__  , total_number - level + 2	):
      current_list.append(__magic_name__	)
      create_all_state(i + 1  , __magic_name__  , level - 1  , __magic_name__  , __magic_name__	)
      current_list.pop()


def _A       (   __magic_name__	):
   for i in total_list:
      print(*__magic_name__	)


if __name__ == "__main__":
  _snake_case           =			4
  _snake_case           =			2
  _snake_case           =			generate_all_combinations(n, k)
  print_all_state(total_list)






