





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 ,							A_ ,							A_=13 ,							A_=7 ,							A_=True ,							A_=True ,							A_=True ,							A_=True ,							A_=99 ,							A_=64 ,							A_=32 ,							A_=5 ,							A_=4 ,							A_=37 ,							A_="gelu" ,							A_=0.1 ,							A_=0.1 ,							A_=512 ,							A_=16 ,							A_=2 ,							A_=0.02 ,							A_=3 ,							A_=4 ,							A_=None ,							):
       '''simple docstring'''




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





      def  __UpperCamelCase(				self     ):
       '''simple docstring'''




       UpperCamelCase			:      Dict 			=     ids_tensor([self.batch_size, self.seq_length] ,							self.vocab_size     )

       UpperCamelCase			:      Dict 			=     None
       if self.use_input_mask:
        UpperCamelCase			:      List[Any] 			=     random_attention_mask([self.batch_size, self.seq_length]     )

       UpperCamelCase			:      List[Any] 			=     None
       if self.use_token_type_ids:
        UpperCamelCase			:      str 			=     ids_tensor([self.batch_size, self.seq_length] ,							self.type_vocab_size     )

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

       UpperCamelCase			:      str 			=     self.get_config()

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





      def  __UpperCamelCase(				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=A_ ,							initializer_range=self.initializer_range ,							)





      def  __UpperCamelCase(				self ,							A_ ,							A_ ,							A_ ,							A_ ,							A_ ,							A_ ,							A_     ):
       '''simple docstring'''




       UpperCamelCase			:      Dict 			=     MegatronBertModel(config=A_     )
       model.to(A_     )
       model.eval()
       UpperCamelCase			:      Any 			=     model(A_ ,							attention_mask=A_ ,							token_type_ids=A_     )
       UpperCamelCase			:      Optional[Any] 			=     model(A_ ,							token_type_ids=A_     )
       UpperCamelCase			:      Optional[Any] 			=     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 ,							A_ ,							A_ ,							A_ ,							A_ ,							A_ ,							A_ ,							A_     ):
       '''simple docstring'''




       UpperCamelCase			:      List[str] 			=     MegatronBertForMaskedLM(config=A_     )
       model.to(A_     )
       model.eval()
       UpperCamelCase			:      Tuple 			=     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 ,							A_ ,							A_ ,							A_ ,							A_ ,							A_ ,							A_ ,							A_     ):
       '''simple docstring'''




       UpperCamelCase			:      int 			=     MegatronBertForCausalLM(config=A_     )
       model.to(A_     )
       model.eval()
       UpperCamelCase			:      str 			=     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 ,							A_ ,							A_ ,							A_ ,							A_ ,							A_ ,							A_ ,							A_     ):
       '''simple docstring'''




       UpperCamelCase			:      Dict 			=     MegatronBertForNextSentencePrediction(config=A_     )
       model.to(A_     )
       model.eval()
       UpperCamelCase			:      Union[str, Any] 			=     model(
           A_ ,							attention_mask=A_ ,							token_type_ids=A_ ,							labels=A_ ,							)
       self.parent.assertEqual(result.logits.shape ,							(self.batch_size, 2)     )





      def  __UpperCamelCase(				self ,							A_ ,							A_ ,							A_ ,							A_ ,							A_ ,							A_ ,							A_     ):
       '''simple docstring'''




       UpperCamelCase			:      str 			=     MegatronBertForPreTraining(config=A_     )
       model.to(A_     )
       model.eval()
       UpperCamelCase			:      List[Any] 			=     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 ,							A_ ,							A_ ,							A_ ,							A_ ,							A_ ,							A_ ,							A_     ):
       '''simple docstring'''




       UpperCamelCase			:      Optional[Any] 			=     MegatronBertForQuestionAnswering(config=A_     )
       model.to(A_     )
       model.eval()
       UpperCamelCase			:      Union[str, Any] 			=     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 ,							A_ ,							A_ ,							A_ ,							A_ ,							A_ ,							A_ ,							A_     ):
       '''simple docstring'''




       UpperCamelCase			:      str 			=     self.num_labels
       UpperCamelCase			:      Any 			=     MegatronBertForSequenceClassification(A_     )
       model.to(A_     )
       model.eval()
       UpperCamelCase			:      List[str] 			=     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 ,							A_ ,							A_ ,							A_ ,							A_ ,							A_ ,							A_ ,							A_     ):
       '''simple docstring'''




       UpperCamelCase			:      str 			=     self.num_labels
       UpperCamelCase			:      int 			=     MegatronBertForTokenClassification(config=A_     )
       model.to(A_     )
       model.eval()
       UpperCamelCase			:      str 			=     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 ,							A_ ,							A_ ,							A_ ,							A_ ,							A_ ,							A_ ,							A_     ):
       '''simple docstring'''




       UpperCamelCase			:      List[str] 			=     self.num_choices
       UpperCamelCase			:      int 			=     MegatronBertForMultipleChoice(config=A_     )
       model.to(A_     )
       model.eval()
       UpperCamelCase			:      Optional[Any] 			=     input_ids.unsqueeze(1     ).expand(-1 ,							self.num_choices ,							-1     ).contiguous()
       UpperCamelCase			:      Optional[int] 			=     token_type_ids.unsqueeze(1     ).expand(-1 ,							self.num_choices ,							-1     ).contiguous()
       UpperCamelCase			:      Tuple 			=     input_mask.unsqueeze(1     ).expand(-1 ,							self.num_choices ,							-1     ).contiguous()
       UpperCamelCase			:      Optional[Any] 			=     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     ):
       '''simple docstring'''




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



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





      def  __UpperCamelCase(				self ,							A_ ,							A_ ,							A_=False     ):
       '''simple docstring'''




       UpperCamelCase			:      Any 			=     super()._prepare_for_class(A_ ,							A_ ,							return_labels=A_     )

       if return_labels:
        if model_class in get_values(A_     ):
         UpperCamelCase			:      Optional[Any] 			=     torch.zeros(
             (self.model_tester.batch_size, self.model_tester.seq_length) ,							dtype=torch.long ,							device=A_     )
         UpperCamelCase			:      Dict 			=     torch.zeros(
             self.model_tester.batch_size ,							dtype=torch.long ,							device=A_     )
       return inputs_dict





      def  __UpperCamelCase(				self     ):
       '''simple docstring'''




       UpperCamelCase			:      List[str] 			=     MegatronBertModelTester(self     )
       UpperCamelCase			:      str 			=     ConfigTester(self ,							config_class=A_ ,							hidden_size=37     )





      def  __UpperCamelCase(				self     ):
       '''simple docstring'''




       self.config_tester.run_common_tests()





      def  __UpperCamelCase(				self     ):
       '''simple docstring'''




       UpperCamelCase			:      Optional[Any] 			=     self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_megatron_bert_model(*A_     )





      def  __UpperCamelCase(				self     ):
       '''simple docstring'''




       UpperCamelCase			:      str 			=     self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_megatron_bert_for_masked_lm(*A_     )





      def  __UpperCamelCase(				self     ):
       '''simple docstring'''




       UpperCamelCase			:      Tuple 			=     self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_megatron_bert_for_multiple_choice(*A_     )





      def  __UpperCamelCase(				self     ):
       '''simple docstring'''




       UpperCamelCase			:      Dict 			=     self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_megatron_bert_for_next_sequence_prediction(*A_     )





      def  __UpperCamelCase(				self     ):
       '''simple docstring'''




       UpperCamelCase			:      Union[str, Any] 			=     self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_megatron_bert_for_pretraining(*A_     )





      def  __UpperCamelCase(				self     ):
       '''simple docstring'''




       UpperCamelCase			:      List[Any] 			=     self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_megatron_bert_for_question_answering(*A_     )





      def  __UpperCamelCase(				self     ):
       '''simple docstring'''




       UpperCamelCase			:      Optional[Any] 			=     self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_megatron_bert_for_sequence_classification(*A_     )





      def  __UpperCamelCase(				self     ):
       '''simple docstring'''




       UpperCamelCase			:      str 			=     self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_megatron_bert_for_token_classification(*A_     )



def 	A_       (  _lowerCAmelCase  )		->				List[Any]:
 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  __UpperCamelCase(				self     ):
       '''simple docstring'''




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

       UpperCamelCase			:      List[str] 			=     [-0.60_40, -0.25_17, -0.10_25, 0.34_20, -0.67_58, -0.00_17, -0.10_89, -0.19_90, 0.57_28]
       for ii in range(3     ):
        for jj in range(3     ):
         UpperCamelCase			:      List[str] 			=     output[0, ii, jj]
         UpperCamelCase			:      int 			=     expected[3 * ii + jj]
         UpperCamelCase			:      int 			=     "ii={} jj={} a={} b={}".format(A_ ,							A_ ,							A_ ,							A_     )
         self.assertTrue(math.isclose(A_ ,							A_ ,							rel_tol=A_ ,							abs_tol=A_     ) ,							msg=A_     )







from typing import List, Optional

from tokenizers import ByteLevelBPETokenizer

from ...tokenization_utils_fast import PreTrainedTokenizerFast
from ...utils import logging
from .tokenization_blenderbot_small import BlenderbotSmallTokenizer


__lowerCamelCase :  Dict										=  logging.get_logger(__name__)

__lowerCamelCase :  Union[str, Any]										=  {
    """vocab_file""": """vocab.json""",
    """merges_file""": """merges.txt""",
    """tokenizer_config_file""": """tokenizer_config.json""",
}

__lowerCamelCase :  Dict										=  {
    """vocab_file""": {
        """facebook/blenderbot_small-90M""": """https://huggingface.co/facebook/blenderbot_small-90M/resolve/main/vocab.json"""
    },
    """merges_file""": {
        """facebook/blenderbot_small-90M""": """https://huggingface.co/facebook/blenderbot_small-90M/resolve/main/merges.txt"""
    },
    """tokenizer_config_file""": {
        """facebook/blenderbot_small-90M""": (
            """https://huggingface.co/facebook/blenderbot_small-90M/resolve/main/tokenizer_config.json"""
        )
    },
}

__lowerCamelCase :  Tuple										=  {
    """facebook/blenderbot_small-90M""": 512,
}



class       A__     ( __snake_case    ):
      _UpperCAmelCase			:Union[str, Any]							=     VOCAB_FILES_NAMES
      _UpperCAmelCase			:Dict							=     PRETRAINED_VOCAB_FILES_MAP
      _UpperCAmelCase			:List[Any]							=     PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
      _UpperCAmelCase			:Optional[Any]							=     BlenderbotSmallTokenizer





      def __init__(				self ,							A_=None ,							A_=None ,							A_="<|endoftext|>" ,							A_="<|endoftext|>" ,							A_="<|endoftext|>" ,							A_=False ,							A_=True ,							**A_ ,							):
       '''simple docstring'''




       super().__init__(
           ByteLevelBPETokenizer(
               vocab=A_ ,							merges=A_ ,							add_prefix_space=A_ ,							trim_offsets=A_ ,							) ,							bos_token=A_ ,							eos_token=A_ ,							unk_token=A_ ,							**A_ ,							)
       UpperCamelCase			:      Union[str, Any] 			=     add_prefix_space





      def  __UpperCamelCase(				self ,							A_ ,							A_=None     ):
       '''simple docstring'''




       UpperCamelCase			:      Dict 			=     [self.bos_token_id] + token_ids_a + [self.eos_token_id]
       if token_ids_a is None:
        return output

       return output + [self.eos_token_id] + token_ids_a + [self.eos_token_id]





      def  __UpperCamelCase(				self ,							A_ ,							A_ = None     ):
       '''simple docstring'''




       UpperCamelCase			:      Tuple 			=     [self.sep_token_id]
       UpperCamelCase			:      int 			=     [self.cls_token_id]

       if token_ids_a is None:
        return len(cls + token_ids_a + sep     ) * [0]
       return len(cls + token_ids_a + sep + sep + token_ids_a + sep     ) * [0]
