

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



       '''simple docstring'''







       def __init__(							self ,   lowercase__ ,   lowercase__=1_3 ,   lowercase__=7 ,   lowercase__=True ,   lowercase__=True ,   lowercase__=True ,   lowercase__=True ,   lowercase__=9_9 ,   lowercase__=6_4 ,   lowercase__=3_2 ,   lowercase__=5 ,   lowercase__=4 ,   lowercase__=3_7 ,   lowercase__="gelu" ,   lowercase__=0.1 ,   lowercase__=0.1 ,   lowercase__=5_1_2 ,   lowercase__=1_6 ,   lowercase__=2 ,   lowercase__=0.02 ,   lowercase__=3 ,   lowercase__=4 ,   lowercase__=None ,   ):







             '''simple docstring'''



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

       def     __UpperCamelCase     (							self     ):







             '''simple docstring'''



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

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

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

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

             __A    =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=lowercase__ ,   initializer_range=self.initializer_range ,   )

       def     __UpperCamelCase     (							self ,   lowercase__ ,   lowercase__ ,   lowercase__ ,   lowercase__ ,   lowercase__ ,   lowercase__ ,   lowercase__     ):







             '''simple docstring'''



             __A    =MegatronBertModel(config=lowercase__     )
             model.to(lowercase__     )
             model.eval()
             __A    =model(lowercase__ ,   attention_mask=lowercase__ ,   token_type_ids=lowercase__     )
             __A    =model(lowercase__ ,   token_type_ids=lowercase__     )
             __A    =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     __UpperCamelCase     (							self ,   lowercase__ ,   lowercase__ ,   lowercase__ ,   lowercase__ ,   lowercase__ ,   lowercase__ ,   lowercase__     ):







             '''simple docstring'''



             __A    =MegatronBertForMaskedLM(config=lowercase__     )
             model.to(lowercase__     )
             model.eval()
             __A    =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     __UpperCamelCase     (							self ,   lowercase__ ,   lowercase__ ,   lowercase__ ,   lowercase__ ,   lowercase__ ,   lowercase__ ,   lowercase__     ):







             '''simple docstring'''



             __A    =MegatronBertForCausalLM(config=lowercase__     )
             model.to(lowercase__     )
             model.eval()
             __A    =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     __UpperCamelCase     (							self ,   lowercase__ ,   lowercase__ ,   lowercase__ ,   lowercase__ ,   lowercase__ ,   lowercase__ ,   lowercase__     ):







             '''simple docstring'''



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

       def     __UpperCamelCase     (							self ,   lowercase__ ,   lowercase__ ,   lowercase__ ,   lowercase__ ,   lowercase__ ,   lowercase__ ,   lowercase__     ):







             '''simple docstring'''



             __A    =MegatronBertForPreTraining(config=lowercase__     )
             model.to(lowercase__     )
             model.eval()
             __A    =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     __UpperCamelCase     (							self ,   lowercase__ ,   lowercase__ ,   lowercase__ ,   lowercase__ ,   lowercase__ ,   lowercase__ ,   lowercase__     ):







             '''simple docstring'''



             __A    =MegatronBertForQuestionAnswering(config=lowercase__     )
             model.to(lowercase__     )
             model.eval()
             __A    =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     __UpperCamelCase     (							self ,   lowercase__ ,   lowercase__ ,   lowercase__ ,   lowercase__ ,   lowercase__ ,   lowercase__ ,   lowercase__     ):







             '''simple docstring'''



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

       def     __UpperCamelCase     (							self ,   lowercase__ ,   lowercase__ ,   lowercase__ ,   lowercase__ ,   lowercase__ ,   lowercase__ ,   lowercase__     ):







             '''simple docstring'''



             __A    =self.num_labels
             __A    =MegatronBertForTokenClassification(config=lowercase__     )
             model.to(lowercase__     )
             model.eval()
             __A    =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     __UpperCamelCase     (							self ,   lowercase__ ,   lowercase__ ,   lowercase__ ,   lowercase__ ,   lowercase__ ,   lowercase__ ,   lowercase__     ):







             '''simple docstring'''



             __A    =self.num_choices
             __A    =MegatronBertForMultipleChoice(config=lowercase__     )
             model.to(lowercase__     )
             model.eval()
             __A    =input_ids.unsqueeze(1     ).expand(-1 ,   self.num_choices ,   -1     ).contiguous()
             __A    =token_type_ids.unsqueeze(1     ).expand(-1 ,   self.num_choices ,   -1     ).contiguous()
             __A    =input_mask.unsqueeze(1     ).expand(-1 ,   self.num_choices ,   -1     ).contiguous()
             __A    =model(
                 lowercase__ ,   attention_mask=lowercase__ ,   token_type_ids=lowercase__ ,   labels=lowercase__ ,   )
             self.parent.assertEqual(result.logits.shape ,   (self.batch_size, self.num_choices)     )


       def     __UpperCamelCase     (							self     ):







             '''simple docstring'''



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





@require_torch
class        lowerCAmelCase__  (					__magic_name__				,							__magic_name__				,							unittest.TestCase			):



       '''simple docstring'''







       lowercase_						=    (
           (
               MegatronBertModel,
               MegatronBertForMaskedLM,
               MegatronBertForCausalLM,
               MegatronBertForMultipleChoice,
               MegatronBertForNextSentencePrediction,
               MegatronBertForPreTraining,
               MegatronBertForQuestionAnswering,
               MegatronBertForSequenceClassification,
               MegatronBertForTokenClassification,
           )
           if is_torch_available()
           else ()
       )
       lowercase_						=    (
           {
               """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 {}
       )
       lowercase_						=    True
       # test_resize_embeddings = False
       lowercase_						=    False

       def     __UpperCamelCase     (							self ,   lowercase__ ,   lowercase__ ,   lowercase__=False     ):







             '''simple docstring'''



             __A    =super()._prepare_for_class(lowercase__ ,   lowercase__ ,   return_labels=lowercase__     )

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

       def     __UpperCamelCase     (							self     ):







             '''simple docstring'''



             __A    =MegatronBertModelTester(self     )
             __A    =ConfigTester(self ,   config_class=lowercase__ ,   hidden_size=3_7     )

       def     __UpperCamelCase     (							self     ):







             '''simple docstring'''



             self.config_tester.run_common_tests()

       def     __UpperCamelCase     (							self     ):







             '''simple docstring'''



             __A    =self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_megatron_bert_model(*lowercase__     )

       def     __UpperCamelCase     (							self     ):







             '''simple docstring'''



             __A    =self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_megatron_bert_for_masked_lm(*lowercase__     )

       def     __UpperCamelCase     (							self     ):







             '''simple docstring'''



             __A    =self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_megatron_bert_for_multiple_choice(*lowercase__     )

       def     __UpperCamelCase     (							self     ):







             '''simple docstring'''



             __A    =self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_megatron_bert_for_next_sequence_prediction(*lowercase__     )

       def     __UpperCamelCase     (							self     ):







             '''simple docstring'''



             __A    =self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_megatron_bert_for_pretraining(*lowercase__     )

       def     __UpperCamelCase     (							self     ):







             '''simple docstring'''



             __A    =self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_megatron_bert_for_question_answering(*lowercase__     )

       def     __UpperCamelCase     (							self     ):







             '''simple docstring'''



             __A    =self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_megatron_bert_for_sequence_classification(*lowercase__     )


       def     __UpperCamelCase     (							self     ):







             '''simple docstring'''



             __A    =self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_megatron_bert_for_token_classification(*lowercase__     )





def      A__					(			__A  :						Any			)					->Optional[int]:
      return torch.tensor(
          __A     ,				dtype=torch.long     ,				device=__A     ,				)


_lowerCamelCase      :					Optional[int]					    =      1e-4





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



       '''simple docstring'''







       @slow
       @unittest.skip('''Model is not available.'''     )
       def     __UpperCamelCase     (							self     ):







             '''simple docstring'''



             __A    ='''nvidia/megatron-bert-uncased-345m'''
             if "MYDIR" in os.environ:
                   __A    =os.path.join(os.environ['''MYDIR'''] ,   lowercase__     )
             __A    =MegatronBertModel.from_pretrained(lowercase__     )
             model.to(lowercase__     )
             model.half()
             __A    =_long_tensor([[1_0_1, 7_1_1_0, 1_0_0_5, 1_0_5_6, 2_0_2_3, 1_1_3_3_3, 1_7_4_1_3, 1_0_2_9, 1_0_2]]     )
             with torch.no_grad():
                   __A    =model(lowercase__     )[0]
             __A    =torch.Size((1, 9, 1_0_2_4)     )
             self.assertEqual(output.shape ,   lowercase__     )

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





import argparse
import pickle

import numpy as np
import torch
from torch import nn

from transformers import ReformerConfig, ReformerModelWithLMHead
from transformers.utils import logging


logging.set_verbosity_info()





def      A__					(			__A  :						Any     ,				__A  :						Dict     ,				__A  :						Optional[int]=None			)					->Tuple:
      # set parameter of one layer
      assert torch_layer.weight.shape == weight.shape, F'''{torch_layer} layer.weight does not match'''
      __A    =nn.Parameter(__A			)
      if bias is not None:
            assert torch_layer.bias.shape == bias.shape, F'''{torch_layer} layer.bias does not match'''
            __A    =nn.Parameter(__A			)





def      A__					(			__A  :						List[Any]     ,				__A  :						Tuple     ,				__A  :						List[Any]			)					->Dict:
      # set torch weights for 1-to-1 comparison
      __A    =np.asarray(weights[0]			)
      __A    =np.asarray(weights[1]			)
      __A    =np.asarray(weights[2]			)

      set_param(
          torch_layer.self_attention.query_key     ,				torch.tensor(__A			).transpose(1     ,				2			).contiguous().view(-1     ,				__A			)     ,				)
      set_param(
          torch_layer.self_attention.value     ,				torch.tensor(__A			).transpose(1     ,				2			).contiguous().view(-1     ,				__A			)     ,				)
      set_param(
          torch_layer.output.dense     ,				torch.tensor(__A			).view(-1     ,				__A			).contiguous().transpose(0     ,				1			)     ,				)





def      A__					(			__A  :						Optional[Any]     ,				__A  :						Tuple     ,				__A  :						Optional[Any]			)					->int:
      # set torch weights for 1-to-1 comparison
      __A    =np.asarray(weights[0]			)
      __A    =np.asarray(weights[1]			)
      __A    =np.asarray(weights[2]			)
      __A    =np.asarray(weights[3]			)

      set_param(
          torch_layer.self_attention.query     ,				torch.tensor(__A			).transpose(1     ,				2			).contiguous().view(-1     ,				__A			)     ,				)
      set_param(
          torch_layer.self_attention.key     ,				torch.tensor(__A			).transpose(1     ,				2			).contiguous().view(-1     ,				__A			)     ,				)
      set_param(
          torch_layer.self_attention.value     ,				torch.tensor(__A			).transpose(1     ,				2			).contiguous().view(-1     ,				__A			)     ,				)
      set_param(
          torch_layer.output.dense     ,				torch.tensor(__A			).view(-1     ,				__A			).contiguous().transpose(0     ,				1			)     ,				)





def      A__					(			__A  :						int     ,				__A  :						List[str]     ,				__A  :						Optional[Any]			)					->Any:
      # layernorm 1
      __A    =weights[0][0][0]
      __A    =np.asarray(layer_norm_a[0]			)
      __A    =np.asarray(layer_norm_a[1]			)
      set_param(
          torch_block.attention.layer_norm     ,				torch.tensor(__A			)     ,				torch.tensor(__A			)     ,				)

      # lsh weights + output
      __A    =weights[0][1]
      if len(__A			) < 4:
            set_layer_weights_in_torch_lsh(__A     ,				torch_block.attention     ,				__A			)
      else:
            set_layer_weights_in_torch_local(__A     ,				torch_block.attention     ,				__A			)

      # intermediate weighs
      __A    =weights[2][0][1][2]

      # Chunked Feed Forward
      if len(__A			) == 4:
            __A    =intermediate_weights[2]

      # layernorm 2
      __A    =np.asarray(intermediate_weights[0][0]			)
      __A    =np.asarray(intermediate_weights[0][1]			)
      set_param(
          torch_block.feed_forward.layer_norm     ,				torch.tensor(__A			)     ,				torch.tensor(__A			)     ,				)

      # intermediate dense
      __A    =np.asarray(intermediate_weights[1][0]			)
      __A    =np.asarray(intermediate_weights[1][1]			)
      set_param(
          torch_block.feed_forward.dense.dense     ,				torch.tensor(__A			).transpose(0     ,				1			).contiguous()     ,				torch.tensor(__A			)     ,				)

      # intermediate out
      __A    =np.asarray(intermediate_weights[4][0]			)
      __A    =np.asarray(intermediate_weights[4][1]			)
      set_param(
          torch_block.feed_forward.output.dense     ,				torch.tensor(__A			).transpose(0     ,				1			).contiguous()     ,				torch.tensor(__A			)     ,				)





def      A__					(			__A  :						Tuple     ,				__A  :						List[Any]     ,				__A  :						Optional[int]			)					->List[Any]:
      # reformer model
      __A    =torch_model.reformer

      # word embeds
      __A    =np.asarray(weights[1]			)
      set_param(
          torch_model_reformer.embeddings.word_embeddings     ,				torch.tensor(__A			)     ,				)

      if isinstance(weights[3]     ,				__A			):
            __A    =torch_model_reformer.embeddings.position_embeddings
            for emb_idx in range(len(position_embeddings.weights			)			):
                  __A    =np.asarray(weights[3][emb_idx][0]			)
                  assert (
                      position_embeddings.weights[emb_idx].shape == emb_weights.shape
                  ), F'''{position_embeddings[emb_idx]} emb does not match'''
                  __A    =nn.Parameter(torch.tensor(__A			)			)

      __A    =weights[5]
      assert len(torch_model_reformer.encoder.layers			) * 4 == len(
          __A			), "HF and trax model do not have the same number of layers"
      for layer_idx, layer in enumerate(torch_model_reformer.encoder.layers			):
            __A    =trax_layer_weights[4 * layer_idx : 4 * (layer_idx + 1)]
            set_block_weights_in_torch(__A     ,				__A     ,				__A			)

      # output layer norm
      __A    =np.asarray(weights[7][0]			)
      __A    =np.asarray(weights[7][1]			)
      set_param(
          torch_model_reformer.encoder.layer_norm     ,				torch.tensor(__A			)     ,				torch.tensor(__A			)     ,				)

      # output embeddings
      __A    =np.asarray(weights[9][0]			)
      __A    =np.asarray(weights[9][1]			)
      set_param(
          torch_model.lm_head.decoder     ,				torch.tensor(__A			).transpose(0     ,				1			).contiguous()     ,				torch.tensor(__A			)     ,				)





def      A__					(			__A  :						int     ,				__A  :						Any     ,				__A  :						Tuple			)					->Union[str, Any]:
      # Initialise PyTorch model
      __A    =ReformerConfig.from_json_file(__A			)
      print(F'''Building PyTorch model from configuration: {config}'''			)
      __A    =ReformerModelWithLMHead(__A			)

      with open(__A     ,				'''rb'''			) as f:
            __A    =pickle.load(__A			)['''weights''']

      set_model_weights_in_torch(__A     ,				__A     ,				config.hidden_size			)

      # Save pytorch-model
      print(F'''Save PyTorch model to {pytorch_dump_path}'''			)
      torch.save(model.state_dict()     ,				__A			)


if __name__ == "__main__":
   _lowerCamelCase      :					Union[str, Any]					    =      argparse.ArgumentParser()
   # Required parameters
   parser.add_argument(
       '''--trax_model_pkl_path''', default=None, type=str, required=True, help='''Path to the TensorFlow checkpoint path.'''
   )
   parser.add_argument(
       '''--config_file''',
       default=None,
       type=str,
       required=True,
       help=(
           '''The config json file corresponding to the pre-trained Reformer model. \n'''
           '''This specifies the model architecture.'''
       ),
   )
   parser.add_argument(
       '''--pytorch_dump_path''', default=None, type=str, required=True, help='''Path to the output PyTorch model.'''
   )
   _lowerCamelCase      :					Tuple					    =      parser.parse_args()
   convert_trax_checkpoint_to_pytorch(args.trax_model_pkl_path, args.config_file, args.pytorch_dump_path)


