"""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 				_A		:

      def __init__(		self	, __lowerCAmelCase	, __lowerCAmelCase=13	, __lowerCAmelCase=7	, __lowerCAmelCase=True	, __lowerCAmelCase=True	, __lowerCAmelCase=True	, __lowerCAmelCase=True	, __lowerCAmelCase=99	, __lowerCAmelCase=64	, __lowerCAmelCase=32	, __lowerCAmelCase=5	, __lowerCAmelCase=4	, __lowerCAmelCase=37	, __lowerCAmelCase="gelu"	, __lowerCAmelCase=0.1	, __lowerCAmelCase=0.1	, __lowerCAmelCase=512	, __lowerCAmelCase=16	, __lowerCAmelCase=2	, __lowerCAmelCase=0.0_2	, __lowerCAmelCase=3	, __lowerCAmelCase=4	, __lowerCAmelCase=None	, ):







            """simple docstring"""



            lowercase														=			parent
            lowercase														=			batch_size
            lowercase														=			seq_length
            lowercase														=			is_training
            lowercase														=			use_input_mask
            lowercase														=			use_token_type_ids
            lowercase														=			use_labels
            lowercase														=			vocab_size
            lowercase														=			hidden_size
            lowercase														=			embedding_size
            lowercase														=			num_hidden_layers
            lowercase														=			num_attention_heads
            lowercase														=			intermediate_size
            lowercase														=			hidden_act
            lowercase														=			hidden_dropout_prob
            lowercase														=			attention_probs_dropout_prob
            lowercase														=			max_position_embeddings
            lowercase														=			type_vocab_size
            lowercase														=			type_sequence_label_size
            lowercase														=			initializer_range
            lowercase														=			num_labels
            lowercase														=			num_choices
            lowercase														=			scope

      def        A__ (		self		):







            """simple docstring"""



            lowercase														=			ids_tensor([self.batch_size, self.seq_length]	, self.vocab_size		)

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

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

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

            lowercase														=			self.get_config()

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

      def        A__ (		self		):







            """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        A__ (		self	, __lowerCAmelCase	, __lowerCAmelCase	, __lowerCAmelCase	, __lowerCAmelCase	, __lowerCAmelCase	, __lowerCAmelCase	, __lowerCAmelCase		):







            """simple docstring"""



            lowercase														=			MegatronBertModel(config=__lowerCAmelCase		)
            model.to(__lowerCAmelCase		)
            model.eval()
            lowercase														=			model(__lowerCAmelCase	, attention_mask=__lowerCAmelCase	, token_type_ids=__lowerCAmelCase		)
            lowercase														=			model(__lowerCAmelCase	, token_type_ids=__lowerCAmelCase		)
            lowercase														=			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        A__ (		self	, __lowerCAmelCase	, __lowerCAmelCase	, __lowerCAmelCase	, __lowerCAmelCase	, __lowerCAmelCase	, __lowerCAmelCase	, __lowerCAmelCase		):







            """simple docstring"""



            lowercase														=			MegatronBertForMaskedLM(config=__lowerCAmelCase		)
            model.to(__lowerCAmelCase		)
            model.eval()
            lowercase														=			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        A__ (		self	, __lowerCAmelCase	, __lowerCAmelCase	, __lowerCAmelCase	, __lowerCAmelCase	, __lowerCAmelCase	, __lowerCAmelCase	, __lowerCAmelCase		):







            """simple docstring"""



            lowercase														=			MegatronBertForCausalLM(config=__lowerCAmelCase		)
            model.to(__lowerCAmelCase		)
            model.eval()
            lowercase														=			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        A__ (		self	, __lowerCAmelCase	, __lowerCAmelCase	, __lowerCAmelCase	, __lowerCAmelCase	, __lowerCAmelCase	, __lowerCAmelCase	, __lowerCAmelCase		):







            """simple docstring"""



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

      def        A__ (		self	, __lowerCAmelCase	, __lowerCAmelCase	, __lowerCAmelCase	, __lowerCAmelCase	, __lowerCAmelCase	, __lowerCAmelCase	, __lowerCAmelCase		):







            """simple docstring"""



            lowercase														=			MegatronBertForPreTraining(config=__lowerCAmelCase		)
            model.to(__lowerCAmelCase		)
            model.eval()
            lowercase														=			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        A__ (		self	, __lowerCAmelCase	, __lowerCAmelCase	, __lowerCAmelCase	, __lowerCAmelCase	, __lowerCAmelCase	, __lowerCAmelCase	, __lowerCAmelCase		):







            """simple docstring"""



            lowercase														=			MegatronBertForQuestionAnswering(config=__lowerCAmelCase		)
            model.to(__lowerCAmelCase		)
            model.eval()
            lowercase														=			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        A__ (		self	, __lowerCAmelCase	, __lowerCAmelCase	, __lowerCAmelCase	, __lowerCAmelCase	, __lowerCAmelCase	, __lowerCAmelCase	, __lowerCAmelCase		):







            """simple docstring"""



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

      def        A__ (		self	, __lowerCAmelCase	, __lowerCAmelCase	, __lowerCAmelCase	, __lowerCAmelCase	, __lowerCAmelCase	, __lowerCAmelCase	, __lowerCAmelCase		):







            """simple docstring"""



            lowercase														=			self.num_labels
            lowercase														=			MegatronBertForTokenClassification(config=__lowerCAmelCase		)
            model.to(__lowerCAmelCase		)
            model.eval()
            lowercase														=			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        A__ (		self	, __lowerCAmelCase	, __lowerCAmelCase	, __lowerCAmelCase	, __lowerCAmelCase	, __lowerCAmelCase	, __lowerCAmelCase	, __lowerCAmelCase		):







            """simple docstring"""



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







      def        A__ (		self		):







            """simple docstring"""



            lowercase														=			self.prepare_config_and_inputs()
            (
                (
                lowercase
            )							,							(
                lowercase
            )							,							(
                lowercase
            )							,							(
                lowercase
            )							,							(
                lowercase
            )							,							(
                lowercase
            )							,							(
                lowercase
            )							,							
            )														=			config_and_inputs
            lowercase														=			{"""input_ids""": input_ids, """token_type_ids""": token_type_ids, """attention_mask""": input_mask}
            return config, inputs_dict






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

      def        A__ (		self	, __lowerCAmelCase	, __lowerCAmelCase	, __lowerCAmelCase=False		):







            """simple docstring"""



            lowercase														=			super()._prepare_for_class(__lowerCAmelCase	, __lowerCAmelCase	, return_labels=__lowerCAmelCase		)

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

      def        A__ (		self		):







            """simple docstring"""



            lowercase														=			MegatronBertModelTester(self		)
            lowercase														=			ConfigTester(self	, config_class=__lowerCAmelCase	, hidden_size=37		)

      def        A__ (		self		):







            """simple docstring"""



            self.config_tester.run_common_tests()

      def        A__ (		self		):







            """simple docstring"""



            lowercase														=			self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_megatron_bert_model(*__lowerCAmelCase		)

      def        A__ (		self		):







            """simple docstring"""



            lowercase														=			self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_megatron_bert_for_masked_lm(*__lowerCAmelCase		)

      def        A__ (		self		):







            """simple docstring"""



            lowercase														=			self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_megatron_bert_for_multiple_choice(*__lowerCAmelCase		)

      def        A__ (		self		):







            """simple docstring"""



            lowercase														=			self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_megatron_bert_for_next_sequence_prediction(*__lowerCAmelCase		)

      def        A__ (		self		):







            """simple docstring"""



            lowercase														=			self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_megatron_bert_for_pretraining(*__lowerCAmelCase		)

      def        A__ (		self		):







            """simple docstring"""



            lowercase														=			self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_megatron_bert_for_question_answering(*__lowerCAmelCase		)

      def        A__ (		self		):







            """simple docstring"""



            lowercase														=			self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_megatron_bert_for_sequence_classification(*__lowerCAmelCase		)







      def        A__ (		self		):







            """simple docstring"""



            lowercase														=			self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_megatron_bert_for_token_classification(*__lowerCAmelCase		)






def     UpperCAmelCase__ (       lowerCAmelCase__  :List[Any]	)				->					Any:
      '''simple docstring'''


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


__lowerCAmelCase					: int        =1E-4






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

      @slow
      @unittest.skip("""Model is not available."""		)
      def        A__ (		self		):







            """simple docstring"""



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

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




"""simple docstring"""






from __future__ import annotations


def     UpperCAmelCase__ (       lowerCAmelCase__  :list[float]      ,	lowerCAmelCase__  :list[float]	)				->					float:
      '''simple docstring'''


      lowercase														=			sorted(numsa + numsa	)
      lowercase							,							lowercase														=			divmod(len(lowerCAmelCase__	)      ,	2	)
      if mod == 1:
            return all_numbers[div]
      else:
            return (all_numbers[div] + all_numbers[div - 1]) / 2


if __name__ == "__main__":
       import doctest

       doctest.testmod()
       __lowerCAmelCase					: List[Any]        =[float(x) for x in input("""Enter the elements of first array: """).split()]
       __lowerCAmelCase					: List[Any]        =[float(x) for x in input("""Enter the elements of second array: """).split()]
       print(F"""The median of two arrays is: {median_of_two_arrays(array_a, array_a)}""")



