





"""simple docstring"""





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





      def __init__(       self      ,						lowercase      ,						lowercase=13      ,						lowercase=7      ,						lowercase=True      ,						lowercase=True      ,						lowercase=True      ,						lowercase=True      ,						lowercase=99      ,						lowercase=64      ,						lowercase=32      ,						lowercase=5      ,						lowercase=4      ,						lowercase=37      ,						lowercase="gelu"      ,						lowercase=0.1      ,						lowercase=0.1      ,						lowercase=512      ,						lowercase=16      ,						lowercase=2      ,						lowercase=0.02      ,						lowercase=3      ,						lowercase=4      ,						lowercase=None      ,						)		->     List[Any]:
            lowerCAmelCase  						=       parent
            lowerCAmelCase  						=       batch_size
            lowerCAmelCase  						=       seq_length
            lowerCAmelCase  						=       is_training
            lowerCAmelCase  						=       use_input_mask
            lowerCAmelCase  						=       use_token_type_ids
            lowerCAmelCase  						=       use_labels
            lowerCAmelCase  						=       vocab_size
            lowerCAmelCase  						=       hidden_size
            lowerCAmelCase  						=       embedding_size
            lowerCAmelCase  						=       num_hidden_layers
            lowerCAmelCase  						=       num_attention_heads
            lowerCAmelCase  						=       intermediate_size
            lowerCAmelCase  						=       hidden_act
            lowerCAmelCase  						=       hidden_dropout_prob
            lowerCAmelCase  						=       attention_probs_dropout_prob
            lowerCAmelCase  						=       max_position_embeddings
            lowerCAmelCase  						=       type_vocab_size
            lowerCAmelCase  						=       type_sequence_label_size
            lowerCAmelCase  						=       initializer_range
            lowerCAmelCase  						=       num_labels
            lowerCAmelCase  						=       num_choices
            lowerCAmelCase  						=       scope





      def     _snake_case			(       self )		->     int:
            lowerCAmelCase  						=       ids_tensor([self.batch_size, self.seq_length]      ,						self.vocab_size )

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

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

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

            lowerCAmelCase  						=       self.get_config()

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





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





      def     _snake_case			(       self      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase )		->     Tuple:
            lowerCAmelCase  						=       MegatronBertForPreTraining(config=lowercase )
            model.to(lowercase )
            model.eval()
            lowerCAmelCase  						=       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     _snake_case			(       self      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase )		->     List[Any]:
            lowerCAmelCase  						=       MegatronBertForQuestionAnswering(config=lowercase )
            model.to(lowercase )
            model.eval()
            lowerCAmelCase  						=       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     _snake_case			(       self      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase )		->     List[Any]:
            lowerCAmelCase  						=       self.num_labels
            lowerCAmelCase  						=       MegatronBertForSequenceClassification(lowercase )
            model.to(lowercase )
            model.eval()
            lowerCAmelCase  						=       model(lowercase      ,						attention_mask=lowercase      ,						token_type_ids=lowercase      ,						labels=lowercase )
            self.parent.assertEqual(result.logits.shape      ,						(self.batch_size, self.num_labels) )





      def     _snake_case			(       self      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase )		->     Union[str, Any]:
            lowerCAmelCase  						=       self.num_labels
            lowerCAmelCase  						=       MegatronBertForTokenClassification(config=lowercase )
            model.to(lowercase )
            model.eval()
            lowerCAmelCase  						=       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     _snake_case			(       self      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase      ,						lowercase )		->     Optional[Any]:
            lowerCAmelCase  						=       self.num_choices
            lowerCAmelCase  						=       MegatronBertForMultipleChoice(config=lowercase )
            model.to(lowercase )
            model.eval()
            lowerCAmelCase  						=       input_ids.unsqueeze(1 ).expand(-1      ,						self.num_choices      ,						-1 ).contiguous()
            lowerCAmelCase  						=       token_type_ids.unsqueeze(1 ).expand(-1      ,						self.num_choices      ,						-1 ).contiguous()
            lowerCAmelCase  						=       input_mask.unsqueeze(1 ).expand(-1      ,						self.num_choices      ,						-1 ).contiguous()
            lowerCAmelCase  						=       model(
                lowercase      ,						attention_mask=lowercase      ,						token_type_ids=lowercase      ,						labels=lowercase      ,						)
            self.parent.assertEqual(result.logits.shape      ,						(self.batch_size, self.num_choices) )





      def     _snake_case			(       self )		->     int:
            lowerCAmelCase  						=       self.prepare_config_and_inputs()
            (
                (
                lowerCAmelCase
            )						,		(
                lowerCAmelCase
            )						,		(
                lowerCAmelCase
            )						,		(
                lowerCAmelCase
            )						,		(
                lowerCAmelCase
            )						,		(
                lowerCAmelCase
            )						,		(
                lowerCAmelCase
            )						,		
            )  						=       config_and_inputs
            lowerCAmelCase  						=       {"""input_ids""": input_ids, """token_type_ids""": token_type_ids, """attention_mask""": input_mask}
            return config, inputs_dict







@require_torch
class 		lowercase					(					_UpperCAmelCase     ,  _UpperCAmelCase     ,  unittest.TestCase					):
      _SCREAMING_SNAKE_CASE					=      (
          (
              MegatronBertModel,
              MegatronBertForMaskedLM,
              MegatronBertForCausalLM,
              MegatronBertForMultipleChoice,
              MegatronBertForNextSentencePrediction,
              MegatronBertForPreTraining,
              MegatronBertForQuestionAnswering,
              MegatronBertForSequenceClassification,
              MegatronBertForTokenClassification,
          )
          if is_torch_available()
          else ()
      )
      _SCREAMING_SNAKE_CASE					=      (
          {
              '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					=      True
      # test_resize_embeddings = False
      _SCREAMING_SNAKE_CASE					=      False





      def     _snake_case			(       self      ,						lowercase      ,						lowercase      ,						lowercase=False )		->     int:
            lowerCAmelCase  						=       super()._prepare_for_class(lowercase      ,						lowercase      ,						return_labels=lowercase )

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





      def     _snake_case			(       self )		->     List[Any]:
            lowerCAmelCase  						=       MegatronBertModelTester(self )
            lowerCAmelCase  						=       ConfigTester(self      ,						config_class=lowercase      ,						hidden_size=37 )





      def     _snake_case			(       self )		->     Tuple:
            self.config_tester.run_common_tests()





      def     _snake_case			(       self )		->     str:
            lowerCAmelCase  						=       self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_megatron_bert_model(*lowercase )





      def     _snake_case			(       self )		->     List[str]:
            lowerCAmelCase  						=       self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_megatron_bert_for_masked_lm(*lowercase )





      def     _snake_case			(       self )		->     Union[str, Any]:
            lowerCAmelCase  						=       self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_megatron_bert_for_multiple_choice(*lowercase )





      def     _snake_case			(       self )		->     int:
            lowerCAmelCase  						=       self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_megatron_bert_for_next_sequence_prediction(*lowercase )





      def     _snake_case			(       self )		->     int:
            lowerCAmelCase  						=       self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_megatron_bert_for_pretraining(*lowercase )





      def     _snake_case			(       self )		->     Dict:
            lowerCAmelCase  						=       self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_megatron_bert_for_question_answering(*lowercase )





      def     _snake_case			(       self )		->     Optional[int]:
            lowerCAmelCase  						=       self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_megatron_bert_for_sequence_classification(*lowercase )





      def     _snake_case			(       self )		->     int:
            lowerCAmelCase  						=       self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_megatron_bert_for_token_classification(*lowercase )







def 			UpperCAmelCase__							(  SCREAMING_SNAKE_CASE     :						List[Any]			):







      '''simple docstring'''




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


SCREAMING_SNAKE_CASE__					=   1e-4






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





      @slow
      @unittest.skip("""Model is not available.""" )
      def     _snake_case			(       self )		->     Any:
            lowerCAmelCase  						=       """nvidia/megatron-bert-uncased-345m"""
            if "MYDIR" in os.environ:
                  lowerCAmelCase  						=       os.path.join(os.environ["""MYDIR"""]      ,						lowercase )
            lowerCAmelCase  						=       MegatronBertModel.from_pretrained(lowercase )
            model.to(lowercase )
            model.half()
            lowerCAmelCase  						=       _long_tensor([[101, 7_110, 1_005, 1_056, 2_023, 11_333, 17_413, 1_029, 102]] )
            with torch.no_grad():
                  lowerCAmelCase  						=       model(lowercase )[0]
            lowerCAmelCase  						=       torch.Size((1, 9, 1_024) )
            self.assertEqual(output.shape      ,						lowercase )

            lowerCAmelCase  						=       [-0.6_040, -0.2_517, -0.1_025, 0.3_420, -0.6_758, -0.0_017, -0.1_089, -0.1_990, 0.5_728]
            for ii in range(3 ):
                  for jj in range(3 ):
                        lowerCAmelCase  						=       output[0, ii, jj]
                        lowerCAmelCase  						=       expected[3 * ii + jj]
                        lowerCAmelCase  						=       """ii={} jj={} a={} b={}""".format(lowercase      ,						lowercase      ,						lowercase      ,						lowercase )
                        self.assertTrue(math.isclose(lowercase      ,						lowercase      ,						rel_tol=lowercase      ,						abs_tol=lowercase )      ,						msg=lowercase )











'''simple docstring'''






import warnings
from contextlib import contextmanager

from ....processing_utils import ProcessorMixin

class 				SCREAMING_SNAKE_CASE(							A__			):







	"""simple docstring"""






	lowerCamelCase__			 =       """MCTCTFeatureExtractor"""
	lowerCamelCase__			 =       """AutoTokenizer"""







	def __init__(						self			:     Dict		,			__snake_case			:     Optional[int]		,			__snake_case			:     List[str]		)			->		str:
								super().__init__(__snake_case		,			__snake_case		)
								UpperCAmelCase :       List[Any] 							=		self.feature_extractor
								UpperCAmelCase :       Union[str, Any] 							=		False







	def __call__(						self			:     Any		,			*__snake_case			:     List[str]		,			**__snake_case			:     Any		)			->		List[Any]:
								# For backward compatibility
								if self._in_target_context_manager:
															return self.current_processor(*__snake_case		,			**__snake_case		)

								if "raw_speech" in kwargs:
															warnings.warn('''Using `raw_speech` as a keyword argument is deprecated. Use `audio` instead.'''		)
															UpperCAmelCase :       int 							=		kwargs.pop('''raw_speech'''		)
								else:
															UpperCAmelCase :       Union[str, Any] 							=		kwargs.pop('''audio'''		,			__snake_case		)
								UpperCAmelCase :       Optional[Any] 							=		kwargs.pop('''sampling_rate'''		,			__snake_case		)
								UpperCAmelCase :       Dict 							=		kwargs.pop('''text'''		,			__snake_case		)
								if len(__snake_case		) > 0:
															UpperCAmelCase :       Any 							=		args[0]
															UpperCAmelCase :       Optional[int] 							=		args[1:]

								if audio is None and text is None:
															raise ValueError('''You need to specify either an `audio` or `text` input to process.'''		)

								if audio is not None:
															UpperCAmelCase :       List[str] 							=		self.feature_extractor(__snake_case		,			*__snake_case		,			sampling_rate=__snake_case		,			**__snake_case		)
								if text is not None:
															UpperCAmelCase :       int 							=		self.tokenizer(__snake_case		,			**__snake_case		)

								if text is None:
															return inputs
								elif audio is None:
															return encodings
								else:
															UpperCAmelCase :       str 							=		encodings['''input_ids''']
															return inputs







	def      A	(						self			:     List[Any]		,			*__snake_case			:     List[Any]		,			**__snake_case			:     List[Any]		)			->		str:
								return self.tokenizer.batch_decode(*__snake_case		,			**__snake_case		)







	def      A	(						self			:     List[Any]		,			*__snake_case			:     int		,			**__snake_case			:     Optional[int]		)			->		Any:
								# For backward compatibility
								if self._in_target_context_manager:
															return self.current_processor.pad(*__snake_case		,			**__snake_case		)

								UpperCAmelCase :       List[Any] 							=		kwargs.pop('''input_features'''		,			__snake_case		)
								UpperCAmelCase :       Optional[Any] 							=		kwargs.pop('''labels'''		,			__snake_case		)
								if len(__snake_case		) > 0:
															UpperCAmelCase :       List[str] 							=		args[0]
															UpperCAmelCase :       List[Any] 							=		args[1:]

								if input_features is not None:
															UpperCAmelCase :       Tuple 							=		self.feature_extractor.pad(__snake_case		,			*__snake_case		,			**__snake_case		)
								if labels is not None:
															UpperCAmelCase :       Optional[int] 							=		self.tokenizer.pad(__snake_case		,			**__snake_case		)

								if labels is None:
															return input_features
								elif input_features is None:
															return labels
								else:
															UpperCAmelCase :       List[str] 							=		labels['''input_ids''']
															return input_features







	def      A	(						self			:     Union[str, Any]		,			*__snake_case			:     Optional[Any]		,			**__snake_case			:     Optional[int]		)			->		Optional[Any]:
								return self.tokenizer.decode(*__snake_case		,			**__snake_case		)







	@contextmanager
	def      A	(						self			:     Any		)			->		Optional[int]:
								warnings.warn(
								    '''`as_target_processor` is deprecated and will be removed in v5 of Transformers. You can process your '''
								    '''labels by using the argument `text` of the regular `__call__` method (either in the same call as '''
								    '''your audio inputs, or in a separate call.'''		)
								UpperCAmelCase :       Dict 							=		True
								UpperCAmelCase :       List[Any] 							=		self.tokenizer
								yield
								UpperCAmelCase :       Tuple 							=		self.feature_extractor
								UpperCAmelCase :       List[Any] 							=		False





