


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


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





      '''simple docstring'''






      UpperCamelCase__					:    str          =   parent
      UpperCamelCase__					:    List[Any]          =   batch_size
      UpperCamelCase__					:    Dict          =   seq_length
      UpperCamelCase__					:    Optional[Any]          =   is_training
      UpperCamelCase__					:    Any          =   use_input_mask
      UpperCamelCase__					:    List[Any]          =   use_token_type_ids
      UpperCamelCase__					:    str          =   use_labels
      UpperCamelCase__					:    List[str]          =   vocab_size
      UpperCamelCase__					:    List[str]          =   hidden_size
      UpperCamelCase__					:    Union[str, Any]          =   embedding_size
      UpperCamelCase__					:    int          =   num_hidden_layers
      UpperCamelCase__					:    str          =   num_attention_heads
      UpperCamelCase__					:    Any          =   intermediate_size
      UpperCamelCase__					:    Dict          =   hidden_act
      UpperCamelCase__					:    Optional[int]          =   hidden_dropout_prob
      UpperCamelCase__					:    int          =   attention_probs_dropout_prob
      UpperCamelCase__					:    str          =   max_position_embeddings
      UpperCamelCase__					:    List[Any]          =   type_vocab_size
      UpperCamelCase__					:    str          =   type_sequence_label_size
      UpperCamelCase__					:    Dict          =   initializer_range
      UpperCamelCase__					:    Optional[Any]          =   num_labels
      UpperCamelCase__					:    Optional[Any]          =   num_choices
      UpperCamelCase__					:    Union[str, Any]          =   scope


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





      '''simple docstring'''






      UpperCamelCase__					:    Optional[Any]          =   ids_tensor([self.batch_size, self.seq_length]							,	self.vocab_size			)

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

      UpperCamelCase__					:    str          =   None
      if self.use_token_type_ids:
       UpperCamelCase__					:    Optional[int]          =   ids_tensor([self.batch_size, self.seq_length]							,	self.type_vocab_size			)

      UpperCamelCase__					:    int          =   None
      UpperCamelCase__					:    Optional[Any]          =   None
      UpperCamelCase__					:    Optional[Any]          =   None
      if self.use_labels:
       UpperCamelCase__					:    int          =   ids_tensor([self.batch_size]							,	self.type_sequence_label_size			)
       UpperCamelCase__					:    Dict          =   ids_tensor([self.batch_size, self.seq_length]							,	self.num_labels			)
       UpperCamelCase__					:    Any          =   ids_tensor([self.batch_size]							,	self.num_choices			)

      UpperCamelCase__					:    int          =   self.get_config()

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


     def      UpperCAmelCase__	(      self							:     Optional[int]			)  ->      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							:     str							,	lowerCamelCase__							:     List[Any]							,	lowerCamelCase__							:     List[str]							,	lowerCamelCase__							:     Union[str, Any]							,	lowerCamelCase__							:     int							,	lowerCamelCase__							:     Optional[Any]							,	lowerCamelCase__							:     Tuple							,	lowerCamelCase__							:     List[Any]			)  ->      Dict:





      '''simple docstring'''






      UpperCamelCase__					:    int          =   MegatronBertModel(config=lowerCamelCase__			)
      model.to(lowerCamelCase__			)
      model.eval()
      UpperCamelCase__					:    List[str]          =   model(lowerCamelCase__							,	attention_mask=lowerCamelCase__							,	token_type_ids=lowerCamelCase__			)
      UpperCamelCase__					:    Optional[int]          =   model(lowerCamelCase__							,	token_type_ids=lowerCamelCase__			)
      UpperCamelCase__					:    int          =   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							:     Union[str, Any]							,	lowerCamelCase__							:     Dict							,	lowerCamelCase__							:     int							,	lowerCamelCase__							:     Union[str, Any]							,	lowerCamelCase__							:     Union[str, Any]							,	lowerCamelCase__							:     int							,	lowerCamelCase__							:     int							,	lowerCamelCase__							:     Optional[int]			)  ->      Union[str, Any]:





      '''simple docstring'''






      UpperCamelCase__					:    List[str]          =   MegatronBertForMaskedLM(config=lowerCamelCase__			)
      model.to(lowerCamelCase__			)
      model.eval()
      UpperCamelCase__					:    int          =   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							:     Optional[Any]							,	lowerCamelCase__							:     List[Any]							,	lowerCamelCase__							:     Optional[Any]							,	lowerCamelCase__							:     List[Any]							,	lowerCamelCase__							:     List[Any]							,	lowerCamelCase__							:     str							,	lowerCamelCase__							:     Any							,	lowerCamelCase__							:     str			)  ->      str:





      '''simple docstring'''






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





      '''simple docstring'''






      UpperCamelCase__					:    str          =   MegatronBertForNextSentencePrediction(config=lowerCamelCase__			)
      model.to(lowerCamelCase__			)
      model.eval()
      UpperCamelCase__					:    Tuple          =   model(
          lowerCamelCase__							,	attention_mask=lowerCamelCase__							,	token_type_ids=lowerCamelCase__							,	labels=lowerCamelCase__							,	)
      self.parent.assertEqual(result.logits.shape							,	(self.batch_size, 2)			)


     def      UpperCAmelCase__	(      self							:     Tuple							,	lowerCamelCase__							:     Optional[int]							,	lowerCamelCase__							:     Optional[int]							,	lowerCamelCase__							:     List[str]							,	lowerCamelCase__							:     Optional[int]							,	lowerCamelCase__							:     Optional[Any]							,	lowerCamelCase__							:     int							,	lowerCamelCase__							:     List[str]			)  ->      Tuple:





      '''simple docstring'''






      UpperCamelCase__					:    Union[str, Any]          =   MegatronBertForPreTraining(config=lowerCamelCase__			)
      model.to(lowerCamelCase__			)
      model.eval()
      UpperCamelCase__					:    Optional[int]          =   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__							:     Union[str, Any]							,	lowerCamelCase__							:     Any							,	lowerCamelCase__							:     List[str]							,	lowerCamelCase__							:     Tuple							,	lowerCamelCase__							:     List[Any]							,	lowerCamelCase__							:     List[Any]			)  ->      Optional[int]:





      '''simple docstring'''






      UpperCamelCase__					:    Dict          =   MegatronBertForQuestionAnswering(config=lowerCamelCase__			)
      model.to(lowerCamelCase__			)
      model.eval()
      UpperCamelCase__					:    int          =   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							:     Any							,	lowerCamelCase__							:     int							,	lowerCamelCase__							:     Optional[Any]							,	lowerCamelCase__							:     int							,	lowerCamelCase__							:     int							,	lowerCamelCase__							:     Dict							,	lowerCamelCase__							:     Tuple							,	lowerCamelCase__							:     int			)  ->      Tuple:





      '''simple docstring'''






      UpperCamelCase__					:    Dict          =   self.num_labels
      UpperCamelCase__					:    Any          =   MegatronBertForSequenceClassification(lowerCamelCase__			)
      model.to(lowerCamelCase__			)
      model.eval()
      UpperCamelCase__					:    Any          =   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__							:     Union[str, Any]							,	lowerCamelCase__							:     Optional[Any]							,	lowerCamelCase__							:     Dict							,	lowerCamelCase__							:     int							,	lowerCamelCase__							:     Any							,	lowerCamelCase__							:     List[Any]							,	lowerCamelCase__							:     Any			)  ->      int:





      '''simple docstring'''






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





      '''simple docstring'''






      UpperCamelCase__					:    Optional[Any]          =   self.num_choices
      UpperCamelCase__					:    Tuple          =   MegatronBertForMultipleChoice(config=lowerCamelCase__			)
      model.to(lowerCamelCase__			)
      model.eval()
      UpperCamelCase__					:    Union[str, Any]          =   input_ids.unsqueeze(1			).expand(-1							,	self.num_choices							,	-1			).contiguous()
      UpperCamelCase__					:    List[str]          =   token_type_ids.unsqueeze(1			).expand(-1							,	self.num_choices							,	-1			).contiguous()
      UpperCamelCase__					:    int          =   input_mask.unsqueeze(1			).expand(-1							,	self.num_choices							,	-1			).contiguous()
      UpperCamelCase__					:    Any          =   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							:     Any			)  ->      Optional[Any]:





      '''simple docstring'''






      UpperCamelCase__					:    str          =   self.prepare_config_and_inputs()
      (
          (
          UpperCamelCase__
      )       ,				(
          UpperCamelCase__
      )       ,				(
          UpperCamelCase__
      )       ,				(
          UpperCamelCase__
      )       ,				(
          UpperCamelCase__
      )       ,				(
          UpperCamelCase__
      )       ,				(
          UpperCamelCase__
      )       ,				
      )					:    Any          =   config_and_inputs
      UpperCamelCase__					:    Tuple          =   {'''input_ids''': input_ids, '''token_type_ids''': token_type_ids, '''attention_mask''': input_mask}
      return config, inputs_dict



@require_torch
class       __magic_name__	(     __lowerCAmelCase					,				__lowerCAmelCase					,				unittest.TestCase):
     A:		List[Any]			=      (
         (
             MegatronBertModel,
             MegatronBertForMaskedLM,
             MegatronBertForCausalLM,
             MegatronBertForMultipleChoice,
             MegatronBertForNextSentencePrediction,
             MegatronBertForPreTraining,
             MegatronBertForQuestionAnswering,
             MegatronBertForSequenceClassification,
             MegatronBertForTokenClassification,
         )
         if is_torch_available()
         else ()
     )
     A:		Optional[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:		Tuple			=      True
     # test_resize_embeddings = False
     A:		List[str]			=      False


     def      UpperCAmelCase__	(      self							:     List[Any]							,	lowerCamelCase__							:     Dict							,	lowerCamelCase__							:     Tuple							,	lowerCamelCase__							:     List[Any]=False			)  ->      List[Any]:





      '''simple docstring'''






      UpperCamelCase__					:    List[Any]          =   super()._prepare_for_class(lowerCamelCase__							,	lowerCamelCase__							,	return_labels=lowerCamelCase__			)

      if return_labels:
       if model_class in get_values(lowerCamelCase__			):
        UpperCamelCase__					:    Dict          =   torch.zeros(
            (self.model_tester.batch_size, self.model_tester.seq_length)							,	dtype=torch.long							,	device=lowerCamelCase__			)
        UpperCamelCase__					:    Optional[Any]          =   torch.zeros(
            self.model_tester.batch_size							,	dtype=torch.long							,	device=lowerCamelCase__			)
      return inputs_dict


     def      UpperCAmelCase__	(      self							:     Optional[Any]			)  ->      Dict:





      '''simple docstring'''






      UpperCamelCase__					:    int          =   MegatronBertModelTester(self			)
      UpperCamelCase__					:    Union[str, Any]          =   ConfigTester(self							,	config_class=lowerCamelCase__							,	hidden_size=37			)


     def      UpperCAmelCase__	(      self							:     Optional[int]			)  ->      Union[str, Any]:





      '''simple docstring'''






      self.config_tester.run_common_tests()


     def      UpperCAmelCase__	(      self							:     Dict			)  ->      Tuple:





      '''simple docstring'''






      UpperCamelCase__					:    Union[str, Any]          =   self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_megatron_bert_model(*lowerCamelCase__			)


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





      '''simple docstring'''






      UpperCamelCase__					:    List[Any]          =   self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_megatron_bert_for_masked_lm(*lowerCamelCase__			)


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





      '''simple docstring'''






      UpperCamelCase__					:    str          =   self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_megatron_bert_for_multiple_choice(*lowerCamelCase__			)


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





      '''simple docstring'''






      UpperCamelCase__					:    Union[str, Any]          =   self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_megatron_bert_for_next_sequence_prediction(*lowerCamelCase__			)


     def      UpperCAmelCase__	(      self							:     Optional[int]			)  ->      Tuple:





      '''simple docstring'''






      UpperCamelCase__					:    Optional[Any]          =   self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_megatron_bert_for_pretraining(*lowerCamelCase__			)


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





      '''simple docstring'''






      UpperCamelCase__					:    List[str]          =   self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_megatron_bert_for_question_answering(*lowerCamelCase__			)


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





      '''simple docstring'''






      UpperCamelCase__					:    List[str]          =   self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_megatron_bert_for_sequence_classification(*lowerCamelCase__			)


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





      '''simple docstring'''






      UpperCamelCase__					:    Any          =   self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_megatron_bert_for_token_classification(*lowerCamelCase__			)



def 			_a  (				SCREAMING_SNAKE_CASE	:				Any     ):







 """simple docstring"""



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


__UpperCamelCase       :       Dict						       =				1E-4






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


     @slow
     @unittest.skip('''Model is not available.'''			)
     def      UpperCAmelCase__	(      self							:     Any			)  ->      Dict:





      '''simple docstring'''






      UpperCamelCase__					:    Any          =   '''nvidia/megatron-bert-uncased-345m'''
      if "MYDIR" in os.environ:
       UpperCamelCase__					:    Union[str, Any]          =   os.path.join(os.environ['''MYDIR''']							,	lowerCamelCase__			)
      UpperCamelCase__					:    Dict          =   MegatronBertModel.from_pretrained(lowerCamelCase__			)
      model.to(lowerCamelCase__			)
      model.half()
      UpperCamelCase__					:    str          =   _long_tensor([[101, 7110, 1005, 1056, 2023, 11333, 17413, 1029, 102]]			)
      with torch.no_grad():
       UpperCamelCase__					:    Any          =   model(lowerCamelCase__			)[0]
      UpperCamelCase__					:    int          =   torch.Size((1, 9, 1024)			)
      self.assertEqual(output.shape							,	lowerCamelCase__			)

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






import functools

def 			_a  (				SCREAMING_SNAKE_CASE	:				str    , SCREAMING_SNAKE_CASE	:				str     ):







 """simple docstring"""



 UpperCamelCase__					:    str          =   len(SCREAMING_SNAKE_CASE     )
 UpperCamelCase__					:    Optional[Any]          =   len(SCREAMING_SNAKE_CASE     )

 @functools.cache
 def min_distance(SCREAMING_SNAKE_CASE	:				int    , SCREAMING_SNAKE_CASE	:				int     ) -> int:
  # if first word index is overflow - delete all from the second word
  if indexa >= len_worda:
   return len_worda - indexa
  # if second word index is overflow - delete all from the first word
  if indexa >= len_worda:
   return len_worda - indexa
  UpperCamelCase__					:    Tuple          =   int(worda[indexa] != worda[indexa]     )  # current letters not identical
  return min(
      1 + min_distance(indexa + 1    , SCREAMING_SNAKE_CASE     )    , 1 + min_distance(SCREAMING_SNAKE_CASE    , indexa + 1     )    , diff + min_distance(indexa + 1    , indexa + 1     )    , )

 return min_distance(0    , 0     )


if __name__ == "__main__":
       import doctest

       doctest.testmod()


