





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







       """simple docstring"""





       def __init__(     self						:						List[str]    ,__A						:						List[str]    ,__A						:						List[str]=13    ,__A						:						Dict=7    ,__A						:						int=True    ,__A						:						Tuple=True    ,__A						:						Tuple=True    ,__A						:						List[Any]=True    ,__A						:						Any=99    ,__A						:						Union[str, Any]=64    ,__A						:						Dict=32    ,__A						:						Any=5    ,__A						:						List[str]=4    ,__A						:						Optional[int]=37    ,__A						:						Optional[int]="gelu"    ,__A						:						Any=0.1    ,__A						:						str=0.1    ,__A						:						str=512    ,__A						:						List[str]=16    ,__A						:						List[str]=2    ,__A						:						List[Any]=0.02    ,__A						:						Optional[int]=3    ,__A						:						Optional[Any]=4    ,__A						:						Optional[Any]=None    ,)    ->       Any:
         _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						:						Dict	)    ->       Union[str, Any]:
         _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						:						Tuple	)    ->       Union[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=__A    ,initializer_range=self.initializer_range    ,)







       def  __UpperCAmelCase			(     self						:						Union[str, Any]    ,__A						:						Optional[Any]    ,__A						:						Any    ,__A						:						Optional[Any]    ,__A						:						Any    ,__A						:						str    ,__A						:						Tuple    ,__A						:						Optional[int]	)    ->       str:
         _lowercase						    =			MegatronBertModel(config=__A	)
         model.to(__A	)
         model.eval()
         _lowercase						    =			model(__A    ,attention_mask=__A    ,token_type_ids=__A	)
         _lowercase						    =			model(__A    ,token_type_ids=__A	)
         _lowercase						    =			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  __UpperCAmelCase			(     self						:						Tuple    ,__A						:						Tuple    ,__A						:						Tuple    ,__A						:						Tuple    ,__A						:						Dict    ,__A						:						Optional[int]    ,__A						:						Dict    ,__A						:						Any	)    ->       int:
         _lowercase						    =			MegatronBertForMaskedLM(config=__A	)
         model.to(__A	)
         model.eval()
         _lowercase						    =			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  __UpperCAmelCase			(     self						:						Tuple    ,__A						:						Any    ,__A						:						Union[str, Any]    ,__A						:						Any    ,__A						:						Union[str, Any]    ,__A						:						Union[str, Any]    ,__A						:						int    ,__A						:						Tuple	)    ->       Any:
         _lowercase						    =			MegatronBertForCausalLM(config=__A	)
         model.to(__A	)
         model.eval()
         _lowercase						    =			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  __UpperCAmelCase			(     self						:						Tuple    ,__A						:						Tuple    ,__A						:						Dict    ,__A						:						Dict    ,__A						:						int    ,__A						:						Optional[Any]    ,__A						:						Tuple    ,__A						:						str	)    ->       Tuple:
         _lowercase						    =			MegatronBertForNextSentencePrediction(config=__A	)
         model.to(__A	)
         model.eval()
         _lowercase						    =			model(
             __A    ,attention_mask=__A    ,token_type_ids=__A    ,labels=__A    ,)
         self.parent.assertEqual(result.logits.shape    ,(self.batch_size, 2)	)







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







       def  __UpperCAmelCase			(     self						:						int    ,__A						:						List[str]    ,__A						:						str    ,__A						:						Any    ,__A						:						Union[str, Any]    ,__A						:						Optional[int]    ,__A						:						List[str]    ,__A						:						List[Any]	)    ->       Optional[int]:
         _lowercase						    =			self.num_labels
         _lowercase						    =			MegatronBertForTokenClassification(config=__A	)
         model.to(__A	)
         model.eval()
         _lowercase						    =			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  __UpperCAmelCase			(     self						:						List[str]    ,__A						:						List[Any]    ,__A						:						List[Any]    ,__A						:						Any    ,__A						:						Optional[Any]    ,__A						:						Dict    ,__A						:						str    ,__A						:						List[Any]	)    ->       Optional[Any]:
         _lowercase						    =			self.num_choices
         _lowercase						    =			MegatronBertForMultipleChoice(config=__A	)
         model.to(__A	)
         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(
             __A    ,attention_mask=__A    ,token_type_ids=__A    ,labels=__A    ,)
         self.parent.assertEqual(result.logits.shape    ,(self.batch_size, self.num_choices)	)







       def  __UpperCAmelCase			(     self						:						Any	)    ->       int:
         _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 					A_					(   UpperCAmelCase       ,						UpperCAmelCase       ,						unittest.TestCase						):







       """simple docstring"""





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







       def  __UpperCAmelCase			(     self						:						List[Any]    ,__A						:						Dict    ,__A						:						Dict    ,__A						:						List[str]=False	)    ->       Optional[Any]:
         _lowercase						    =			super()._prepare_for_class(__A    ,__A    ,return_labels=__A	)

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







       def  __UpperCAmelCase			(     self						:						Optional[Any]	)    ->       Optional[int]:
         _lowercase						    =			MegatronBertModelTester(self	)
         _lowercase						    =			ConfigTester(self    ,config_class=__A    ,hidden_size=37	)







       def  __UpperCAmelCase			(     self						:						Dict	)    ->       str:
         self.config_tester.run_common_tests()







       def  __UpperCAmelCase			(     self						:						List[Any]	)    ->       Optional[Any]:
         _lowercase						    =			self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_megatron_bert_model(*__A	)







       def  __UpperCAmelCase			(     self						:						int	)    ->       str:
         _lowercase						    =			self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_megatron_bert_for_masked_lm(*__A	)







       def  __UpperCAmelCase			(     self						:						List[str]	)    ->       int:
         _lowercase						    =			self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_megatron_bert_for_multiple_choice(*__A	)







       def  __UpperCAmelCase			(     self						:						Union[str, Any]	)    ->       Dict:
         _lowercase						    =			self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_megatron_bert_for_next_sequence_prediction(*__A	)







       def  __UpperCAmelCase			(     self						:						List[Any]	)    ->       Any:
         _lowercase						    =			self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_megatron_bert_for_pretraining(*__A	)







       def  __UpperCAmelCase			(     self						:						Optional[int]	)    ->       Union[str, Any]:
         _lowercase						    =			self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_megatron_bert_for_question_answering(*__A	)







       def  __UpperCAmelCase			(     self						:						Dict	)    ->       Dict:
         _lowercase						    =			self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_megatron_bert_for_sequence_classification(*__A	)







       def  __UpperCAmelCase			(     self						:						List[Any]	)    ->       int:
         _lowercase						    =			self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_megatron_bert_for_token_classification(*__A	)

def 				SCREAMING_SNAKE_CASE__							(	snake_case__	:Optional[int]  )   ->    Optional[Any]:
  return torch.tensor(
      snake_case__      ,				dtype=torch.long      ,				device=snake_case__      ,				)


snake_case								=							1e-4







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







       """simple docstring"""





       @slow
       @unittest.skip('Model is not available.'	)
       def  __UpperCAmelCase			(     self						:						Any	)    ->       Tuple:
         _lowercase						    =			'nvidia/megatron-bert-uncased-345m'
         if "MYDIR" in os.environ:
           _lowercase						    =			os.path.join(os.environ['MYDIR']    ,__A	)
         _lowercase						    =			MegatronBertModel.from_pretrained(__A	)
         model.to(__A	)
         model.half()
         _lowercase						    =			_long_tensor([[101, 7110, 1005, 1056, 2023, 1_1333, 1_7413, 1029, 102]]	)
         with torch.no_grad():
           _lowercase						    =			model(__A	)[0]
         _lowercase						    =			torch.Size((1, 9, 1024)	)
         self.assertEqual(output.shape    ,__A	)

         _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(__A    ,__A    ,__A    ,__A	)
             self.assertTrue(math.isclose(__A    ,__A    ,rel_tol=__A    ,abs_tol=__A	)    ,msg=__A	)





'''simple docstring'''



import argparse
import os
import re


UpperCAmelCase_							:	List[str]	     =		'''src/transformers/models/auto'''


# re pattern that matches mapping introductions:
#    SUPER_MODEL_MAPPING_NAMES = OrderedDict or SUPER_MODEL_MAPPING = OrderedDict
UpperCAmelCase_							:	Tuple	     =		re.compile(R'''[A-Z_]+_MAPPING(\s+|_[A-Z_]+\s+)=\s+OrderedDict''')
# re pattern that matches identifiers in mappings
UpperCAmelCase_							:	Dict	     =		re.compile(R'''\s*\(\s*"(\S[^"]+)"''')







def       _UpperCamelCase			(_lowerCamelCase				:  List[Any] ,						_lowerCamelCase				:  bool = False						)-> str:




     '''simple docstring'''





     with open(_lowerCamelCase ,						'''r''' ,						encoding='''utf-8'''						) as f:
          __snake_case            =	f.read()

     __snake_case            =	content.split('''\n'''						)
     __snake_case            =	[]
     __snake_case            =	0
     while line_idx < len(_lowerCamelCase						):
          if _re_intro_mapping.search(lines[line_idx]						) is not None:
               __snake_case            =	len(re.search(R'''^(\s*)\S''' ,						lines[line_idx]						).groups()[0]						) + 8
               # Start of a new mapping!
               while not lines[line_idx].startswith(''' ''' * indent + '''('''						):
                    new_lines.append(lines[line_idx]						)
                    line_idx += 1

               __snake_case            =	[]
               while lines[line_idx].strip() != "]":
                    # Blocks either fit in one line or not
                    if lines[line_idx].strip() == "(":
                         __snake_case            =	line_idx
                         while not lines[line_idx].startswith(''' ''' * indent + ''')'''						):
                              line_idx += 1
                         blocks.append('''\n'''.join(lines[start_idx : line_idx + 1]						)						)
                    else:
                         blocks.append(lines[line_idx]						)
                    line_idx += 1

               # Sort blocks by their identifiers
               __snake_case            =	sorted(_lowerCamelCase ,						key=lambda _lowerCamelCase						: _re_identifier.search(_lowerCamelCase						).groups()[0]						)
               new_lines += blocks
          else:
               new_lines.append(lines[line_idx]						)
               line_idx += 1

     if overwrite:
          with open(_lowerCamelCase ,						'''w''' ,						encoding='''utf-8'''						) as f:
               f.write('''\n'''.join(_lowerCamelCase						)						)
     elif "\n".join(_lowerCamelCase						) != content:
          return True







def       _UpperCamelCase			(_lowerCamelCase				:  bool = False						)-> Tuple:




     '''simple docstring'''





     __snake_case            =	[os.path.join(_lowerCamelCase ,						_lowerCamelCase						) for f in os.listdir(_lowerCamelCase						) if f.endswith('''.py'''						)]
     __snake_case            =	[sort_auto_mapping(_lowerCamelCase ,						overwrite=_lowerCamelCase						) for fname in fnames]

     if not overwrite and any(_lowerCamelCase						):
          __snake_case            =	[f for f, d in zip(_lowerCamelCase ,						_lowerCamelCase						) if d]
          raise ValueError(
              f'''The following files have auto mappings that need sorting: {", ".join(_lowerCamelCase						)}. Run `make style` to fix'''
              ''' this.'''						)


if __name__ == "__main__":
     UpperCAmelCase_							:	str	     =		argparse.ArgumentParser()
     parser.add_argument('''--check_only''', action='''store_true''', help='''Whether to only check or fix style.''')
     UpperCAmelCase_							:	List[Any]	     =		parser.parse_args()

     sort_all_auto_mappings(not args.check_only)






