



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




       """simple docstring"""






       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     ,):

             """simple docstring"""





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



       def  A_   (     self):

             """simple docstring"""





             UpperCAmelCase_  :	Union[str, Any]				      =		ids_tensor([self.batch_size, self.seq_length]     ,self.vocab_size)

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

             UpperCAmelCase_  :	Optional[int]				      =		None
             if self.use_token_type_ids:
                   UpperCAmelCase_  :	List[Any]				      =		ids_tensor([self.batch_size, self.seq_length]     ,self.type_vocab_size)

             UpperCAmelCase_  :	Tuple				      =		None
             UpperCAmelCase_  :	Optional[int]				      =		None
             UpperCAmelCase_  :	str				      =		None
             if self.use_labels:
                   UpperCAmelCase_  :	List[Any]				      =		ids_tensor([self.batch_size]     ,self.type_sequence_label_size)
                   UpperCAmelCase_  :	Union[str, Any]				      =		ids_tensor([self.batch_size, self.seq_length]     ,self.num_labels)
                   UpperCAmelCase_  :	str				      =		ids_tensor([self.batch_size]     ,self.num_choices)

             UpperCAmelCase_  :	Any				      =		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=lowercase     ,initializer_range=self.initializer_range     ,)



       def  A_   (     self     ,lowercase     ,lowercase     ,lowercase     ,lowercase     ,lowercase     ,lowercase     ,lowercase):

             """simple docstring"""





             UpperCAmelCase_  :	str				      =		MegatronBertModel(config=lowercase)
             model.to(lowercase)
             model.eval()
             UpperCAmelCase_  :	List[Any]				      =		model(lowercase     ,attention_mask=lowercase     ,token_type_ids=lowercase)
             UpperCAmelCase_  :	Dict				      =		model(lowercase     ,token_type_ids=lowercase)
             UpperCAmelCase_  :	int				      =		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  A_   (     self     ,lowercase     ,lowercase     ,lowercase     ,lowercase     ,lowercase     ,lowercase     ,lowercase):

             """simple docstring"""





             UpperCAmelCase_  :	Union[str, Any]				      =		MegatronBertForMaskedLM(config=lowercase)
             model.to(lowercase)
             model.eval()
             UpperCAmelCase_  :	int				      =		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  A_   (     self     ,lowercase     ,lowercase     ,lowercase     ,lowercase     ,lowercase     ,lowercase     ,lowercase):

             """simple docstring"""





             UpperCAmelCase_  :	Optional[Any]				      =		MegatronBertForCausalLM(config=lowercase)
             model.to(lowercase)
             model.eval()
             UpperCAmelCase_  :	str				      =		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  A_   (     self     ,lowercase     ,lowercase     ,lowercase     ,lowercase     ,lowercase     ,lowercase     ,lowercase):

             """simple docstring"""





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



       def  A_   (     self     ,lowercase     ,lowercase     ,lowercase     ,lowercase     ,lowercase     ,lowercase     ,lowercase):

             """simple docstring"""





             UpperCAmelCase_  :	Optional[int]				      =		MegatronBertForPreTraining(config=lowercase)
             model.to(lowercase)
             model.eval()
             UpperCAmelCase_  :	Optional[Any]				      =		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  A_   (     self     ,lowercase     ,lowercase     ,lowercase     ,lowercase     ,lowercase     ,lowercase     ,lowercase):

             """simple docstring"""





             UpperCAmelCase_  :	Any				      =		MegatronBertForQuestionAnswering(config=lowercase)
             model.to(lowercase)
             model.eval()
             UpperCAmelCase_  :	Optional[Any]				      =		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  A_   (     self     ,lowercase     ,lowercase     ,lowercase     ,lowercase     ,lowercase     ,lowercase     ,lowercase):

             """simple docstring"""





             UpperCAmelCase_  :	Any				      =		self.num_labels
             UpperCAmelCase_  :	List[Any]				      =		MegatronBertForSequenceClassification(lowercase)
             model.to(lowercase)
             model.eval()
             UpperCAmelCase_  :	int				      =		model(lowercase     ,attention_mask=lowercase     ,token_type_ids=lowercase     ,labels=lowercase)
             self.parent.assertEqual(result.logits.shape     ,(self.batch_size, self.num_labels))



       def  A_   (     self     ,lowercase     ,lowercase     ,lowercase     ,lowercase     ,lowercase     ,lowercase     ,lowercase):

             """simple docstring"""





             UpperCAmelCase_  :	Optional[Any]				      =		self.num_labels
             UpperCAmelCase_  :	int				      =		MegatronBertForTokenClassification(config=lowercase)
             model.to(lowercase)
             model.eval()
             UpperCAmelCase_  :	List[Any]				      =		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  A_   (     self     ,lowercase     ,lowercase     ,lowercase     ,lowercase     ,lowercase     ,lowercase     ,lowercase):

             """simple docstring"""





             UpperCAmelCase_  :	int				      =		self.num_choices
             UpperCAmelCase_  :	Any				      =		MegatronBertForMultipleChoice(config=lowercase)
             model.to(lowercase)
             model.eval()
             UpperCAmelCase_  :	Union[str, Any]				      =		input_ids.unsqueeze(1).expand(-1     ,self.num_choices     ,-1).contiguous()
             UpperCAmelCase_  :	Tuple				      =		token_type_ids.unsqueeze(1).expand(-1     ,self.num_choices     ,-1).contiguous()
             UpperCAmelCase_  :	List[Any]				      =		input_mask.unsqueeze(1).expand(-1     ,self.num_choices     ,-1).contiguous()
             UpperCAmelCase_  :	Dict				      =		model(
                 lowercase     ,attention_mask=lowercase     ,token_type_ids=lowercase     ,labels=lowercase     ,)
             self.parent.assertEqual(result.logits.shape     ,(self.batch_size, self.num_choices))



       def  A_   (     self):

             """simple docstring"""





             UpperCAmelCase_  :	List[Any]				      =		self.prepare_config_and_inputs()
             (
                 (
                 UpperCAmelCase_
             )      ,      (
                 UpperCAmelCase_
             )      ,      (
                 UpperCAmelCase_
             )      ,      (
                 UpperCAmelCase_
             )      ,      (
                 UpperCAmelCase_
             )      ,      (
                 UpperCAmelCase_
             )      ,      (
                 UpperCAmelCase_
             )      ,      
             )  :	List[str]				      =		config_and_inputs
             UpperCAmelCase_  :	Optional[Any]				      =		{"input_ids": input_ids, "token_type_ids": token_type_ids, "attention_mask": input_mask}
             return config, inputs_dict


@require_torch
class   snake_case_ (lowercase__ , lowercase__ , 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     ,lowercase     ,lowercase     ,lowercase=False):

             """simple docstring"""





             UpperCAmelCase_  :	Tuple				      =		super()._prepare_for_class(lowercase     ,lowercase     ,return_labels=lowercase)

             if return_labels:
                   if model_class in get_values(lowercase):
                         UpperCAmelCase_  :	Union[str, Any]				      =		torch.zeros(
                             (self.model_tester.batch_size, self.model_tester.seq_length)     ,dtype=torch.long     ,device=lowercase)
                         UpperCAmelCase_  :	Any				      =		torch.zeros(
                             self.model_tester.batch_size     ,dtype=torch.long     ,device=lowercase)
             return inputs_dict



       def  A_   (     self):

             """simple docstring"""





             UpperCAmelCase_  :	Tuple				      =		MegatronBertModelTester(self)
             UpperCAmelCase_  :	Dict				      =		ConfigTester(self     ,config_class=lowercase     ,hidden_size=37)



       def  A_   (     self):

             """simple docstring"""





             self.config_tester.run_common_tests()



       def  A_   (     self):

             """simple docstring"""





             UpperCAmelCase_  :	Any				      =		self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_megatron_bert_model(*lowercase)



       def  A_   (     self):

             """simple docstring"""





             UpperCAmelCase_  :	Optional[Any]				      =		self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_megatron_bert_for_masked_lm(*lowercase)



       def  A_   (     self):

             """simple docstring"""





             UpperCAmelCase_  :	Dict				      =		self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_megatron_bert_for_multiple_choice(*lowercase)



       def  A_   (     self):

             """simple docstring"""





             UpperCAmelCase_  :	List[Any]				      =		self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_megatron_bert_for_next_sequence_prediction(*lowercase)



       def  A_   (     self):

             """simple docstring"""





             UpperCAmelCase_  :	List[Any]				      =		self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_megatron_bert_for_pretraining(*lowercase)



       def  A_   (     self):

             """simple docstring"""





             UpperCAmelCase_  :	Any				      =		self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_megatron_bert_for_question_answering(*lowercase)



       def  A_   (     self):

             """simple docstring"""





             UpperCAmelCase_  :	int				      =		self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_megatron_bert_for_sequence_classification(*lowercase)



       def  A_   (     self):

             """simple docstring"""





             UpperCAmelCase_  :	int				      =		self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_megatron_bert_for_token_classification(*lowercase)







def    _snake_case (       __snake_case    )     ->     Dict:






      '''simple docstring'''







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


__lowerCamelCase         =			1e-4





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




       """simple docstring"""






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

             """simple docstring"""





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

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





from ...configuration_utils import PretrainedConfig
from ...utils import logging
from ...utils.backbone_utils import BackboneConfigMixin, get_aligned_output_features_output_indices


__lowerCamelCase         =			logging.get_logger(__name__)

__lowerCamelCase         =			{
    '''microsoft/focalnet-tiny''': '''https://huggingface.co/microsoft/focalnet-tiny/resolve/main/config.json''',
}


class   snake_case_ (lowercase__ , lowercase__  ):




       """simple docstring"""






       _lowerCamelCase			    =				"""focalnet"""



       def __init__(     self     ,lowercase=224     ,lowercase=4     ,lowercase=3     ,lowercase=96     ,lowercase=False     ,lowercase=[192, 384, 768, 768]     ,lowercase=[2, 2, 6, 2]     ,lowercase=[2, 2, 2, 2]     ,lowercase=[3, 3, 3, 3]     ,lowercase="gelu"     ,lowercase=4.0     ,lowercase=0.0     ,lowercase=0.1     ,lowercase=False     ,lowercase=1E-4     ,lowercase=False     ,lowercase=False     ,lowercase=False     ,lowercase=0.02     ,lowercase=1E-5     ,lowercase=32     ,lowercase=None     ,lowercase=None     ,**lowercase     ,):

             """simple docstring"""





             super().__init__(**lowercase)

             UpperCAmelCase_  :	int				      =		image_size
             UpperCAmelCase_  :	Optional[Any]				      =		patch_size
             UpperCAmelCase_  :	Dict				      =		num_channels
             UpperCAmelCase_  :	Dict				      =		embed_dim
             UpperCAmelCase_  :	Optional[int]				      =		use_conv_embed
             UpperCAmelCase_  :	int				      =		hidden_sizes
             UpperCAmelCase_  :	Optional[int]				      =		depths
             UpperCAmelCase_  :	Optional[Any]				      =		focal_levels
             UpperCAmelCase_  :	Tuple				      =		focal_windows
             UpperCAmelCase_  :	int				      =		hidden_act
             UpperCAmelCase_  :	Any				      =		mlp_ratio
             UpperCAmelCase_  :	List[Any]				      =		hidden_dropout_prob
             UpperCAmelCase_  :	Tuple				      =		drop_path_rate
             UpperCAmelCase_  :	List[str]				      =		use_layerscale
             UpperCAmelCase_  :	List[str]				      =		layerscale_value
             UpperCAmelCase_  :	List[str]				      =		use_post_layernorm
             UpperCAmelCase_  :	Dict				      =		use_post_layernorm_in_modulation
             UpperCAmelCase_  :	int				      =		normalize_modulator
             UpperCAmelCase_  :	Dict				      =		initializer_range
             UpperCAmelCase_  :	str				      =		layer_norm_eps
             UpperCAmelCase_  :	Union[str, Any]				      =		encoder_stride
             UpperCAmelCase_  :	str				      =		["stem"] + [F"""stage{idx}""" for idx in range(1     ,len(self.depths) + 1)]
             UpperCAmelCase_      ,      UpperCAmelCase_  :	Any				      =		get_aligned_output_features_output_indices(
                 out_features=lowercase     ,out_indices=lowercase     ,stage_names=self.stage_names)
