





import gc
import unittest

import torch
from parameterized import parameterized

from diffusers import AutoencoderKL
from diffusers.utils import floats_tensor, load_hf_numpy, require_torch_gpu, slow, torch_all_close, torch_device
from diffusers.utils.import_utils import is_xformers_available
from diffusers.utils.testing_utils import enable_full_determinism

from .test_modeling_common import ModelTesterMixin, UNetTesterMixin


enable_full_determinism()




class   UpperCamelCase_       (     UpperCamelCase__				,  UpperCamelCase__				,  unittest.TestCase    ):
							lowerCamelCase_           =					AutoencoderKL
							lowerCamelCase_           =					"sample"
							lowerCamelCase_           =					1e-2


							@property
							def      _snake_case  (							self    :List[str]		)   ->   List[str]:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__			=      4
									SCREAMING_SNAKE_CASE__			=      3
									SCREAMING_SNAKE_CASE__			=      (32, 32)

									SCREAMING_SNAKE_CASE__			=      floats_tensor((batch_size, num_channels) + sizes		).to(__A		)

									return {"sample": image}


							@property
							def      _snake_case  (							self    :Any		)   ->   List[Any]:



									"""simple docstring"""



									return (3, 32, 32)


							@property
							def      _snake_case  (							self    :Optional[Any]		)   ->   Union[str, Any]:



									"""simple docstring"""



									return (3, 32, 32)


							def      _snake_case  (							self    :str		)   ->   Tuple:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__			=      {
									    """block_out_channels""": [32, 64],
									    """in_channels""": 3,
									    """out_channels""": 3,
									    """down_block_types""": ["""DownEncoderBlock2D""", """DownEncoderBlock2D"""],
									    """up_block_types""": ["""UpDecoderBlock2D""", """UpDecoderBlock2D"""],
									    """latent_channels""": 4,
									}
									SCREAMING_SNAKE_CASE__			=      self.dummy_input
									return init_dict, inputs_dict


							def      _snake_case  (							self    :List[Any]		)   ->   Optional[int]:



									"""simple docstring"""



									pass


							def      _snake_case  (							self    :Optional[Any]		)   ->   List[Any]:



									"""simple docstring"""



									pass


							@unittest.skipIf(torch_device == """mps"""							,				"""Gradient checkpointing skipped on MPS"""		)
							def      _snake_case  (							self    :Any		)   ->   Union[str, Any]:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__	,						SCREAMING_SNAKE_CASE__			=      self.prepare_init_args_and_inputs_for_common()
									SCREAMING_SNAKE_CASE__			=      self.model_class(**__A		)
									model.to(__A		)

									assert not model.is_gradient_checkpointing and model.training

									SCREAMING_SNAKE_CASE__			=      model(**__A		).sample
									# run the backwards pass on the model. For backwards pass, for simplicity purpose,
									# we won't calculate the loss and rather backprop on out.sum()
									model.zero_grad()

									SCREAMING_SNAKE_CASE__			=      torch.randn_like(__A		)
									SCREAMING_SNAKE_CASE__			=      (out - labels).mean()
									loss.backward()

									# re-instantiate the model now enabling gradient checkpointing
									SCREAMING_SNAKE_CASE__			=      self.model_class(**__A		)
									# clone model
									model_a.load_state_dict(model.state_dict()		)
									model_a.to(__A		)
									model_a.enable_gradient_checkpointing()

									assert model_a.is_gradient_checkpointing and model_a.training

									SCREAMING_SNAKE_CASE__			=      model_a(**__A		).sample
									# run the backwards pass on the model. For backwards pass, for simplicity purpose,
									# we won't calculate the loss and rather backprop on out.sum()
									model_a.zero_grad()
									SCREAMING_SNAKE_CASE__			=      (out_a - labels).mean()
									loss_a.backward()

									# compare the output and parameters gradients
									self.assertTrue((loss - loss_a).abs() < 1E-5		)
									SCREAMING_SNAKE_CASE__			=      dict(model.named_parameters()		)
									SCREAMING_SNAKE_CASE__			=      dict(model_a.named_parameters()		)
									for name, param in named_params.items():
											self.assertTrue(torch_all_close(param.grad.data							,				named_params_a[name].grad.data							,				atol=5E-5		)		)


							def      _snake_case  (							self    :Any		)   ->   str:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__	,						SCREAMING_SNAKE_CASE__			=      AutoencoderKL.from_pretrained("""fusing/autoencoder-kl-dummy"""							,				output_loading_info=__A		)
									self.assertIsNotNone(__A		)
									self.assertEqual(len(loading_info["""missing_keys"""]		)							,				0		)

									model.to(__A		)
									SCREAMING_SNAKE_CASE__			=      model(**self.dummy_input		)

									assert image is not None, "Make sure output is not None"


							def      _snake_case  (							self    :List[Any]		)   ->   List[Any]:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__			=      AutoencoderKL.from_pretrained("""fusing/autoencoder-kl-dummy"""		)
									SCREAMING_SNAKE_CASE__			=      model.to(__A		)
									model.eval()

									if torch_device == "mps":
											SCREAMING_SNAKE_CASE__			=      torch.manual_seed(0		)
									else:
											SCREAMING_SNAKE_CASE__			=      torch.Generator(device=__A		).manual_seed(0		)

									SCREAMING_SNAKE_CASE__			=      torch.randn(
									    1							,				model.config.in_channels							,				model.config.sample_size							,				model.config.sample_size							,				generator=torch.manual_seed(0		)							,				)
									SCREAMING_SNAKE_CASE__			=      image.to(__A		)
									with torch.no_grad():
											SCREAMING_SNAKE_CASE__			=      model(__A							,				sample_posterior=__A							,				generator=__A		).sample

									SCREAMING_SNAKE_CASE__			=      output[0, -1, -3:, -3:].flatten().cpu()

									# Since the VAE Gaussian prior's generator is seeded on the appropriate device,
									# the expected output slices are not the same for CPU and GPU.
									if torch_device == "mps":
											SCREAMING_SNAKE_CASE__			=      torch.tensor(
											    [
											        -4.0_078E-01,
											        -3.8_323E-04,
											        -1.2_681E-01,
											        -1.1_462E-01,
											        2.0_095E-01,
											        1.0_893E-01,
											        -8.8_247E-02,
											        -3.0_361E-01,
											        -9.8_644E-03,
											    ]		)
									elif torch_device == "cpu":
											SCREAMING_SNAKE_CASE__			=      torch.tensor(
											    [-0.1_3_5_2, 0.0_8_7_8, 0.0_4_1_9, -0.0_8_1_8, -0.1_0_6_9, 0.0_6_8_8, -0.1_4_5_8, -0.4_4_4_6, -0.0_0_2_6]		)
									else:
											SCREAMING_SNAKE_CASE__			=      torch.tensor(
											    [-0.2_4_2_1, 0.4_6_4_2, 0.2_5_0_7, -0.0_4_3_8, 0.0_6_8_2, 0.3_1_6_0, -0.2_0_1_8, -0.0_7_2_7, 0.2_4_8_5]		)

									self.assertTrue(torch_all_close(__A							,				__A							,				rtol=1E-2		)		)




@slow
class   UpperCamelCase_       (     unittest.TestCase    ):


							def      _snake_case  (							self    :Optional[int]							,				__A    :int							,				__A    :Optional[Any]		)   ->   Tuple:



									"""simple docstring"""



									return f'''gaussian_noise_s={seed}_shape={'_'.join([str(__A		) for s in shape]		)}.npy'''


							def      _snake_case  (							self    :List[str]		)   ->   int:



									"""simple docstring"""



									super().tearDown()
									gc.collect()
									torch.cuda.empty_cache()


							def      _snake_case  (							self    :int							,				__A    :Any=0							,				__A    :List[str]=(4, 3, 512, 512)							,				__A    :int=False		)   ->   Any:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__			=      torch.floataa if fpaa else torch.floataa
									SCREAMING_SNAKE_CASE__			=      torch.from_numpy(load_hf_numpy(self.get_file_format(__A							,				__A		)		)		).to(__A		).to(__A		)
									return image


							def      _snake_case  (							self    :Tuple							,				__A    :int="CompVis/stable-diffusion-v1-4"							,				__A    :Any=False		)   ->   List[str]:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__			=      """fp16""" if fpaa else None
									SCREAMING_SNAKE_CASE__			=      torch.floataa if fpaa else torch.floataa

									SCREAMING_SNAKE_CASE__			=      AutoencoderKL.from_pretrained(
									    __A							,				subfolder="""vae"""							,				torch_dtype=__A							,				revision=__A							,				)
									model.to(__A		).eval()

									return model


							def      _snake_case  (							self    :Optional[int]							,				__A    :Optional[Any]=0		)   ->   Optional[int]:



									"""simple docstring"""



									if torch_device == "mps":
											return torch.manual_seed(__A		)
									return torch.Generator(device=__A		).manual_seed(__A		)


							@parameterized.expand(
							    [
							        # fmt: off
							        [33, [-0.1_6_0_3, 0.9_8_7_8, -0.0_4_9_5, -0.0_7_9_0, -0.2_7_0_9, 0.8_3_7_5, -0.2_0_6_0, -0.0_8_2_4], [-0.2_3_9_5, 0.0_0_9_8, 0.0_1_0_2, -0.0_7_0_9, -0.2_8_4_0, -0.0_2_7_4, -0.0_7_1_8, -0.1_8_2_4]],
							        [47, [-0.2_3_7_6, 0.1_1_6_8, 0.1_3_3_2, -0.4_8_4_0, -0.2_5_0_8, -0.0_7_9_1, -0.0_4_9_3, -0.4_0_8_9], [0.0_3_5_0, 0.0_8_4_7, 0.0_4_6_7, 0.0_3_4_4, -0.0_8_4_2, -0.0_5_4_7, -0.0_6_3_3, -0.1_1_3_1]],
							        # fmt: on
							    ]		)
							def      _snake_case  (							self    :List[Any]							,				__A    :int							,				__A    :str							,				__A    :Optional[int]		)   ->   Dict:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__			=      self.get_sd_vae_model()
									SCREAMING_SNAKE_CASE__			=      self.get_sd_image(__A		)
									SCREAMING_SNAKE_CASE__			=      self.get_generator(__A		)

									with torch.no_grad():
											SCREAMING_SNAKE_CASE__			=      model(__A							,				generator=__A							,				sample_posterior=__A		).sample

									assert sample.shape == image.shape

									SCREAMING_SNAKE_CASE__			=      sample[-1, -2:, -2:, :2].flatten().float().cpu()
									SCREAMING_SNAKE_CASE__			=      torch.tensor(expected_slice_mps if torch_device == """mps""" else expected_slice		)

									assert torch_all_close(__A							,				__A							,				atol=3E-3		)


							@parameterized.expand(
							    [
							        # fmt: off
							        [33, [-0.0_5_1_3, 0.0_2_8_9, 1.3_7_9_9, 0.2_1_6_6, -0.2_5_7_3, -0.0_8_7_1, 0.5_1_0_3, -0.0_9_9_9]],
							        [47, [-0.4_1_2_8, -0.1_3_2_0, -0.3_7_0_4, 0.1_9_6_5, -0.4_1_1_6, -0.2_3_3_2, -0.3_3_4_0, 0.2_2_4_7]],
							        # fmt: on
							    ]		)
							@require_torch_gpu
							def      _snake_case  (							self    :str							,				__A    :Optional[Any]							,				__A    :Any		)   ->   List[Any]:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__			=      self.get_sd_vae_model(fpaa=__A		)
									SCREAMING_SNAKE_CASE__			=      self.get_sd_image(__A							,				fpaa=__A		)
									SCREAMING_SNAKE_CASE__			=      self.get_generator(__A		)

									with torch.no_grad():
											SCREAMING_SNAKE_CASE__			=      model(__A							,				generator=__A							,				sample_posterior=__A		).sample

									assert sample.shape == image.shape

									SCREAMING_SNAKE_CASE__			=      sample[-1, -2:, :2, -2:].flatten().float().cpu()
									SCREAMING_SNAKE_CASE__			=      torch.tensor(__A		)

									assert torch_all_close(__A							,				__A							,				atol=1E-2		)


							@parameterized.expand(
							    [
							        # fmt: off
							        [33, [-0.1_6_0_9, 0.9_8_6_6, -0.0_4_8_7, -0.0_7_7_7, -0.2_7_1_6, 0.8_3_6_8, -0.2_0_5_5, -0.0_8_1_4], [-0.2_3_9_5, 0.0_0_9_8, 0.0_1_0_2, -0.0_7_0_9, -0.2_8_4_0, -0.0_2_7_4, -0.0_7_1_8, -0.1_8_2_4]],
							        [47, [-0.2_3_7_7, 0.1_1_4_7, 0.1_3_3_3, -0.4_8_4_1, -0.2_5_0_6, -0.0_8_0_5, -0.0_4_9_1, -0.4_0_8_5], [0.0_3_5_0, 0.0_8_4_7, 0.0_4_6_7, 0.0_3_4_4, -0.0_8_4_2, -0.0_5_4_7, -0.0_6_3_3, -0.1_1_3_1]],
							        # fmt: on
							    ]		)
							def      _snake_case  (							self    :Any							,				__A    :Tuple							,				__A    :List[str]							,				__A    :str		)   ->   List[str]:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__			=      self.get_sd_vae_model()
									SCREAMING_SNAKE_CASE__			=      self.get_sd_image(__A		)

									with torch.no_grad():
											SCREAMING_SNAKE_CASE__			=      model(__A		).sample

									assert sample.shape == image.shape

									SCREAMING_SNAKE_CASE__			=      sample[-1, -2:, -2:, :2].flatten().float().cpu()
									SCREAMING_SNAKE_CASE__			=      torch.tensor(expected_slice_mps if torch_device == """mps""" else expected_slice		)

									assert torch_all_close(__A							,				__A							,				atol=3E-3		)


							@parameterized.expand(
							    [
							        # fmt: off
							        [13, [-0.2_0_5_1, -0.1_8_0_3, -0.2_3_1_1, -0.2_1_1_4, -0.3_2_9_2, -0.3_5_7_4, -0.2_9_5_3, -0.3_3_2_3]],
							        [37, [-0.2_6_3_2, -0.2_6_2_5, -0.2_1_9_9, -0.2_7_4_1, -0.4_5_3_9, -0.4_9_9_0, -0.3_7_2_0, -0.4_9_2_5]],
							        # fmt: on
							    ]		)
							@require_torch_gpu
							def      _snake_case  (							self    :Dict							,				__A    :Any							,				__A    :List[Any]		)   ->   Union[str, Any]:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__			=      self.get_sd_vae_model()
									SCREAMING_SNAKE_CASE__			=      self.get_sd_image(__A							,				shape=(3, 4, 64, 64)		)

									with torch.no_grad():
											SCREAMING_SNAKE_CASE__			=      model.decode(__A		).sample

									assert list(sample.shape		) == [3, 3, 512, 512]

									SCREAMING_SNAKE_CASE__			=      sample[-1, -2:, :2, -2:].flatten().cpu()
									SCREAMING_SNAKE_CASE__			=      torch.tensor(__A		)

									assert torch_all_close(__A							,				__A							,				atol=1E-3		)


							@parameterized.expand(
							    [
							        # fmt: off
							        [27, [-0.0_3_6_9, 0.0_2_0_7, -0.0_7_7_6, -0.0_6_8_2, -0.1_7_4_7, -0.1_9_3_0, -0.1_4_6_5, -0.2_0_3_9]],
							        [16, [-0.1_6_2_8, -0.2_1_3_4, -0.2_7_4_7, -0.2_6_4_2, -0.3_7_7_4, -0.4_4_0_4, -0.3_6_8_7, -0.4_2_7_7]],
							        # fmt: on
							    ]		)
							@require_torch_gpu
							def      _snake_case  (							self    :int							,				__A    :Any							,				__A    :Optional[int]		)   ->   List[Any]:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__			=      self.get_sd_vae_model(fpaa=__A		)
									SCREAMING_SNAKE_CASE__			=      self.get_sd_image(__A							,				shape=(3, 4, 64, 64)							,				fpaa=__A		)

									with torch.no_grad():
											SCREAMING_SNAKE_CASE__			=      model.decode(__A		).sample

									assert list(sample.shape		) == [3, 3, 512, 512]

									SCREAMING_SNAKE_CASE__			=      sample[-1, -2:, :2, -2:].flatten().float().cpu()
									SCREAMING_SNAKE_CASE__			=      torch.tensor(__A		)

									assert torch_all_close(__A							,				__A							,				atol=5E-3		)


							@parameterized.expand([(13,), (16,), (27,)]		)
							@require_torch_gpu
							@unittest.skipIf(not is_xformers_available()							,				reason="""xformers is not required when using PyTorch 2.0."""		)
							def      _snake_case  (							self    :Optional[int]							,				__A    :Any		)   ->   int:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__			=      self.get_sd_vae_model(fpaa=__A		)
									SCREAMING_SNAKE_CASE__			=      self.get_sd_image(__A							,				shape=(3, 4, 64, 64)							,				fpaa=__A		)

									with torch.no_grad():
											SCREAMING_SNAKE_CASE__			=      model.decode(__A		).sample

									model.enable_xformers_memory_efficient_attention()
									with torch.no_grad():
											SCREAMING_SNAKE_CASE__			=      model.decode(__A		).sample

									assert list(sample.shape		) == [3, 3, 512, 512]

									assert torch_all_close(__A							,				__A							,				atol=1E-1		)


							@parameterized.expand([(13,), (16,), (37,)]		)
							@require_torch_gpu
							@unittest.skipIf(not is_xformers_available()							,				reason="""xformers is not required when using PyTorch 2.0."""		)
							def      _snake_case  (							self    :str							,				__A    :int		)   ->   Optional[int]:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__			=      self.get_sd_vae_model()
									SCREAMING_SNAKE_CASE__			=      self.get_sd_image(__A							,				shape=(3, 4, 64, 64)		)

									with torch.no_grad():
											SCREAMING_SNAKE_CASE__			=      model.decode(__A		).sample

									model.enable_xformers_memory_efficient_attention()
									with torch.no_grad():
											SCREAMING_SNAKE_CASE__			=      model.decode(__A		).sample

									assert list(sample.shape		) == [3, 3, 512, 512]

									assert torch_all_close(__A							,				__A							,				atol=1E-2		)


							@parameterized.expand(
							    [
							        # fmt: off
							        [33, [-0.3_0_0_1, 0.0_9_1_8, -2.6_9_8_4, -3.9_7_2_0, -3.2_0_9_9, -5.0_3_5_3, 1.7_3_3_8, -0.2_0_6_5, 3.4_2_6_7]],
							        [47, [-1.5_0_3_0, -4.3_8_7_1, -6.0_3_5_5, -9.1_1_5_7, -1.6_6_6_1, -2.7_8_5_3, 2.1_6_0_7, -5.0_8_2_3, 2.5_6_3_3]],
							        # fmt: on
							    ]		)
							def      _snake_case  (							self    :Tuple							,				__A    :Union[str, Any]							,				__A    :Any		)   ->   Union[str, Any]:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__			=      self.get_sd_vae_model()
									SCREAMING_SNAKE_CASE__			=      self.get_sd_image(__A		)
									SCREAMING_SNAKE_CASE__			=      self.get_generator(__A		)

									with torch.no_grad():
											SCREAMING_SNAKE_CASE__			=      model.encode(__A		).latent_dist
											SCREAMING_SNAKE_CASE__			=      dist.sample(generator=__A		)

									assert list(sample.shape		) == [image.shape[0], 4] + [i // 8 for i in image.shape[2:]]

									SCREAMING_SNAKE_CASE__			=      sample[0, -1, -3:, -3:].flatten().cpu()
									SCREAMING_SNAKE_CASE__			=      torch.tensor(__A		)

									SCREAMING_SNAKE_CASE__			=      3E-3 if torch_device != """mps""" else 1E-2
									assert torch_all_close(__A							,				__A							,				atol=__A		)






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






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

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

				SCREAMING_SNAKE_CASE__			=      load_tf_weights_in_transfo_xl(UpperCamelCase__					,   UpperCamelCase__					,   UpperCamelCase__    )
				# Save pytorch-model
				SCREAMING_SNAKE_CASE__			=      os.path.join(UpperCamelCase__					,   UpperCamelCase__    )
				SCREAMING_SNAKE_CASE__			=      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__":
	_lowerCamelCase   				=			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.',
	)
	_lowerCamelCase   				=			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,
	)