






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


				'''simple docstring'''

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

					'''simple docstring'''




					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    a__				(					self     :     str     )							->							Union[str, Any]:

					'''simple docstring'''




					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    a__				(					self     :     str     )							->							List[Any]:

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

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

					'''simple docstring'''




					lowerCamelCase__  =     MegatronBertModel(config=UpperCamelCase__     )
					model.to(UpperCamelCase__     )
					model.eval()
					lowerCamelCase__  =     model(UpperCamelCase__				,     attention_mask=UpperCamelCase__				,     token_type_ids=UpperCamelCase__     )
					lowerCamelCase__  =     model(UpperCamelCase__				,     token_type_ids=UpperCamelCase__     )
					lowerCamelCase__  =     model(UpperCamelCase__     )

					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    a__				(					self     :     int				,     __lowerCamelCase     :     Union[str, Any]				,     __lowerCamelCase     :     List[str]				,     __lowerCamelCase     :     str				,     __lowerCamelCase     :     Optional[int]				,     __lowerCamelCase     :     int				,     __lowerCamelCase     :     Tuple				,     __lowerCamelCase     :     List[str]     )							->							Any:

					'''simple docstring'''




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

				def    a__				(					self     :     Any				,     __lowerCamelCase     :     Union[str, Any]				,     __lowerCamelCase     :     Tuple				,     __lowerCamelCase     :     int				,     __lowerCamelCase     :     Union[str, Any]				,     __lowerCamelCase     :     Tuple				,     __lowerCamelCase     :     str				,     __lowerCamelCase     :     Any     )							->							List[Any]:

					'''simple docstring'''




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

				def    a__				(					self     :     Union[str, Any]				,     __lowerCamelCase     :     Optional[Any]				,     __lowerCamelCase     :     Dict				,     __lowerCamelCase     :     Dict				,     __lowerCamelCase     :     List[str]				,     __lowerCamelCase     :     Dict				,     __lowerCamelCase     :     Any				,     __lowerCamelCase     :     int     )							->							str:

					'''simple docstring'''




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

				def    a__				(					self     :     Any				,     __lowerCamelCase     :     int				,     __lowerCamelCase     :     Tuple				,     __lowerCamelCase     :     Optional[Any]				,     __lowerCamelCase     :     Optional[Any]				,     __lowerCamelCase     :     Union[str, Any]				,     __lowerCamelCase     :     Union[str, Any]				,     __lowerCamelCase     :     List[Any]     )							->							Any:

					'''simple docstring'''




					lowerCamelCase__  =     MegatronBertForPreTraining(config=UpperCamelCase__     )
					model.to(UpperCamelCase__     )
					model.eval()
					lowerCamelCase__  =     model(
					    UpperCamelCase__				,     attention_mask=UpperCamelCase__				,     token_type_ids=UpperCamelCase__				,     labels=UpperCamelCase__				,     next_sentence_label=UpperCamelCase__				,     )
					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    a__				(					self     :     Optional[Any]				,     __lowerCamelCase     :     Dict				,     __lowerCamelCase     :     Tuple				,     __lowerCamelCase     :     Optional[int]				,     __lowerCamelCase     :     List[str]				,     __lowerCamelCase     :     List[str]				,     __lowerCamelCase     :     List[Any]				,     __lowerCamelCase     :     List[str]     )							->							Dict:

					'''simple docstring'''




					lowerCamelCase__  =     MegatronBertForQuestionAnswering(config=UpperCamelCase__     )
					model.to(UpperCamelCase__     )
					model.eval()
					lowerCamelCase__  =     model(
					    UpperCamelCase__				,     attention_mask=UpperCamelCase__				,     token_type_ids=UpperCamelCase__				,     start_positions=UpperCamelCase__				,     end_positions=UpperCamelCase__				,     )
					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    a__				(					self     :     int				,     __lowerCamelCase     :     Dict				,     __lowerCamelCase     :     Optional[int]				,     __lowerCamelCase     :     List[Any]				,     __lowerCamelCase     :     Dict				,     __lowerCamelCase     :     Optional[int]				,     __lowerCamelCase     :     Optional[int]				,     __lowerCamelCase     :     Tuple     )							->							Dict:

					'''simple docstring'''




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

				def    a__				(					self     :     str				,     __lowerCamelCase     :     Union[str, Any]				,     __lowerCamelCase     :     int				,     __lowerCamelCase     :     Tuple				,     __lowerCamelCase     :     str				,     __lowerCamelCase     :     Optional[int]				,     __lowerCamelCase     :     Optional[int]				,     __lowerCamelCase     :     Optional[Any]     )							->							Dict:

					'''simple docstring'''




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

				def    a__				(					self     :     List[str]				,     __lowerCamelCase     :     Dict				,     __lowerCamelCase     :     int				,     __lowerCamelCase     :     List[str]				,     __lowerCamelCase     :     int				,     __lowerCamelCase     :     Union[str, Any]				,     __lowerCamelCase     :     int				,     __lowerCamelCase     :     Any     )							->							Optional[int]:

					'''simple docstring'''




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




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

					'''simple docstring'''




					lowerCamelCase__  =     self.prepare_config_and_inputs()
					(
					    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       (							__A ,       __A ,       unittest.TestCase      ):


				'''simple docstring'''

				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    a__				(					self     :     Union[str, Any]				,     __lowerCamelCase     :     List[Any]				,     __lowerCamelCase     :     str				,     __lowerCamelCase     :     List[Any]=False     )							->							Any:

					'''simple docstring'''




					lowerCamelCase__  =     super()._prepare_for_class(UpperCamelCase__				,     UpperCamelCase__				,     return_labels=UpperCamelCase__     )

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

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

					'''simple docstring'''




					lowerCamelCase__  =     MegatronBertModelTester(self     )
					lowerCamelCase__  =     ConfigTester(self				,     config_class=UpperCamelCase__				,     hidden_size=37     )

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

					'''simple docstring'''




					self.config_tester.run_common_tests()

				def    a__				(					self     :     List[Any]     )							->							List[Any]:

					'''simple docstring'''




					lowerCamelCase__  =     self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_megatron_bert_model(*UpperCamelCase__     )

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

					'''simple docstring'''




					lowerCamelCase__  =     self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_megatron_bert_for_masked_lm(*UpperCamelCase__     )

				def    a__				(					self     :     Dict     )							->							str:

					'''simple docstring'''




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

				def    a__				(					self     :     int     )							->							int:

					'''simple docstring'''




					lowerCamelCase__  =     self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_megatron_bert_for_next_sequence_prediction(*UpperCamelCase__     )

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

					'''simple docstring'''




					lowerCamelCase__  =     self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_megatron_bert_for_pretraining(*UpperCamelCase__     )

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

					'''simple docstring'''




					lowerCamelCase__  =     self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_megatron_bert_for_question_answering(*UpperCamelCase__     )

				def    a__				(					self     :     Tuple     )							->							Tuple:

					'''simple docstring'''




					lowerCamelCase__  =     self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_megatron_bert_for_sequence_classification(*UpperCamelCase__     )




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

					'''simple docstring'''




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






def    lowerCamelCase_     (    lowercase__):
	return torch.tensor(
	    lowerCAmelCase__				,     dtype=torch.long				,     device=lowerCAmelCase__				,     )


__A							:							str  							=		1E-4
@require_torch
@require_sentencepiece
@require_tokenizers
class 				lowercase       (							unittest.TestCase      ):


				'''simple docstring'''

				@slow
				@unittest.skip("Model is not available."     )
				def    a__				(					self     :     Any     )							->							Tuple:

					'''simple docstring'''




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

					lowerCamelCase__  =     [-0.6_0_4_0, -0.2_5_1_7, -0.1_0_2_5, 0.3_4_2_0, -0.6_7_5_8, -0.0_0_1_7, -0.1_0_8_9, -0.1_9_9_0, 0.5_7_2_8]
					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(UpperCamelCase__				,     UpperCamelCase__				,     UpperCamelCase__				,     UpperCamelCase__     )
							self.assertTrue(math.isclose(UpperCamelCase__				,     UpperCamelCase__				,     rel_tol=UpperCamelCase__				,     abs_tol=UpperCamelCase__     )				,     msg=UpperCamelCase__     )









'''simple docstring'''





from __future__ import annotations
class 				lowercase       :


    '''simple docstring'''

    def __init__(					self     :     Optional[int]				,     __lowerCamelCase     :     int     )							->							None:

     '''simple docstring'''




     lowerCamelCase__  =     order

     # a_{0} ... a_{k}
     lowerCamelCase__  =     [1.0] + [0.0] * order
     # b_{0} ... b_{k}
     lowerCamelCase__  =     [1.0] + [0.0] * order

     # x[n-1] ... x[n-k]
     lowerCamelCase__  =     [0.0] * self.order
     # y[n-1] ... y[n-k]
     lowerCamelCase__  =     [0.0] * self.order

    def    a__				(					self     :     Dict				,     __lowerCamelCase     :     list[float]				,     __lowerCamelCase     :     list[float]     )							->							None:

     '''simple docstring'''




     if len(__lowerCamelCase     ) < self.order:
      lowerCamelCase__  =     [1.0, *a_coeffs]

     if len(__lowerCamelCase     ) != self.order + 1:
      lowerCamelCase__  =     (
          f'''Expected a_coeffs to have {self.order + 1} elements '''
          f'''for {self.order}-order filter, got {len(__lowerCamelCase     )}'''
      )
      raise ValueError(__lowerCamelCase     )

     if len(__lowerCamelCase     ) != self.order + 1:
      lowerCamelCase__  =     (
          f'''Expected b_coeffs to have {self.order + 1} elements '''
          f'''for {self.order}-order filter, got {len(__lowerCamelCase     )}'''
      )
      raise ValueError(__lowerCamelCase     )

     lowerCamelCase__  =     a_coeffs
     lowerCamelCase__  =     b_coeffs




    def    a__				(					self     :     Dict				,     __lowerCamelCase     :     float     )							->							float:

     '''simple docstring'''




     lowerCamelCase__  =     0.0

     # Start at index 1 and do index 0 at the end.
     for i in range(1				,     self.order + 1     ):
      result += (
          self.b_coeffs[i] * self.input_history[i - 1]
          - self.a_coeffs[i] * self.output_history[i - 1]
      )

     lowerCamelCase__  =     (result + self.b_coeffs[0] * sample) / self.a_coeffs[0]

     lowerCamelCase__  =     self.input_history[:-1]
     lowerCamelCase__  =     self.output_history[:-1]

     lowerCamelCase__  =     sample
     lowerCamelCase__  =     result

     return result

