



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__  :  str								=					data_utils.TransfoXLTokenizer
A__  :  Any								=					data_utils.TransfoXLCorpus
A__  :  str								=					data_utils
A__  :  int								=					data_utils


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

				lowerCAmelCase_			:     Union[str, Any]									= corpus.__dict__
				corpus_dict_no_vocab.pop('''vocab'''     ,__UpperCamelCase )
				lowerCAmelCase_			:     List[Any]									= pytorch_dump_folder_path + '''/''' + CORPUS_NAME
				print(f"""Save dataset to {pytorch_dataset_dump_path}""" )
				torch.save(__UpperCamelCase     ,__UpperCamelCase )

		if tf_checkpoint_path:
				# Convert a pre-trained TensorFlow model
				lowerCAmelCase_			:     Optional[Any]									= os.path.abspath(__UpperCamelCase )
				lowerCAmelCase_			:     Dict									= os.path.abspath(__UpperCamelCase )

				print(f"""Converting Transformer XL checkpoint from {tf_path} with config at {config_path}.""" )
				# Initialise PyTorch model
				if transfo_xl_config_file == "":
						lowerCAmelCase_			:     str									= TransfoXLConfig()
				else:
						lowerCAmelCase_			:     int									= TransfoXLConfig.from_json_file(__UpperCamelCase )
				print(f"""Building PyTorch model from configuration: {config}""" )
				lowerCAmelCase_			:     Optional[int]									= TransfoXLLMHeadModel(__UpperCamelCase )

				lowerCAmelCase_			:     List[Any]									= load_tf_weights_in_transfo_xl(__UpperCamelCase     ,__UpperCamelCase     ,__UpperCamelCase )
				# Save pytorch-model
				lowerCAmelCase_			:     str									= os.path.join(__UpperCamelCase     ,__UpperCamelCase )
				lowerCAmelCase_			:     List[str]									= os.path.join(__UpperCamelCase     ,__UpperCamelCase )
				print(f"""Save PyTorch model to {os.path.abspath(__UpperCamelCase )}""" )
				torch.save(model.state_dict()     ,__UpperCamelCase )
				print(f"""Save configuration file to {os.path.abspath(__UpperCamelCase )}""" )
				with open(__UpperCamelCase     ,'''w'''     ,encoding='''utf-8''' ) as f:
						f.write(config.to_json_string() )


if __name__ == "__main__":
						A__  :  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.''',
						)
						A__  :  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,
						)










import importlib
import sys
from argparse import REMAINDER, ArgumentParser
from pathlib import Path

import torch_xla.distributed.xla_multiprocessing as xmp


def 			UpperCamelCase(  ):
		lowerCAmelCase_			:     List[str]									= ArgumentParser(
		    description=(
		        '''PyTorch TPU distributed training launch '''
		        '''helper utility that will spawn up '''
		        '''multiple distributed processes'''
		    ) )

		# Optional arguments for the launch helper
		parser.add_argument('''--num_cores'''     ,type=__UpperCamelCase     ,default=1     ,help='''Number of TPU cores to use (1 or 8).''' )

		# positional
		parser.add_argument(
		    '''training_script'''     ,type=__UpperCamelCase     ,help=(
		        '''The full path to the single TPU training '''
		        '''program/script to be launched in parallel, '''
		        '''followed by all the arguments for the '''
		        '''training script'''
		    )     ,)

		# rest from the training program
		parser.add_argument('''training_script_args'''     ,nargs=__UpperCamelCase )

		return parser.parse_args()


def 			UpperCamelCase(  ):
		lowerCAmelCase_			:     str									= parse_args()

		# Import training_script as a module.
		lowerCAmelCase_			:     str									= Path(args.training_script )
		sys.path.append(str(script_fpath.parent.resolve() ) )
		lowerCAmelCase_			:     Tuple									= script_fpath.stem
		lowerCAmelCase_			:     Union[str, Any]									= importlib.import_module(__UpperCamelCase )

		# Patch sys.argv
		lowerCAmelCase_			:     Optional[int]									= [args.training_script] + args.training_script_args + ['''--tpu_num_cores''', str(args.num_cores )]
		xmp.spawn(mod._mp_fn     ,args=()     ,nprocs=args.num_cores )


if __name__ == "__main__":
						main()





