




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
UpperCamelCase__					    =    data_utils.TransfoXLTokenizer
UpperCamelCase__					    =    data_utils.TransfoXLCorpus
UpperCamelCase__					    =    data_utils
UpperCamelCase__					    =    data_utils







def 			_a   (		SCREAMING_SNAKE_CASE_							:       Optional[Any]   ,   SCREAMING_SNAKE_CASE_							:       int   ,   SCREAMING_SNAKE_CASE_							:       Dict   ,   SCREAMING_SNAKE_CASE_							:       int  )     ->   Optional[int]:





					if transfo_xl_dataset_file:
										# Convert a pre-processed corpus (see original TensorFlow repo)
										with open(__snake_case   ,   "rb"  ) as fp:
															__lowerCAmelCase            = pickle.load(__snake_case   ,   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(__snake_case   ,   __snake_case  )

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

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

										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(__snake_case  )
										print(F"""Building PyTorch model from configuration: {config}"""  )
										__lowerCAmelCase            = TransfoXLLMHeadModel(__snake_case  )

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


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


# All paths are set with the intent you should run this script from the root of the repo with the command
# python utils/check_doctest_list.py
UpperCamelCase__					    =    """."""


if __name__ == "__main__":
					UpperCamelCase__					    =    os.path.join(REPO_PATH, """utils/documentation_tests.txt""")
					UpperCamelCase__					    =    []
					UpperCamelCase__					    =    []
					with open(doctest_file_path) as fp:
										for line in fp:
															UpperCamelCase__					    =    line.strip()
															UpperCamelCase__					    =    os.path.join(REPO_PATH, line)
															if not (os.path.isfile(path) or os.path.isdir(path)):
																				non_existent_paths.append(line)
															all_paths.append(path)
					if len(non_existent_paths) > 0:
										UpperCamelCase__					    =    """\n""".join(non_existent_paths)
										raise ValueError(f'''`utils/documentation_tests.txt` contains non-existent paths:\n{non_existent_paths}''')
					if all_paths != sorted(all_paths):
										raise ValueError("""Files in `utils/documentation_tests.txt` are not in alphabetical order.""")




