
"""simple docstring"""

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 (		a_,						a_,						a_,						a_      ):





		'''simple docstring'''


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

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

				lowerCamelCase					:      Any          =   load_tf_weights_in_transfo_xl(a_,						a_,						a_      )
				# Save pytorch-model
				lowerCamelCase					:      Dict          =   os.path.join(a_,						a_      )
				lowerCamelCase					:      Tuple          =   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"""

# Lint as: python3
import dataclasses
import re
from dataclasses import dataclass
from functools import total_ordering
from typing import Optional, Union


_A						=			re.compile(r'^(?P<major>\d+)' r'\.(?P<minor>\d+)' r'\.(?P<patch>\d+)$')






@total_ordering
@dataclass
class 	_lowercase       :
	lowercase_				=     42
	lowercase_				=     None
	lowercase_				=     None
	lowercase_				=     None
	lowercase_				=     None




	def       _UpperCamelCase   (			self )		->							List[str]:
			lowerCamelCase     ,   lowerCamelCase     ,   lowerCamelCase					:      List[str]          =   _str_to_version_tuple(self.version_str )




	def __repr__(			self )		->							int:
			return F"""{self.tuple[0]}.{self.tuple[1]}.{self.tuple[2]}"""




	@property
	def       _UpperCamelCase   (			self )		->							Dict:
			return self.major, self.minor, self.patch




	def       _UpperCamelCase   (			self		,			UpperCAmelCase_ )		->							Optional[int]:
			if isinstance(UpperCAmelCase_		,			UpperCAmelCase_ ):
					return Version(UpperCAmelCase_ )
			elif isinstance(UpperCAmelCase_		,			UpperCAmelCase_ ):
					return other
			raise TypeError(F"""{other} (type {type(UpperCAmelCase_ )}) cannot be compared to version.""" )




	def __eq__(			self		,			UpperCAmelCase_ )		->							Optional[Any]:
			try:
					lowerCamelCase					:      List[str]          =   self._validate_operand(UpperCAmelCase_ )
			except (TypeError, ValueError):
					return False
			else:
					return self.tuple == other.tuple




	def __lt__(			self		,			UpperCAmelCase_ )		->							Optional[int]:
			lowerCamelCase					:      Optional[int]          =   self._validate_operand(UpperCAmelCase_ )
			return self.tuple < other.tuple




	def __hash__(			self )		->							Optional[Any]:
			return hash(_version_tuple_to_str(self.tuple ) )




	@classmethod
	def       _UpperCamelCase   (			cls		,			UpperCAmelCase_ )		->							Union[str, Any]:
			lowerCamelCase					:      str          =   {f.name for f in dataclasses.fields(cls )}
			return cls(**{k: v for k, v in dic.items() if k in field_names} )




	def       _UpperCamelCase   (			self )		->							str:
			return self.version_str


def 		UpperCAmelCase (		a_      ):





		'''simple docstring'''


		lowerCamelCase					:      Tuple          =   _VERSION_REG.match(a_      )
		if not res:
				raise ValueError(F"""Invalid version '{version_str}'. Format should be x.y.z with {{x,y,z}} being digits."""      )
		return tuple(int(a_      ) for v in [res.group('major'      ), res.group('minor'      ), res.group('patch'      )]      )





def 		UpperCAmelCase (		a_      ):





		'''simple docstring'''


		return ".".join(str(a_      ) for v in version_tuple      )






