




import contextlib
from multiprocessing import Pool, RLock

from tqdm.auto import tqdm

from ..utils import experimental, logging


__A								=       logging.get_logger(__name__)







class        __lowerCAmelCase					:

			"""simple docstring"""







			snake_case_  						=			None







@experimental
def      lowerCamelCase_						(			UpperCamelCase__							: Optional[Any]					,	UpperCamelCase__							: Union[str, Any]					,	UpperCamelCase__							: List[str]					,	UpperCamelCase__							: Tuple					,	UpperCamelCase__							: List[Any]					,	UpperCamelCase__							: List[Any]					,	UpperCamelCase__							: Tuple )	->			Optional[Any]:





				"""simple docstring"""


				if ParallelBackendConfig.backend_name is None:
								return _map_with_multiprocessing_pool(
								    UpperCamelCase__					,	UpperCamelCase__					,	UpperCamelCase__					,	UpperCamelCase__					,	UpperCamelCase__					,	UpperCamelCase__					,	UpperCamelCase__ )

				return _map_with_joblib(UpperCamelCase__					,	UpperCamelCase__					,	UpperCamelCase__					,	UpperCamelCase__					,	UpperCamelCase__					,	UpperCamelCase__					,	UpperCamelCase__ )







def      lowerCamelCase_						(			UpperCamelCase__							: Union[str, Any]					,	UpperCamelCase__							: Union[str, Any]					,	UpperCamelCase__							: List[Any]					,	UpperCamelCase__							: Optional[Any]					,	UpperCamelCase__							: List[str]					,	UpperCamelCase__							: Tuple					,	UpperCamelCase__							: List[str] )	->			List[str]:





				"""simple docstring"""


				__lowerCamelCase									=       num_proc if num_proc <= len(UpperCamelCase__ ) else len(UpperCamelCase__ )
				__lowerCamelCase									=       []  # We organize the splits ourselve (contiguous splits)
				for index in range(UpperCamelCase__ ):
								__lowerCamelCase									=       len(UpperCamelCase__ ) // num_proc
								__lowerCamelCase									=       len(UpperCamelCase__ ) % num_proc
								__lowerCamelCase									=       div * index + min(UpperCamelCase__					,	UpperCamelCase__ )
								__lowerCamelCase									=       start + div + (1 if index < mod else 0)
								split_kwds.append((function, iterable[start:end], types, index, disable_tqdm, desc) )

				if len(UpperCamelCase__ ) != sum(len(i[1] ) for i in split_kwds ):
								raise ValueError(
								    F"""Error dividing inputs iterable among processes. """
								    F"""Total number of objects {len(UpperCamelCase__ )}, """
								    F"""length: {sum(len(i[1] ) for i in split_kwds )}""" )

				logger.info(
				    F"""Spawning {num_proc} processes for {len(UpperCamelCase__ )} objects in slices of {[len(i[1] ) for i in split_kwds]}""" )
				__lowerCamelCase			,		__lowerCamelCase									=       None, None
				if not disable_tqdm:
								__lowerCamelCase			,		__lowerCamelCase									=       (RLock(),), tqdm.set_lock
				with Pool(UpperCamelCase__					,	initargs=UpperCamelCase__					,	initializer=UpperCamelCase__ ) as pool:
								__lowerCamelCase									=       pool.map(UpperCamelCase__					,	UpperCamelCase__ )
				logger.info(F"""Finished {num_proc} processes""" )
				__lowerCamelCase									=       [obj for proc_res in mapped for obj in proc_res]
				logger.info(F"""Unpacked {len(UpperCamelCase__ )} objects""" )

				return mapped







def      lowerCamelCase_						(			UpperCamelCase__							: str					,	UpperCamelCase__							: int					,	UpperCamelCase__							: int					,	UpperCamelCase__							: List[str]					,	UpperCamelCase__							: int					,	UpperCamelCase__							: Union[str, Any]					,	UpperCamelCase__							: Optional[Any] )	->			str:





				"""simple docstring"""


				import joblib

				with joblib.parallel_backend(ParallelBackendConfig.backend_name					,	n_jobs=UpperCamelCase__ ):
								return joblib.Parallel()(
								    joblib.delayed(UpperCamelCase__ )((function, obj, types, None, True, None) ) for obj in iterable )







@experimental
@contextlib.contextmanager
def      lowerCamelCase_						(			UpperCamelCase__							: str )	->			str:





				"""simple docstring"""


				__lowerCamelCase									=       backend_name

				if backend_name == "spark":
								from joblibspark import register_spark

								register_spark()

								# TODO: call create_cache_and_write_probe if "download" in steps
								# TODO: raise NotImplementedError when Dataset.map etc is called

				try:
								yield
				finally:
								__lowerCamelCase									=       None








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      lowerCamelCase_						(			UpperCamelCase__							: int					,	UpperCamelCase__							: str					,	UpperCamelCase__							: Optional[Any]					,	UpperCamelCase__							: Tuple )	->			Tuple:





				"""simple docstring"""


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

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

								__lowerCamelCase									=       load_tf_weights_in_transfo_xl(UpperCamelCase__					,	UpperCamelCase__					,	UpperCamelCase__ )
								# Save pytorch-model
								__lowerCamelCase									=       os.path.join(UpperCamelCase__					,	UpperCamelCase__ )
								__lowerCamelCase									=       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								=       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,
					)


