






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							:   Tuple							=  data_utils.TransfoXLTokenizer
__UpperCAmelCase							:   Union[str, Any]							=  data_utils.TransfoXLCorpus
__UpperCAmelCase							:   Optional[Any]							=  data_utils
__UpperCAmelCase							:   Optional[Any]							=  data_utils





def 					lowercase_ (    __snake_case     :							List[str]	,	__snake_case     :							Any	,	__snake_case     :							Optional[Any]	,	__snake_case     :							List[Any]      )     -> Union[str, Any]:





				'''simple docstring'''


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

								snake_case__      :str         =   corpus.__dict__
								corpus_dict_no_vocab.pop("vocab"	,	__snake_case      )
								snake_case__      :str         =   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
								snake_case__      :Tuple         =   os.path.abspath(__snake_case      )
								snake_case__      :int         =   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 == "":
												snake_case__      :Dict         =   TransfoXLConfig()
								else:
												snake_case__      :List[Any]         =   TransfoXLConfig.from_json_file(__snake_case      )
								print(F'Building PyTorch model from configuration: {config}'      )
								snake_case__      :str         =   TransfoXLLMHeadModel(__snake_case      )

								snake_case__      :Any         =   load_tf_weights_in_transfo_xl(__snake_case	,	__snake_case	,	__snake_case      )
								# Save pytorch-model
								snake_case__      :Dict         =   os.path.join(__snake_case	,	__snake_case      )
								snake_case__      :int         =   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							:   Dict							=  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							:   Tuple							=  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 pytest


__UpperCAmelCase							:   int							=  "__dummy_dataset1__"

__UpperCAmelCase							:   int							=  "\nimport json\nimport os\n\nimport datasets\n\n\nREPO_URL = \"https://huggingface.co/datasets/albertvillanova/tests-raw-jsonl/resolve/main/\"\nURLS = {\"train\": REPO_URL + \"wikiann-bn-train.jsonl\", \"validation\": REPO_URL + \"wikiann-bn-validation.jsonl\"}\n\n\nclass __DummyDataset1__(datasets.GeneratorBasedBuilder):\n\n    def _info(self):\n        features = datasets.Features(\n            {\n                \"tokens\": datasets.Sequence(datasets.Value(\"string\")),\n                \"ner_tags\": datasets.Sequence(\n                    datasets.features.ClassLabel(\n                        names=[\n                            \"O\",\n                            \"B-PER\",\n                            \"I-PER\",\n                            \"B-ORG\",\n                            \"I-ORG\",\n                            \"B-LOC\",\n                            \"I-LOC\",\n                        ]\n                    )\n                ),\n                \"langs\": datasets.Sequence(datasets.Value(\"string\")),\n                \"spans\": datasets.Sequence(datasets.Value(\"string\")),\n            }\n        )\n        return datasets.DatasetInfo(features=features)\n\n    def _split_generators(self, dl_manager):\n        dl_path = dl_manager.download(URLS)\n        return [\n            datasets.SplitGenerator(datasets.Split.TRAIN, gen_kwargs={\"filepath\": dl_path[\"train\"]}),\n            datasets.SplitGenerator(datasets.Split.VALIDATION, gen_kwargs={\"filepath\": dl_path[\"validation\"]}),\n        ]\n\n    def _generate_examples(self, filepath):\n        with open(filepath, \"r\", encoding=\"utf-8\") as f:\n            for i, line in enumerate(f):\n                yield i, json.loads(line)\n"





@pytest.fixture
def 					lowercase_ (    )     -> Optional[Any]:





				'''simple docstring'''


				return DATASET_LOADING_SCRIPT_NAME





@pytest.fixture
def 					lowercase_ (    )     -> Optional[int]:





				'''simple docstring'''


				return DATASET_LOADING_SCRIPT_CODE





@pytest.fixture
def 					lowercase_ (    __snake_case     :							Optional[int]	,	__snake_case     :							List[Any]	,	__snake_case     :							Any      )     -> Dict:





				'''simple docstring'''


				snake_case__      :Optional[Any]         =   dataset_loading_script_name
				snake_case__      :Optional[Any]         =   tmp_path / "datasets" / script_name
				script_dir.mkdir(parents=__snake_case      )
				snake_case__      :List[Any]         =   script_dir / F'{script_name}.py'
				with open(__snake_case	,	"w"      ) as f:
								f.write(__snake_case      )
				return str(__snake_case      )