


'''simple docstring'''






import itertools
import os
import random
import tempfile
import unittest

import numpy as np

from transformers import TvltFeatureExtractor, is_datasets_available
from transformers.testing_utils import check_json_file_has_correct_format, require_torch, require_torchaudio
from transformers.utils.import_utils import is_torch_available

from ...test_sequence_feature_extraction_common import SequenceFeatureExtractionTestMixin


if is_torch_available():
		import torch

if is_datasets_available():
		from datasets import load_dataset

__UpperCAmelCase       =							random.Random()


def  _snake_case       (					A			,   A=1.0			,   A=None			,   A=None	)						->   List[Any]:
							if rng is None:
														lowerCAmelCase__	       =     global_rng

							lowerCAmelCase__	       =     []
							for batch_idx in range(shape[0]	):
														values.append([]	)
														for _ in range(shape[1]	):
																					values[-1].append(rng.random() * scale	)

							return values




class 		a__    (							unittest.TestCase  ):





							'''simple docstring'''






							def __init__(		self ,   lowerCamelCase_ ,   lowerCamelCase_=7 ,   lowerCamelCase_=4_00 ,   lowerCamelCase_=20_00 ,   lowerCamelCase_=20_48 ,   lowerCamelCase_=1_28 ,   lowerCamelCase_=1 ,   lowerCamelCase_=5_12 ,   lowerCamelCase_=30 ,   lowerCamelCase_=4_41_00 ,   )	->							Any:
														lowerCAmelCase__	       =     parent
														lowerCAmelCase__	       =     batch_size
														lowerCAmelCase__	       =     min_seq_length
														lowerCAmelCase__	       =     max_seq_length
														lowerCAmelCase__	       =     (self.max_seq_length - self.min_seq_length) // (self.batch_size - 1)
														lowerCAmelCase__	       =     spectrogram_length
														lowerCAmelCase__	       =     feature_size
														lowerCAmelCase__	       =     num_audio_channels
														lowerCAmelCase__	       =     hop_length
														lowerCAmelCase__	       =     chunk_length
														lowerCAmelCase__	       =     sampling_rate







							def   __SCREAMING_SNAKE_CASE							(		self       )	->							str:
														return {
														    "spectrogram_length": self.spectrogram_length,
														    "feature_size": self.feature_size,
														    "num_audio_channels": self.num_audio_channels,
														    "hop_length": self.hop_length,
														    "chunk_length": self.chunk_length,
														    "sampling_rate": self.sampling_rate,
														}







							def   __SCREAMING_SNAKE_CASE							(		self ,   lowerCamelCase_=False ,   lowerCamelCase_=False       )	->							Dict:
														def _flatten(lowerCamelCase_       ):
																					return list(itertools.chain(*lowerCamelCase_       )       )

														if equal_length:
																					lowerCAmelCase__	       =     [floats_list((self.max_seq_length, self.feature_size)       ) for _ in range(self.batch_size       )]
														else:
																					# make sure that inputs increase in size
																					lowerCAmelCase__	       =     [
																					    floats_list((x, self.feature_size)       )
																					    for x in range(self.min_seq_length ,   self.max_seq_length ,   self.seq_length_diff       )
																					]
														if numpify:
																					lowerCAmelCase__	       =     [np.asarray(lowerCamelCase_       ) for x in speech_inputs]
														return speech_inputs




@require_torch
@require_torchaudio
class 		a__    (							a__    ,    unittest.TestCase  ):





							'''simple docstring'''






							lowercase__      :       int     						=				TvltFeatureExtractor







							def   __SCREAMING_SNAKE_CASE							(		self       )	->							List[str]:
														lowerCAmelCase__	       =     TvltFeatureExtractionTester(self       )







							def   __SCREAMING_SNAKE_CASE							(		self       )	->							Optional[Any]:
														lowerCAmelCase__	       =     self.feature_extraction_class(**self.feat_extract_dict       )
														self.assertTrue(hasattr(lowerCamelCase_ ,   '''spectrogram_length'''       )       )
														self.assertTrue(hasattr(lowerCamelCase_ ,   '''feature_size'''       )       )
														self.assertTrue(hasattr(lowerCamelCase_ ,   '''num_audio_channels'''       )       )
														self.assertTrue(hasattr(lowerCamelCase_ ,   '''hop_length'''       )       )
														self.assertTrue(hasattr(lowerCamelCase_ ,   '''chunk_length'''       )       )
														self.assertTrue(hasattr(lowerCamelCase_ ,   '''sampling_rate'''       )       )







							def   __SCREAMING_SNAKE_CASE							(		self       )	->							int:
														lowerCAmelCase__	       =     self.feature_extraction_class(**self.feat_extract_dict       )

														with tempfile.TemporaryDirectory() as tmpdirname:
																					lowerCAmelCase__	       =     feat_extract_first.save_pretrained(lowerCamelCase_       )[0]
																					check_json_file_has_correct_format(lowerCamelCase_       )
																					lowerCAmelCase__	       =     self.feature_extraction_class.from_pretrained(lowerCamelCase_       )

														lowerCAmelCase__	       =     feat_extract_first.to_dict()
														lowerCAmelCase__	       =     feat_extract_second.to_dict()
														lowerCAmelCase__	       =     dict_first.pop('''mel_filters'''       )
														lowerCAmelCase__	       =     dict_second.pop('''mel_filters'''       )
														self.assertTrue(np.allclose(lowerCamelCase_ ,   lowerCamelCase_       )       )
														self.assertEqual(lowerCamelCase_ ,   lowerCamelCase_       )







							def   __SCREAMING_SNAKE_CASE							(		self       )	->							str:
														lowerCAmelCase__	       =     self.feature_extraction_class(**self.feat_extract_dict       )

														with tempfile.TemporaryDirectory() as tmpdirname:
																					lowerCAmelCase__	       =     os.path.join(lowerCamelCase_ ,   '''feat_extract.json'''       )
																					feat_extract_first.to_json_file(lowerCamelCase_       )
																					lowerCAmelCase__	       =     self.feature_extraction_class.from_json_file(lowerCamelCase_       )

														lowerCAmelCase__	       =     feat_extract_first.to_dict()
														lowerCAmelCase__	       =     feat_extract_second.to_dict()
														lowerCAmelCase__	       =     dict_first.pop('''mel_filters'''       )
														lowerCAmelCase__	       =     dict_second.pop('''mel_filters'''       )
														self.assertTrue(np.allclose(lowerCamelCase_ ,   lowerCamelCase_       )       )
														self.assertEqual(lowerCamelCase_ ,   lowerCamelCase_       )







							def   __SCREAMING_SNAKE_CASE							(		self       )	->							Optional[Any]:
														# Initialize feature_extractor
														lowerCAmelCase__	       =     self.feature_extraction_class(**self.feat_extract_dict       )

														# create three inputs of length 800, 1000, and 1200
														lowerCAmelCase__	       =     [floats_list((1, x)       )[0] for x in range(8_00 ,   14_00 ,   2_00       )]
														lowerCAmelCase__	       =     [np.asarray(lowerCamelCase_       ) for speech_input in speech_inputs]

														# Test not batched input
														lowerCAmelCase__	       =     feature_extractor(np_speech_inputs[0] ,   return_tensors='''np''' ,   sampling_rate=4_41_00       ).audio_values

														self.assertTrue(encoded_audios.ndim == 4       )
														self.assertTrue(encoded_audios.shape[-1] == feature_extractor.feature_size       )
														self.assertTrue(encoded_audios.shape[-2] <= feature_extractor.spectrogram_length       )
														self.assertTrue(encoded_audios.shape[-3] == feature_extractor.num_channels       )

														# Test batched
														lowerCAmelCase__	       =     feature_extractor(lowerCamelCase_ ,   return_tensors='''np''' ,   sampling_rate=4_41_00       ).audio_values

														self.assertTrue(encoded_audios.ndim == 4       )
														self.assertTrue(encoded_audios.shape[-1] == feature_extractor.feature_size       )
														self.assertTrue(encoded_audios.shape[-2] <= feature_extractor.spectrogram_length       )
														self.assertTrue(encoded_audios.shape[-3] == feature_extractor.num_channels       )

														# Test audio masking
														lowerCAmelCase__	       =     feature_extractor(
														    lowerCamelCase_ ,   return_tensors='''np''' ,   sampling_rate=4_41_00 ,   mask_audio=lowerCamelCase_       ).audio_values

														self.assertTrue(encoded_audios.ndim == 4       )
														self.assertTrue(encoded_audios.shape[-1] == feature_extractor.feature_size       )
														self.assertTrue(encoded_audios.shape[-2] <= feature_extractor.spectrogram_length       )
														self.assertTrue(encoded_audios.shape[-3] == feature_extractor.num_channels       )

														# Test 2-D numpy arrays are batched.
														lowerCAmelCase__	       =     [floats_list((1, x)       )[0] for x in (8_00, 8_00, 8_00)]
														lowerCAmelCase__	       =     np.asarray(lowerCamelCase_       )
														lowerCAmelCase__	       =     feature_extractor(lowerCamelCase_ ,   return_tensors='''np''' ,   sampling_rate=4_41_00       ).audio_values
														self.assertTrue(encoded_audios.ndim == 4       )
														self.assertTrue(encoded_audios.shape[-1] == feature_extractor.feature_size       )
														self.assertTrue(encoded_audios.shape[-2] <= feature_extractor.spectrogram_length       )
														self.assertTrue(encoded_audios.shape[-3] == feature_extractor.num_channels       )







							def   __SCREAMING_SNAKE_CASE							(		self ,   lowerCamelCase_       )	->							List[str]:
														lowerCAmelCase__	       =     load_dataset('''hf-internal-testing/librispeech_asr_dummy''' ,   '''clean''' ,   split='''validation'''       )
														# automatic decoding with librispeech
														lowerCAmelCase__	       =     ds.sort('''id'''       ).select(range(lowerCamelCase_       )       )[:num_samples]['''audio''']

														return [x["array"] for x in speech_samples]







							def   __SCREAMING_SNAKE_CASE							(		self       )	->							List[str]:
														lowerCAmelCase__	       =     self._load_datasamples(1       )
														lowerCAmelCase__	       =     TvltFeatureExtractor()
														lowerCAmelCase__	       =     feature_extractor(lowerCamelCase_ ,   return_tensors='''pt'''       ).audio_values

														self.assertEquals(audio_values.shape ,   (1, 1, 1_92, 1_28)       )

														lowerCAmelCase__	       =     torch.tensor([[-0.3_032, -0.2_708], [-0.4_434, -0.4_007]]       )
														self.assertTrue(torch.allclose(audio_values[0, 0, :2, :2] ,   lowerCamelCase_ ,   atol=1e-4       )       )



'''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()
__UpperCAmelCase       =							logging.get_logger(__name__)


__UpperCAmelCase       =							'''https://openaipublic.azureedge.net/jukebox/models/'''
__UpperCAmelCase       =							{
    '''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  _snake_case       (					A	)						->   Union[str, Any]:
							if key.endswith('''.model.1.bias'''	) and len(key.split('''.'''	)	) > 10:
														lowerCAmelCase__	       =     key.replace('''.model.1.bias'''			,   '''.conv1d_1.bias'''	)
							elif key.endswith('''.model.1.weight'''	) and len(key.split('''.'''	)	) > 10:
														lowerCAmelCase__	       =     key.replace('''.model.1.weight'''			,   '''.conv1d_1.weight'''	)
							elif key.endswith('''.model.3.bias'''	) and len(key.split('''.'''	)	) > 10:
														lowerCAmelCase__	       =     key.replace('''.model.3.bias'''			,   '''.conv1d_2.bias'''	)
							elif key.endswith('''.model.3.weight'''	) and len(key.split('''.'''	)	) > 10:
														lowerCAmelCase__	       =     key.replace('''.model.3.weight'''			,   '''.conv1d_2.weight'''	)

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

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

							if ".emb." in key and "total" not in key and "absolute" not in key and "relative" not in key:
														lowerCAmelCase__	       =     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:
														lowerCAmelCase__	       =     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  _snake_case       (					A			,   A			,   A			,   A	)						->   Optional[int]:
							lowerCAmelCase__	       =     {}
							import re

							lowerCAmelCase__	       =     re.compile(R'''encoders.(\d*).level_blocks.(\d*).model.(\d*).(\d).(bias|weight)'''	)
							lowerCAmelCase__	       =     re.compile(
							    R'''encoders.(\d*).level_blocks.(\d*).model.(\d*).(\d).model.(\d*).model.(\d*).(bias|weight)'''	)
							lowerCAmelCase__	       =     re.compile(R'''encoders.(\d*).level_blocks.(\d*).model.(\d*).(bias|weight)'''	)

							lowerCAmelCase__	       =     re.compile(R'''decoders.(\d*).level_blocks.(\d*).model.(\d*).(\d).(bias|weight)'''	)
							lowerCAmelCase__	       =     re.compile(
							    R'''decoders.(\d*).level_blocks.(\d*).model.(\d*).(\d).model.(\d*).model.(\d*).(bias|weight)'''	)
							lowerCAmelCase__	       =     re.compile(R'''decoders.(\d*).level_blocks.(\d*).model.(\d*).(bias|weight)'''	)

							lowerCAmelCase__	       =     re.compile(R'''conditioner_blocks.(\d*).cond.model.(\d*).(\d).(bias|weight)'''	)
							lowerCAmelCase__	       =     re.compile(
							    R'''conditioner_blocks.(\d*).cond.model.(\d*).(\d).model.(\d*).model.(\d*).(bias|weight)'''	)
							lowerCAmelCase__	       =     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(A	):
																					lowerCAmelCase__	       =     re_encoder_block_conv_in.match(A	)
																					lowerCAmelCase__	       =     regex_match.groups()
																					lowerCAmelCase__	       =     int(groups[2]	) * 2 + int(groups[3]	)
																					lowerCAmelCase__	       =     F"""encoders.{groups[0]}.level_blocks.{groups[1]}.downsample_block.{block_index}.{groups[-1]}"""
																					lowerCAmelCase__	       =     re_encoder_block_conv_in.sub(A			,   A	)

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

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

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

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

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

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

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

														elif re_prior_cond_proj_in.fullmatch(A	):
																					lowerCAmelCase__	       =     re_prior_cond_proj_in.match(A	)
																					lowerCAmelCase__	       =     regex_match.groups()
																					lowerCAmelCase__	       =     F"""conditioner_blocks.upsampler.proj_in.{groups[-1]}"""
																					lowerCAmelCase__	       =     re_prior_cond_proj_in.sub(A			,   A	)

														# keep original key
														else:
																					lowerCAmelCase__	       =     original_key

														lowerCAmelCase__	       =     replace_key(A	)

														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:
																					lowerCAmelCase__	       =     model_state_dict[F"""{key_prefix}.{key}"""]
																					print(F"""{original_key}-> {key} : \nshape {val.shape} and { value.shape}, do not match"""	)
																					lowerCAmelCase__	       =     original_key

														lowerCAmelCase__	       =     original_key
														lowerCAmelCase__	       =     value

							return new_dict


@torch.no_grad()
def  _snake_case       (					A=None			,   A=None	)						->   str:
							for file in MODEL_MAPPING[model_name]:
														if not os.path.isfile(F"""{pytorch_dump_folder_path}/{file.split("/"	)[-1]}"""	):
																					lowerCAmelCase__	       =     requests.get(F"""{PREFIX}{file}"""			,   allow_redirects=A	)
																					os.makedirs(F"""{pytorch_dump_folder_path}/"""			,   exist_ok=A	)
																					open(F"""{pytorch_dump_folder_path}/{file.split("/"	)[-1]}"""			,   '''wb'''	).write(r.content	)

							lowerCAmelCase__	       =     MODEL_MAPPING[model_name.split('''/'''	)[-1]]

							lowerCAmelCase__	       =     JukeboxConfig.from_pretrained(A	)
							lowerCAmelCase__	       =     JukeboxModel(A	)

							lowerCAmelCase__	       =     []
							lowerCAmelCase__	       =     {}
							for i, dict_name in enumerate(A	):
														lowerCAmelCase__	       =     torch.load(F"""{pytorch_dump_folder_path}/{dict_name.split("/"	)[-1]}"""	)['''model''']

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

														lowerCAmelCase__	       =     '''vqvae''' if i == 0 else F"""priors.{3 - i}"""
														lowerCAmelCase__	       =     fix_jukebox_keys(A			,   model.state_dict()			,   A			,   A	)
														weight_dict.append(A	)

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

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

							print(F"""Saving model {model_name} to {pytorch_dump_folder_path}"""	)
							model.save_pretrained(A	)

							return weight_dict


if __name__ == "__main__":
		__UpperCAmelCase       =							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.''',
		)
		__UpperCAmelCase       =							parser.parse_args()
		convert_openai_checkpoint(args.model_name, args.pytorch_dump_folder_path)