


"""simple docstring"""
import json
import os
import shutil
import tempfile
import unittest

import numpy as np
import pytest

from transformers import BertTokenizer, BertTokenizerFast
from transformers.models.bert.tokenization_bert import VOCAB_FILES_NAMES
from transformers.testing_utils import require_vision
from transformers.utils import IMAGE_PROCESSOR_NAME, is_vision_available


if is_vision_available():
				from PIL import Image

				from transformers import AlignProcessor, EfficientNetImageProcessor





@require_vision
class      __A     (unittest.TestCase):





							'''simple docstring'''

							def 		lowerCAmelCase       (				self							:      Union[str, Any]       )					->int:



														"""simple docstring"""



														snake_case_           =		tempfile.mkdtemp()

														snake_case_           =		[
														    """[UNK]""",
														    """[CLS]""",
														    """[SEP]""",
														    """[PAD]""",
														    """[MASK]""",
														    """want""",
														    """##want""",
														    """##ed""",
														    """wa""",
														    """un""",
														    """runn""",
														    """##ing""",
														    """,""",
														    """low""",
														    """lowest""",
														]
														snake_case_           =		os.path.join(self.tmpdirname		,       VOCAB_FILES_NAMES["""vocab_file"""]       )
														with open(self.vocab_file		,       """w"""		,       encoding="""utf-8"""       ) as vocab_writer:
																					vocab_writer.write("""""".join([x + """\n""" for x in vocab_tokens]       )       )

														snake_case_           =		{
														    """do_resize""": True,
														    """size""": 20,
														    """do_center_crop""": True,
														    """crop_size""": 18,
														    """do_normalize""": True,
														    """image_mean""": [0.48_145_466, 0.4_578_275, 0.40_821_073],
														    """image_std""": [0.26_862_954, 0.26_130_258, 0.27_577_711],
														}
														snake_case_           =		os.path.join(self.tmpdirname		,       UpperCAmelCase_       )
														with open(self.image_processor_file		,       """w"""		,       encoding="""utf-8"""       ) as fp:
																					json.dump(UpperCAmelCase_		,       UpperCAmelCase_       )



							def 		lowerCAmelCase       (				self							:      List[str]		,       **UpperCAmelCase_							:      Dict       )					->Dict:



														"""simple docstring"""



														return BertTokenizer.from_pretrained(self.tmpdirname		,       **UpperCAmelCase_       )



							def 		lowerCAmelCase       (				self							:      Optional[int]		,       **UpperCAmelCase_							:      Tuple       )					->str:



														"""simple docstring"""



														return BertTokenizerFast.from_pretrained(self.tmpdirname		,       **UpperCAmelCase_       )



							def 		lowerCAmelCase       (				self							:      Optional[int]		,       **UpperCAmelCase_							:      Tuple       )					->Dict:



														"""simple docstring"""



														return EfficientNetImageProcessor.from_pretrained(self.tmpdirname		,       **UpperCAmelCase_       )



							def 		lowerCAmelCase       (				self							:      int       )					->Dict:



														"""simple docstring"""



														shutil.rmtree(self.tmpdirname       )



							def 		lowerCAmelCase       (				self							:      List[str]       )					->List[str]:



														"""simple docstring"""



														snake_case_           =		[np.random.randint(255		,       size=(3, 30, 400)		,       dtype=np.uinta       )]
														snake_case_           =		[Image.fromarray(np.moveaxis(UpperCAmelCase_		,       0		,       -1       )       ) for x in image_inputs]
														return image_inputs



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



														"""simple docstring"""



														snake_case_           =		self.get_tokenizer()
														snake_case_           =		self.get_rust_tokenizer()
														snake_case_           =		self.get_image_processor()

														snake_case_           =		AlignProcessor(tokenizer=UpperCAmelCase_		,       image_processor=UpperCAmelCase_       )
														processor_slow.save_pretrained(self.tmpdirname       )
														snake_case_           =		AlignProcessor.from_pretrained(self.tmpdirname		,       use_fast=UpperCAmelCase_       )

														snake_case_           =		AlignProcessor(tokenizer=UpperCAmelCase_		,       image_processor=UpperCAmelCase_       )
														processor_fast.save_pretrained(self.tmpdirname       )
														snake_case_           =		AlignProcessor.from_pretrained(self.tmpdirname       )

														self.assertEqual(processor_slow.tokenizer.get_vocab()		,       tokenizer_slow.get_vocab()       )
														self.assertEqual(processor_fast.tokenizer.get_vocab()		,       tokenizer_fast.get_vocab()       )
														self.assertEqual(tokenizer_slow.get_vocab()		,       tokenizer_fast.get_vocab()       )
														self.assertIsInstance(processor_slow.tokenizer		,       UpperCAmelCase_       )
														self.assertIsInstance(processor_fast.tokenizer		,       UpperCAmelCase_       )

														self.assertEqual(processor_slow.image_processor.to_json_string()		,       image_processor.to_json_string()       )
														self.assertEqual(processor_fast.image_processor.to_json_string()		,       image_processor.to_json_string()       )
														self.assertIsInstance(processor_slow.image_processor		,       UpperCAmelCase_       )
														self.assertIsInstance(processor_fast.image_processor		,       UpperCAmelCase_       )



							def 		lowerCAmelCase       (				self							:      int       )					->int:



														"""simple docstring"""



														snake_case_           =		AlignProcessor(tokenizer=self.get_tokenizer()		,       image_processor=self.get_image_processor()       )
														processor.save_pretrained(self.tmpdirname       )

														snake_case_           =		self.get_tokenizer(bos_token="""(BOS)"""		,       eos_token="""(EOS)"""       )
														snake_case_           =		self.get_image_processor(do_normalize=UpperCAmelCase_		,       padding_value=1.0       )

														snake_case_           =		AlignProcessor.from_pretrained(
														    self.tmpdirname		,       bos_token="""(BOS)"""		,       eos_token="""(EOS)"""		,       do_normalize=UpperCAmelCase_		,       padding_value=1.0       )

														self.assertEqual(processor.tokenizer.get_vocab()		,       tokenizer_add_kwargs.get_vocab()       )
														self.assertIsInstance(processor.tokenizer		,       UpperCAmelCase_       )

														self.assertEqual(processor.image_processor.to_json_string()		,       image_processor_add_kwargs.to_json_string()       )
														self.assertIsInstance(processor.image_processor		,       UpperCAmelCase_       )



							def 		lowerCAmelCase       (				self							:      Tuple       )					->List[str]:



														"""simple docstring"""



														snake_case_           =		self.get_image_processor()
														snake_case_           =		self.get_tokenizer()

														snake_case_           =		AlignProcessor(tokenizer=UpperCAmelCase_		,       image_processor=UpperCAmelCase_       )

														snake_case_           =		self.prepare_image_inputs()

														snake_case_           =		image_processor(UpperCAmelCase_		,       return_tensors="""np"""       )
														snake_case_           =		processor(images=UpperCAmelCase_		,       return_tensors="""np"""       )

														for key in input_image_proc.keys():
																					self.assertAlmostEqual(input_image_proc[key].sum()		,       input_processor[key].sum()		,       delta=1E-2       )



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



														"""simple docstring"""



														snake_case_           =		self.get_image_processor()
														snake_case_           =		self.get_tokenizer()

														snake_case_           =		AlignProcessor(tokenizer=UpperCAmelCase_		,       image_processor=UpperCAmelCase_       )

														snake_case_           =		"""lower newer"""

														snake_case_           =		processor(text=UpperCAmelCase_       )

														snake_case_           =		tokenizer(UpperCAmelCase_		,       padding="""max_length"""		,       max_length=64       )

														for key in encoded_tok.keys():
																					self.assertListEqual(encoded_tok[key]		,       encoded_processor[key]       )



							def 		lowerCAmelCase       (				self							:      Dict       )					->Optional[Any]:



														"""simple docstring"""



														snake_case_           =		self.get_image_processor()
														snake_case_           =		self.get_tokenizer()

														snake_case_           =		AlignProcessor(tokenizer=UpperCAmelCase_		,       image_processor=UpperCAmelCase_       )

														snake_case_           =		"""lower newer"""
														snake_case_           =		self.prepare_image_inputs()

														snake_case_           =		processor(text=UpperCAmelCase_		,       images=UpperCAmelCase_       )

														self.assertListEqual(list(inputs.keys()       )		,       ["""input_ids""", """token_type_ids""", """attention_mask""", """pixel_values"""]       )

														# test if it raises when no input is passed
														with pytest.raises(UpperCAmelCase_       ):
																					processor()



							def 		lowerCAmelCase       (				self							:      int       )					->int:



														"""simple docstring"""



														snake_case_           =		self.get_image_processor()
														snake_case_           =		self.get_tokenizer()

														snake_case_           =		AlignProcessor(tokenizer=UpperCAmelCase_		,       image_processor=UpperCAmelCase_       )

														snake_case_           =		[[1, 4, 5, 8, 1, 0, 8], [3, 4, 3, 1, 1, 8, 9]]

														snake_case_           =		processor.batch_decode(UpperCAmelCase_       )
														snake_case_           =		tokenizer.batch_decode(UpperCAmelCase_       )

														self.assertListEqual(UpperCAmelCase_		,       UpperCAmelCase_       )



							def 		lowerCAmelCase       (				self							:      int       )					->Union[str, Any]:



														"""simple docstring"""



														snake_case_           =		self.get_image_processor()
														snake_case_           =		self.get_tokenizer()

														snake_case_           =		AlignProcessor(tokenizer=UpperCAmelCase_		,       image_processor=UpperCAmelCase_       )

														snake_case_           =		"""lower newer"""
														snake_case_           =		self.prepare_image_inputs()

														snake_case_           =		processor(text=UpperCAmelCase_		,       images=UpperCAmelCase_       )

														self.assertListEqual(list(inputs.keys()       )		,       processor.model_input_names       )







"""simple docstring"""
import argparse
import json
import os
from pathlib import Path

import requests
import torch

from transformers import JukeboxConfig, JukeboxModel
from transformers.utils import logging


logging.set_verbosity_info()
__SCREAMING_SNAKE_CASE							:    str		=    logging.get_logger(__name__)


__SCREAMING_SNAKE_CASE							:    Optional[int]		=    'https://openaipublic.azureedge.net/jukebox/models/'
__SCREAMING_SNAKE_CASE							:    List[Any]		=    {
    'jukebox-1b-lyrics': [
        '5b/vqvae.pth.tar',
        '5b/prior_level_0.pth.tar',
        '5b/prior_level_1.pth.tar',
        '1b_lyrics/prior_level_2.pth.tar',
    ],
    'jukebox-5b-lyrics': [
        '5b/vqvae.pth.tar',
        '5b/prior_level_0.pth.tar',
        '5b/prior_level_1.pth.tar',
        '5b_lyrics/prior_level_2.pth.tar',
    ],
}







def 						_a				(      _SCREAMING_SNAKE_CASE				)    ->				int:
							if key.endswith(""".model.1.bias"""				) and len(key.split("""."""				)				) > 10:
														snake_case_           =		key.replace(""".model.1.bias"""				, """.conv1d_1.bias"""				)
							elif key.endswith(""".model.1.weight"""				) and len(key.split("""."""				)				) > 10:
														snake_case_           =		key.replace(""".model.1.weight"""				, """.conv1d_1.weight"""				)
							elif key.endswith(""".model.3.bias"""				) and len(key.split("""."""				)				) > 10:
														snake_case_           =		key.replace(""".model.3.bias"""				, """.conv1d_2.bias"""				)
							elif key.endswith(""".model.3.weight"""				) and len(key.split("""."""				)				) > 10:
														snake_case_           =		key.replace(""".model.3.weight"""				, """.conv1d_2.weight"""				)

							if "conditioner_blocks.0." in key:
														snake_case_           =		key.replace("""conditioner_blocks.0"""				, """conditioner_blocks"""				)

							if "prime_prior" in key:
														snake_case_           =		key.replace("""prime_prior"""				, """encoder"""				)

							if ".emb." in key and "total" not in key and "absolute" not in key and "relative" not in key:
														snake_case_           =		key.replace(""".emb."""				, """."""				)

							if key.endswith("""k"""				):  # replace vqvae.X.k with vqvae.X.codebook
														return key.replace(""".k"""				, """.codebook"""				)
							if "y_emb." in key:
														return key.replace("""y_emb."""				, """metadata_embedding."""				)

							if "x_emb.emb." in key:
														snake_case_           =		key.replace("""0.x_emb.emb"""				, """embed_tokens"""				)

							if "prime_state_ln" in key:
														return key.replace("""prime_state_ln"""				, """encoder.final_layer_norm"""				)
							if ".ln" in key:
														return key.replace(""".ln"""				, """.layer_norm"""				)
							if "_ln" in key:
														return key.replace("""_ln"""				, """_layer_norm"""				)

							if "prime_state_proj" in key:
														return key.replace("""prime_state_proj"""				, """encoder.proj_in"""				)
							if "prime_x_out" in key:
														return key.replace("""prime_x_out"""				, """encoder.lm_head"""				)
							if "prior.x_out" in key:
														return key.replace("""x_out"""				, """fc_proj_out"""				)
							if "x_emb" in key:
														return key.replace("""x_emb"""				, """embed_tokens"""				)

							return key







def 						_a				(      _SCREAMING_SNAKE_CASE				, _SCREAMING_SNAKE_CASE				, _SCREAMING_SNAKE_CASE				, _SCREAMING_SNAKE_CASE				)    ->				Union[str, Any]:
							snake_case_           =		{}
							import re

							snake_case_           =		re.compile(r"""encoders.(\d*).level_blocks.(\d*).model.(\d*).(\d).(bias|weight)"""				)
							snake_case_           =		re.compile(
							    r"""encoders.(\d*).level_blocks.(\d*).model.(\d*).(\d).model.(\d*).model.(\d*).(bias|weight)"""				)
							snake_case_           =		re.compile(r"""encoders.(\d*).level_blocks.(\d*).model.(\d*).(bias|weight)"""				)

							snake_case_           =		re.compile(r"""decoders.(\d*).level_blocks.(\d*).model.(\d*).(\d).(bias|weight)"""				)
							snake_case_           =		re.compile(
							    r"""decoders.(\d*).level_blocks.(\d*).model.(\d*).(\d).model.(\d*).model.(\d*).(bias|weight)"""				)
							snake_case_           =		re.compile(r"""decoders.(\d*).level_blocks.(\d*).model.(\d*).(bias|weight)"""				)

							snake_case_           =		re.compile(r"""conditioner_blocks.(\d*).cond.model.(\d*).(\d).(bias|weight)"""				)
							snake_case_           =		re.compile(
							    r"""conditioner_blocks.(\d*).cond.model.(\d*).(\d).model.(\d*).model.(\d*).(bias|weight)"""				)
							snake_case_           =		re.compile(r"""conditioner_blocks.(\d*).cond.model.(\d*).(bias|weight)"""				)

							for original_key, value in state_dict.items():
														# rename vqvae.encoder keys
														if re_encoder_block_conv_in.fullmatch(_SCREAMING_SNAKE_CASE				):
																					snake_case_           =		re_encoder_block_conv_in.match(_SCREAMING_SNAKE_CASE				)
																					snake_case_           =		regex_match.groups()
																					snake_case_           =		int(groups[2]				) * 2 + int(groups[3]				)
																					snake_case_           =		f"""encoders.{groups[0]}.level_blocks.{groups[1]}.downsample_block.{block_index}.{groups[-1]}"""
																					snake_case_           =		re_encoder_block_conv_in.sub(_SCREAMING_SNAKE_CASE				, _SCREAMING_SNAKE_CASE				)

														elif re_encoder_block_resnet.fullmatch(_SCREAMING_SNAKE_CASE				):
																					snake_case_           =		re_encoder_block_resnet.match(_SCREAMING_SNAKE_CASE				)
																					snake_case_           =		regex_match.groups()
																					snake_case_           =		int(groups[2]				) * 2 + int(groups[3]				)
																					snake_case_           =		{"""1""": 1, """3""": 2}[groups[-2]]
																					snake_case_           =		f"""encoders.{groups[0]}.level_blocks.{groups[1]}.downsample_block.{block_index}."""
																					snake_case_           =		f"""resnet_block.{groups[-3]}.conv1d_{conv_index}.{groups[-1]}"""
																					snake_case_           =		prefix + resnet_block
																					snake_case_           =		re_encoder_block_resnet.sub(_SCREAMING_SNAKE_CASE				, _SCREAMING_SNAKE_CASE				)

														elif re_encoder_block_proj_out.fullmatch(_SCREAMING_SNAKE_CASE				):
																					snake_case_           =		re_encoder_block_proj_out.match(_SCREAMING_SNAKE_CASE				)
																					snake_case_           =		regex_match.groups()
																					snake_case_           =		f"""encoders.{groups[0]}.level_blocks.{groups[1]}.proj_out.{groups[-1]}"""
																					snake_case_           =		re_encoder_block_proj_out.sub(_SCREAMING_SNAKE_CASE				, _SCREAMING_SNAKE_CASE				)

														# rename vqvae.decoder keys
														elif re_decoder_block_conv_out.fullmatch(_SCREAMING_SNAKE_CASE				):
																					snake_case_           =		re_decoder_block_conv_out.match(_SCREAMING_SNAKE_CASE				)
																					snake_case_           =		regex_match.groups()
																					snake_case_           =		int(groups[2]				) * 2 + int(groups[3]				) - 2
																					snake_case_           =		f"""decoders.{groups[0]}.level_blocks.{groups[1]}.upsample_block.{block_index}.{groups[-1]}"""
																					snake_case_           =		re_decoder_block_conv_out.sub(_SCREAMING_SNAKE_CASE				, _SCREAMING_SNAKE_CASE				)

														elif re_decoder_block_resnet.fullmatch(_SCREAMING_SNAKE_CASE				):
																					snake_case_           =		re_decoder_block_resnet.match(_SCREAMING_SNAKE_CASE				)
																					snake_case_           =		regex_match.groups()
																					snake_case_           =		int(groups[2]				) * 2 + int(groups[3]				) - 2
																					snake_case_           =		{"""1""": 1, """3""": 2}[groups[-2]]
																					snake_case_           =		f"""decoders.{groups[0]}.level_blocks.{groups[1]}.upsample_block.{block_index}."""
																					snake_case_           =		f"""resnet_block.{groups[-3]}.conv1d_{conv_index}.{groups[-1]}"""
																					snake_case_           =		prefix + resnet_block
																					snake_case_           =		re_decoder_block_resnet.sub(_SCREAMING_SNAKE_CASE				, _SCREAMING_SNAKE_CASE				)

														elif re_decoder_block_proj_in.fullmatch(_SCREAMING_SNAKE_CASE				):
																					snake_case_           =		re_decoder_block_proj_in.match(_SCREAMING_SNAKE_CASE				)
																					snake_case_           =		regex_match.groups()
																					snake_case_           =		f"""decoders.{groups[0]}.level_blocks.{groups[1]}.proj_in.{groups[-1]}"""
																					snake_case_           =		re_decoder_block_proj_in.sub(_SCREAMING_SNAKE_CASE				, _SCREAMING_SNAKE_CASE				)

														# rename prior cond.model to upsampler.upsample_block and resnet
														elif re_prior_cond_conv_out.fullmatch(_SCREAMING_SNAKE_CASE				):
																					snake_case_           =		re_prior_cond_conv_out.match(_SCREAMING_SNAKE_CASE				)
																					snake_case_           =		regex_match.groups()
																					snake_case_           =		int(groups[1]				) * 2 + int(groups[2]				) - 2
																					snake_case_           =		f"""conditioner_blocks.upsampler.upsample_block.{block_index}.{groups[-1]}"""
																					snake_case_           =		re_prior_cond_conv_out.sub(_SCREAMING_SNAKE_CASE				, _SCREAMING_SNAKE_CASE				)

														elif re_prior_cond_resnet.fullmatch(_SCREAMING_SNAKE_CASE				):
																					snake_case_           =		re_prior_cond_resnet.match(_SCREAMING_SNAKE_CASE				)
																					snake_case_           =		regex_match.groups()
																					snake_case_           =		int(groups[1]				) * 2 + int(groups[2]				) - 2
																					snake_case_           =		{"""1""": 1, """3""": 2}[groups[-2]]
																					snake_case_           =		f"""conditioner_blocks.upsampler.upsample_block.{block_index}."""
																					snake_case_           =		f"""resnet_block.{groups[-3]}.conv1d_{conv_index}.{groups[-1]}"""
																					snake_case_           =		prefix + resnet_block
																					snake_case_           =		re_prior_cond_resnet.sub(_SCREAMING_SNAKE_CASE				, _SCREAMING_SNAKE_CASE				)

														elif re_prior_cond_proj_in.fullmatch(_SCREAMING_SNAKE_CASE				):
																					snake_case_           =		re_prior_cond_proj_in.match(_SCREAMING_SNAKE_CASE				)
																					snake_case_           =		regex_match.groups()
																					snake_case_           =		f"""conditioner_blocks.upsampler.proj_in.{groups[-1]}"""
																					snake_case_           =		re_prior_cond_proj_in.sub(_SCREAMING_SNAKE_CASE				, _SCREAMING_SNAKE_CASE				)

														# keep original key
														else:
																					snake_case_           =		original_key

														snake_case_           =		replace_key(_SCREAMING_SNAKE_CASE				)

														if f"""{key_prefix}.{key}""" not in model_state_dict or key is None:
																					print(f"""failed converting {original_key} to {key}, does not match"""				)

														# handle missmatched shape
														elif value.shape != model_state_dict[f"""{key_prefix}.{key}"""].shape:
																					snake_case_           =		model_state_dict[f"""{key_prefix}.{key}"""]
																					print(f"""{original_key}-> {key} : \nshape {val.shape} and { value.shape}, do not match"""				)
																					snake_case_           =		original_key

														snake_case_           =		original_key
														snake_case_           =		value

							return new_dict







@torch.no_grad()
def 						_a				(      _SCREAMING_SNAKE_CASE=None				, _SCREAMING_SNAKE_CASE=None				)    ->				Optional[int]:
							for file in MODEL_MAPPING[model_name]:
														if not os.path.isfile(f"""{pytorch_dump_folder_path}/{file.split("/"				)[-1]}"""				):
																					snake_case_           =		requests.get(f"""{PREFIX}{file}"""				, allow_redirects=_SCREAMING_SNAKE_CASE				)
																					os.makedirs(f"""{pytorch_dump_folder_path}/"""				, exist_ok=_SCREAMING_SNAKE_CASE				)
																					open(f"""{pytorch_dump_folder_path}/{file.split("/"				)[-1]}"""				, """wb"""				).write(r.content				)

							snake_case_           =		MODEL_MAPPING[model_name.split("""/"""				)[-1]]

							snake_case_           =		JukeboxConfig.from_pretrained(_SCREAMING_SNAKE_CASE				)
							snake_case_           =		JukeboxModel(_SCREAMING_SNAKE_CASE				)

							snake_case_           =		[]
							snake_case_           =		{}
							for i, dict_name in enumerate(_SCREAMING_SNAKE_CASE				):
														snake_case_           =		torch.load(f"""{pytorch_dump_folder_path}/{dict_name.split("/"				)[-1]}"""				)["""model"""]

														snake_case_           =		{}
														for k in old_dic.keys():
																					if k.endswith(""".b"""				):
																												snake_case_           =		old_dic[k]
																					elif k.endswith(""".w"""				):
																												snake_case_           =		old_dic[k]
																					elif "level_2" not in dict_name and "cond.model." in k:
																												snake_case_           =		old_dic[k]
																					else:
																												snake_case_           =		old_dic[k]

														snake_case_           =		"""vqvae""" if i == 0 else f"""priors.{3 - i}"""
														snake_case_           =		fix_jukebox_keys(_SCREAMING_SNAKE_CASE				, model.state_dict()				, _SCREAMING_SNAKE_CASE				, _SCREAMING_SNAKE_CASE				)
														weight_dict.append(_SCREAMING_SNAKE_CASE				)

							snake_case_           =		weight_dict.pop(0				)
							model.vqvae.load_state_dict(_SCREAMING_SNAKE_CASE				)
							for i in range(len(_SCREAMING_SNAKE_CASE				)				):
														model.priors[i].load_state_dict(weight_dict[2 - i]				)

							Path(_SCREAMING_SNAKE_CASE				).mkdir(exist_ok=_SCREAMING_SNAKE_CASE				)
							with open(f"""{pytorch_dump_folder_path}/mapping.json"""				, """w"""				) as txtfile:
														json.dump(_SCREAMING_SNAKE_CASE				, _SCREAMING_SNAKE_CASE				)

							print(f"""Saving model {model_name} to {pytorch_dump_folder_path}"""				)
							model.save_pretrained(_SCREAMING_SNAKE_CASE				)

							return weight_dict


if __name__ == "__main__":
				__SCREAMING_SNAKE_CASE							:    Union[str, Any]		=    argparse.ArgumentParser()
				# Required parameters
				parser.add_argument(
				    '--model_name',
				    default='jukebox-5b-lyrics',
				    type=str,
				    help='Name of the model you\'d like to convert.',
				)
				parser.add_argument(
				    '--pytorch_dump_folder_path',
				    default='jukebox-5b-lyrics-converted',
				    type=str,
				    help='Path to the output PyTorch model directory.',
				)
				__SCREAMING_SNAKE_CASE							:    str		=    parser.parse_args()
				convert_openai_checkpoint(args.model_name, args.pytorch_dump_folder_path)



