

from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available


_SCREAMING_SNAKE_CASE					:  int             =      {
    '''configuration_megatron_bert''': ['''MEGATRON_BERT_PRETRAINED_CONFIG_ARCHIVE_MAP''', '''MegatronBertConfig'''],
}

try:
       if not is_torch_available():
              raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
       pass
else:
       _SCREAMING_SNAKE_CASE					:  List[Any]             =      [
           '''MEGATRON_BERT_PRETRAINED_MODEL_ARCHIVE_LIST''',
           '''MegatronBertForCausalLM''',
           '''MegatronBertForMaskedLM''',
           '''MegatronBertForMultipleChoice''',
           '''MegatronBertForNextSentencePrediction''',
           '''MegatronBertForPreTraining''',
           '''MegatronBertForQuestionAnswering''',
           '''MegatronBertForSequenceClassification''',
           '''MegatronBertForTokenClassification''',
           '''MegatronBertModel''',
           '''MegatronBertPreTrainedModel''',
       ]

if TYPE_CHECKING:
       from .configuration_megatron_bert import MEGATRON_BERT_PRETRAINED_CONFIG_ARCHIVE_MAP, MegatronBertConfig

       try:
              if not is_torch_available():
                     raise OptionalDependencyNotAvailable()
       except OptionalDependencyNotAvailable:
              pass
       else:
              from .modeling_megatron_bert import (
                  MEGATRON_BERT_PRETRAINED_MODEL_ARCHIVE_LIST,
                  MegatronBertForCausalLM,
                  MegatronBertForMaskedLM,
                  MegatronBertForMultipleChoice,
                  MegatronBertForNextSentencePrediction,
                  MegatronBertForPreTraining,
                  MegatronBertForQuestionAnswering,
                  MegatronBertForSequenceClassification,
                  MegatronBertForTokenClassification,
                  MegatronBertModel,
                  MegatronBertPreTrainedModel,
              )

else:
       import sys

       _SCREAMING_SNAKE_CASE					:  Any             =      _LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)





import argparse
import os
import pickle
import sys

import torch

from transformers import TransfoXLConfig, TransfoXLLMHeadModel, load_tf_weights_in_transfo_xl
from transformers.models.transfo_xl import tokenization_transfo_xl as data_utils
from transformers.models.transfo_xl.tokenization_transfo_xl import CORPUS_NAME, VOCAB_FILES_NAMES
from transformers.utils import CONFIG_NAME, WEIGHTS_NAME, logging


logging.set_verbosity_info()

# We do this to be able to load python 2 datasets pickles
# See e.g. https://stackoverflow.com/questions/2121874/python-pickling-after-changing-a-modules-directory/2121918#2121918
_SCREAMING_SNAKE_CASE					:  Tuple             =      data_utils.TransfoXLTokenizer
_SCREAMING_SNAKE_CASE					:  Dict             =      data_utils.TransfoXLCorpus
_SCREAMING_SNAKE_CASE					:  Union[str, Any]             =      data_utils
_SCREAMING_SNAKE_CASE					:  Any             =      data_utils
def     UpperCAmelCase_       (						_A				,		_A				,		_A				,		_A    ):





     '''simple docstring'''




     if transfo_xl_dataset_file:
          # Convert a pre-processed corpus (see original TensorFlow repo)
          with open(_A				,		'''rb'''    ) as fp:
               SCREAMING_SNAKE_CASE__   =   pickle.load(_A				,		encoding='''latin1'''    )
          # Save vocabulary and dataset cache as Dictionaries (should be better than pickles for the long-term)
          SCREAMING_SNAKE_CASE__   =   pytorch_dump_folder_path + '''/''' + VOCAB_FILES_NAMES['''pretrained_vocab_file''']
          print(F'''Save vocabulary to {pytorch_vocab_dump_path}'''    )
          SCREAMING_SNAKE_CASE__   =   corpus.vocab.__dict__
          torch.save(_A				,		_A    )

          SCREAMING_SNAKE_CASE__   =   corpus.__dict__
          corpus_dict_no_vocab.pop('''vocab'''				,		_A    )
          SCREAMING_SNAKE_CASE__   =   pytorch_dump_folder_path + '''/''' + CORPUS_NAME
          print(F'''Save dataset to {pytorch_dataset_dump_path}'''    )
          torch.save(_A				,		_A    )

     if tf_checkpoint_path:
          # Convert a pre-trained TensorFlow model
          SCREAMING_SNAKE_CASE__   =   os.path.abspath(_A    )
          SCREAMING_SNAKE_CASE__   =   os.path.abspath(_A    )

          print(F'''Converting Transformer XL checkpoint from {tf_path} with config at {config_path}.'''    )
          # Initialise PyTorch model
          if transfo_xl_config_file == "":
               SCREAMING_SNAKE_CASE__   =   TransfoXLConfig()
          else:
               SCREAMING_SNAKE_CASE__   =   TransfoXLConfig.from_json_file(_A    )
          print(F'''Building PyTorch model from configuration: {config}'''    )
          SCREAMING_SNAKE_CASE__   =   TransfoXLLMHeadModel(_A    )

          SCREAMING_SNAKE_CASE__   =   load_tf_weights_in_transfo_xl(_A				,		_A				,		_A    )
          # Save pytorch-model
          SCREAMING_SNAKE_CASE__   =   os.path.join(_A				,		_A    )
          SCREAMING_SNAKE_CASE__   =   os.path.join(_A				,		_A    )
          print(F'''Save PyTorch model to {os.path.abspath(_A    )}'''    )
          torch.save(model.state_dict()				,		_A    )
          print(F'''Save configuration file to {os.path.abspath(_A    )}'''    )
          with open(_A				,		'''w'''				,		encoding='''utf-8'''    ) as f:
               f.write(config.to_json_string()    )


if __name__ == "__main__":
       _SCREAMING_SNAKE_CASE					:  Any             =      argparse.ArgumentParser()
       parser.add_argument(
           '''--pytorch_dump_folder_path''',
           default=None,
           type=str,
           required=True,
           help='''Path to the folder to store the PyTorch model or dataset/vocab.''',
       )
       parser.add_argument(
           '''--tf_checkpoint_path''',
           default='''''',
           type=str,
           help='''An optional path to a TensorFlow checkpoint path to be converted.''',
       )
       parser.add_argument(
           '''--transfo_xl_config_file''',
           default='''''',
           type=str,
           help=(
               '''An optional config json file corresponding to the pre-trained BERT model. \n'''
               '''This specifies the model architecture.'''
           ),
       )
       parser.add_argument(
           '''--transfo_xl_dataset_file''',
           default='''''',
           type=str,
           help='''An optional dataset file to be converted in a vocabulary.''',
       )
       _SCREAMING_SNAKE_CASE					:  Union[str, Any]             =      parser.parse_args()
       convert_transfo_xl_checkpoint_to_pytorch(
           args.tf_checkpoint_path,
           args.transfo_xl_config_file,
           args.pytorch_dump_folder_path,
           args.transfo_xl_dataset_file,
       )


