






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      ,    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__=5_12      ,    a__=16      ,    a__=2      ,    a__=0.02      ,    a__=3      ,    a__=4      ,    a__=None      ,    )			->   Dict:






            '''simple docstring'''





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



       def        __lowercase    (				self     )			->   List[str]:






            '''simple docstring'''





            __snake_case      :List[Any] 							=    ids_tensor([self.batch_size, self.seq_length]      ,    self.vocab_size     )

            __snake_case      :Any 							=    None
            if self.use_input_mask:
                 __snake_case      :Optional[int] 							=    random_attention_mask([self.batch_size, self.seq_length]     )

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

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

            __snake_case      :str 							=    self.get_config()

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



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



       def        __lowercase    (				self      ,    a__      ,    a__      ,    a__      ,    a__      ,    a__      ,    a__      ,    a__     )			->   Dict:






            '''simple docstring'''





            __snake_case      :Optional[int] 							=    MegatronBertModel(config=a__     )
            model.to(a__     )
            model.eval()
            __snake_case      :str 							=    model(a__      ,    attention_mask=a__      ,    token_type_ids=a__     )
            __snake_case      :Dict 							=    model(a__      ,    token_type_ids=a__     )
            __snake_case      :List[str] 							=    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        __lowercase    (				self      ,    a__      ,    a__      ,    a__      ,    a__      ,    a__      ,    a__      ,    a__     )			->   Union[str, Any]:






            '''simple docstring'''





            __snake_case      :Dict 							=    MegatronBertForMaskedLM(config=a__     )
            model.to(a__     )
            model.eval()
            __snake_case      :Optional[Any] 							=    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        __lowercase    (				self      ,    a__      ,    a__      ,    a__      ,    a__      ,    a__      ,    a__      ,    a__     )			->   Optional[Any]:






            '''simple docstring'''





            __snake_case      :Tuple 							=    MegatronBertForCausalLM(config=a__     )
            model.to(a__     )
            model.eval()
            __snake_case      :Optional[int] 							=    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        __lowercase    (				self      ,    a__      ,    a__      ,    a__      ,    a__      ,    a__      ,    a__      ,    a__     )			->   str:






            '''simple docstring'''





            __snake_case      :Tuple 							=    MegatronBertForNextSentencePrediction(config=a__     )
            model.to(a__     )
            model.eval()
            __snake_case      :Optional[int] 							=    model(
                a__      ,    attention_mask=a__      ,    token_type_ids=a__      ,    labels=a__      ,    )
            self.parent.assertEqual(result.logits.shape      ,    (self.batch_size, 2)     )



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






            '''simple docstring'''





            __snake_case      :Optional[Any] 							=    MegatronBertForPreTraining(config=a__     )
            model.to(a__     )
            model.eval()
            __snake_case      :Dict 							=    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        __lowercase    (				self      ,    a__      ,    a__      ,    a__      ,    a__      ,    a__      ,    a__      ,    a__     )			->   List[str]:






            '''simple docstring'''





            __snake_case      :Optional[int] 							=    MegatronBertForQuestionAnswering(config=a__     )
            model.to(a__     )
            model.eval()
            __snake_case      :str 							=    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        __lowercase    (				self      ,    a__      ,    a__      ,    a__      ,    a__      ,    a__      ,    a__      ,    a__     )			->   List[str]:






            '''simple docstring'''





            __snake_case      :Union[str, Any] 							=    self.num_labels
            __snake_case      :str 							=    MegatronBertForSequenceClassification(a__     )
            model.to(a__     )
            model.eval()
            __snake_case      :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        __lowercase    (				self      ,    a__      ,    a__      ,    a__      ,    a__      ,    a__      ,    a__      ,    a__     )			->   int:






            '''simple docstring'''





            __snake_case      :Optional[Any] 							=    self.num_labels
            __snake_case      :str 							=    MegatronBertForTokenClassification(config=a__     )
            model.to(a__     )
            model.eval()
            __snake_case      :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        __lowercase    (				self      ,    a__      ,    a__      ,    a__      ,    a__      ,    a__      ,    a__      ,    a__     )			->   str:






            '''simple docstring'''





            __snake_case      :List[Any] 							=    self.num_choices
            __snake_case      :Tuple 							=    MegatronBertForMultipleChoice(config=a__     )
            model.to(a__     )
            model.eval()
            __snake_case      :str 							=    input_ids.unsqueeze(1     ).expand(-1      ,    self.num_choices      ,    -1     ).contiguous()
            __snake_case      :Any 							=    token_type_ids.unsqueeze(1     ).expand(-1      ,    self.num_choices      ,    -1     ).contiguous()
            __snake_case      :str 							=    input_mask.unsqueeze(1     ).expand(-1      ,    self.num_choices      ,    -1     ).contiguous()
            __snake_case      :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        __lowercase    (				self     )			->   Any:






            '''simple docstring'''





            __snake_case      :Optional[int] 							=    self.prepare_config_and_inputs()
            (
                (
                __snake_case
            )							,			(
                __snake_case
            )							,			(
                __snake_case
            )							,			(
                __snake_case
            )							,			(
                __snake_case
            )							,			(
                __snake_case
            )							,			(
                __snake_case
            )							,			
            )      :Optional[int] 							=    config_and_inputs
            __snake_case      :str 							=    {"""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				:					str								=			(
           (
               MegatronBertModel,
               MegatronBertForMaskedLM,
               MegatronBertForCausalLM,
               MegatronBertForMultipleChoice,
               MegatronBertForNextSentencePrediction,
               MegatronBertForPreTraining,
               MegatronBertForQuestionAnswering,
               MegatronBertForSequenceClassification,
               MegatronBertForTokenClassification,
           )
           if is_torch_available()
           else ()
       )
       lowerCamelCase				:					Optional[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 {}
       )
       lowerCamelCase				:					List[str]								=			True
       # test_resize_embeddings = False
       lowerCamelCase				:					Any								=			False



       def        __lowercase    (				self      ,    a__      ,    a__      ,    a__=False     )			->   Tuple:






            '''simple docstring'''





            __snake_case      :List[Any] 							=    super()._prepare_for_class(a__      ,    a__      ,    return_labels=a__     )

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



       def        __lowercase    (				self     )			->   Any:






            '''simple docstring'''





            __snake_case      :Dict 							=    MegatronBertModelTester(self     )
            __snake_case      :Optional[Any] 							=    ConfigTester(self      ,    config_class=a__      ,    hidden_size=37     )



       def        __lowercase    (				self     )			->   Optional[int]:






            '''simple docstring'''





            self.config_tester.run_common_tests()



       def        __lowercase    (				self     )			->   Dict:






            '''simple docstring'''





            __snake_case      :str 							=    self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_megatron_bert_model(*a__     )



       def        __lowercase    (				self     )			->   Dict:






            '''simple docstring'''





            __snake_case      :List[Any] 							=    self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_megatron_bert_for_masked_lm(*a__     )



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






            '''simple docstring'''





            __snake_case      :Optional[Any] 							=    self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_megatron_bert_for_multiple_choice(*a__     )



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






            '''simple docstring'''





            __snake_case      :Optional[int] 							=    self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_megatron_bert_for_next_sequence_prediction(*a__     )



       def        __lowercase    (				self     )			->   Optional[int]:






            '''simple docstring'''





            __snake_case      :Any 							=    self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_megatron_bert_for_pretraining(*a__     )



       def        __lowercase    (				self     )			->   Optional[int]:






            '''simple docstring'''





            __snake_case      :int 							=    self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_megatron_bert_for_question_answering(*a__     )



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






            '''simple docstring'''





            __snake_case      :int 							=    self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_megatron_bert_for_sequence_classification(*a__     )



       def        __lowercase    (				self     )			->   Any:






            '''simple docstring'''





            __snake_case      :str 							=    self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_megatron_bert_for_token_classification(*a__     )






def 	UpperCamelCase	(							snake_case__  : Union[str, Any]	):




     '''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        __lowercase    (				self     )			->   Optional[Any]:






            '''simple docstring'''





            __snake_case      :Optional[int] 							=    """nvidia/megatron-bert-uncased-345m"""
            if "MYDIR" in os.environ:
                 __snake_case      :Dict 							=    os.path.join(os.environ["""MYDIR"""]      ,    a__     )
            __snake_case      :Optional[int] 							=    MegatronBertModel.from_pretrained(a__     )
            model.to(a__     )
            model.half()
            __snake_case      :List[str] 							=    _long_tensor([[1_01, 71_10, 10_05, 10_56, 20_23, 1_13_33, 1_74_13, 10_29, 1_02]]     )
            with torch.no_grad():
                 __snake_case      :List[str] 							=    model(a__     )[0]
            __snake_case      :List[str] 							=    torch.Size((1, 9, 10_24)     )
            self.assertEqual(output.shape      ,    a__     )

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










from ...configuration_utils import PretrainedConfig


lowerCamelCase__       		=  {
    """google/tapas-base-finetuned-sqa""": (
        """https://huggingface.co/google/tapas-base-finetuned-sqa/resolve/main/config.json"""
    ),
    """google/tapas-base-finetuned-wtq""": (
        """https://huggingface.co/google/tapas-base-finetuned-wtq/resolve/main/config.json"""
    ),
    """google/tapas-base-finetuned-wikisql-supervised""": (
        """https://huggingface.co/google/tapas-base-finetuned-wikisql-supervised/resolve/main/config.json"""
    ),
    """google/tapas-base-finetuned-tabfact""": (
        """https://huggingface.co/google/tapas-base-finetuned-tabfact/resolve/main/config.json"""
    ),
}




class       snake_case__ (				lowercase_):

       '''simple docstring'''
       lowerCamelCase				:					Tuple								=			"tapas"



       def __init__(				self      ,    a__=3_05_22      ,    a__=7_68      ,    a__=12      ,    a__=12      ,    a__=30_72      ,    a__="gelu"      ,    a__=0.1      ,    a__=0.1      ,    a__=10_24      ,    a__=[3, 2_56, 2_56, 2, 2_56, 2_56, 10]      ,    a__=0.02      ,    a__=1e-12      ,    a__=0      ,    a__=10.0      ,    a__=0      ,    a__=1.0      ,    a__=None      ,    a__=1.0      ,    a__=False      ,    a__=None      ,    a__=1.0      ,    a__=1.0      ,    a__=False      ,    a__=False      ,    a__="ratio"      ,    a__=None      ,    a__=None      ,    a__=64      ,    a__=32      ,    a__=False      ,    a__=True      ,    a__=False      ,    a__=False      ,    a__=True      ,    a__=False      ,    a__=None      ,    a__=None      ,    **a__      ,    )			->   Any:






            '''simple docstring'''





            super().__init__(pad_token_id=a__      ,    **a__     )

            # BERT hyperparameters (with updated max_position_embeddings and type_vocab_sizes)
            __snake_case      :Optional[Any] 							=    vocab_size
            __snake_case      :Tuple 							=    hidden_size
            __snake_case      :Tuple 							=    num_hidden_layers
            __snake_case      :Any 							=    num_attention_heads
            __snake_case      :List[str] 							=    hidden_act
            __snake_case      :List[str] 							=    intermediate_size
            __snake_case      :int 							=    hidden_dropout_prob
            __snake_case      :List[Any] 							=    attention_probs_dropout_prob
            __snake_case      :Any 							=    max_position_embeddings
            __snake_case      :Optional[int] 							=    type_vocab_sizes
            __snake_case      :Tuple 							=    initializer_range
            __snake_case      :Any 							=    layer_norm_eps

            # Fine-tuning task hyperparameters
            __snake_case      :Optional[int] 							=    positive_label_weight
            __snake_case      :Optional[Any] 							=    num_aggregation_labels
            __snake_case      :Dict 							=    aggregation_loss_weight
            __snake_case      :str 							=    use_answer_as_supervision
            __snake_case      :int 							=    answer_loss_importance
            __snake_case      :List[Any] 							=    use_normalized_answer_loss
            __snake_case      :List[str] 							=    huber_loss_delta
            __snake_case      :Any 							=    temperature
            __snake_case      :Optional[int] 							=    aggregation_temperature
            __snake_case      :List[str] 							=    use_gumbel_for_cells
            __snake_case      :Optional[int] 							=    use_gumbel_for_aggregation
            __snake_case      :Optional[Any] 							=    average_approximation_function
            __snake_case      :Union[str, Any] 							=    cell_selection_preference
            __snake_case      :Tuple 							=    answer_loss_cutoff
            __snake_case      :Union[str, Any] 							=    max_num_rows
            __snake_case      :Optional[int] 							=    max_num_columns
            __snake_case      :str 							=    average_logits_per_cell
            __snake_case      :Union[str, Any] 							=    select_one_column
            __snake_case      :Optional[int] 							=    allow_empty_column_selection
            __snake_case      :Optional[Any] 							=    init_cell_selection_weights_to_zero
            __snake_case      :Tuple 							=    reset_position_index_per_cell
            __snake_case      :Any 							=    disable_per_token_loss

            # Aggregation hyperparameters
            __snake_case      :List[str] 							=    aggregation_labels
            __snake_case      :int 							=    no_aggregation_label_index

            if isinstance(self.aggregation_labels      ,    a__     ):
                 __snake_case      :Any 							=    {int(a__     ): v for k, v in aggregation_labels.items()}


