


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
__a	 =      data_utils.TransfoXLTokenizer
__a	 =      data_utils.TransfoXLCorpus
__a	 =      data_utils
__a	 =      data_utils
def 	_UpperCamelCase					(    lowerCAmelCase_	,    lowerCAmelCase_	,    lowerCAmelCase_	,    lowerCAmelCase_ )    ->Dict:




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

      UpperCAmelCase      					=  corpus.__dict__
      corpus_dict_no_vocab.pop("""vocab"""	,    __A )
      UpperCAmelCase      					=  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
      UpperCAmelCase      					=  os.path.abspath(__A )
      UpperCAmelCase      					=  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 == "":
         UpperCAmelCase      					=  TransfoXLConfig()
      else:
         UpperCAmelCase      					=  TransfoXLConfig.from_json_file(__A )
      print(F"""Building PyTorch model from configuration: {config}""" )
      UpperCAmelCase      					=  TransfoXLLMHeadModel(__A )

      UpperCAmelCase      					=  load_tf_weights_in_transfo_xl(__A	,    __A	,    __A )
      # Save pytorch-model
      UpperCAmelCase      					=  os.path.join(__A	,    __A )
      UpperCAmelCase      					=  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__":
 __a	 =      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.""",
 )
 __a	 =      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,
 )











'''simple docstring'''



import unittest
from dataclasses import dataclass

import pytest

from accelerate.commands.config.config_args import SageMakerConfig
from accelerate.utils import ComputeEnvironment
from accelerate.utils.launch import _convert_nargs_to_dict
@dataclass
class UpperCAmelCase_       (     __A					):




      """simple docstring"""




      UpperCamelCase_             =      ComputeEnvironment.AMAZON_SAGEMAKER
      UpperCamelCase_             =      True
      UpperCamelCase_             =      '''ml.p3.2xlarge'''
      UpperCamelCase_             =      '''accelerate_sagemaker_execution_role'''
      UpperCamelCase_             =      '''hf-sm'''
      UpperCamelCase_             =      '''us-east-1'''
      UpperCamelCase_             =      1
      UpperCamelCase_             =      '''accelerate-sagemaker-1'''
      UpperCamelCase_             =      '''1.6'''
      UpperCamelCase_             =      '''4.4'''
      UpperCamelCase_             =      '''train.py'''
      UpperCamelCase_             =      [
          '''--model_name_or_path''',
          '''bert''',
          '''--do_train''',
          '''False''',
          '''--epochs''',
          '''3''',
          '''--learning_rate''',
          '''5e-5''',
          '''--max_steps''',
          '''50.5''',
      ]
      UpperCamelCase_             =      [
          '''--model_name_or_path''',
          '''bert''',
          '''--do_train''',
          '''--do_test''',
          '''False''',
          '''--do_predict''',
          '''--epochs''',
          '''3''',
          '''--learning_rate''',
          '''5e-5''',
          '''--max_steps''',
          '''50.5''',
      ]





class UpperCAmelCase_       (     unittest.TestCase					):




      """simple docstring"""




      def 					A__   (			self      :		int		)	->	Dict:


       '''simple docstring'''

       lowercase       : str						=_convert_nargs_to_dict(MockLaunchConfig.success_training_script_args		)
       assert isinstance(converted_args['''model_name_or_path''']  ,   UpperCAmelCase		)
       assert isinstance(converted_args['''do_train''']  ,   UpperCAmelCase		)
       assert isinstance(converted_args['''epochs''']  ,   UpperCAmelCase		)
       assert isinstance(converted_args['''learning_rate''']  ,   UpperCAmelCase		)
       assert isinstance(converted_args['''max_steps''']  ,   UpperCAmelCase		)

       with pytest.raises(UpperCAmelCase		):
        _convert_nargs_to_dict(MockLaunchConfig.fail_training_script_args		)



