

'''simple docstring'''





import argparse
from pathlib import Path
from typing import Dict, OrderedDict, Tuple

import torch
from audiocraft.models import MusicGen

from transformers import (
    AutoFeatureExtractor,
    AutoTokenizer,
    EncodecModel,
    MusicgenDecoderConfig,
    MusicgenForConditionalGeneration,
    MusicgenProcessor,
    TaEncoderModel,
)
from transformers.models.musicgen.modeling_musicgen import MusicgenForCausalLM
from transformers.utils import logging


logging.set_verbosity_info()
UpperCamelCase_    						=       logging.get_logger(__name__)


UpperCamelCase_    						=       ["""model.decoder.embed_positions.weights"""]






def 				_lowerCAmelCase			(   __magic_name__						:     str						)				->  Optional[int]:
						if "emb" in name:
												lowercase :							List[str]    		=name.replace('''emb'''   , '''model.decoder.embed_tokens'''						)
						if "transformer" in name:
												lowercase :							List[str]    		=name.replace('''transformer'''   , '''model.decoder'''						)
						if "cross_attention" in name:
												lowercase :							List[str]    		=name.replace('''cross_attention'''   , '''encoder_attn'''						)
						if "linear1" in name:
												lowercase :							Optional[int]    		=name.replace('''linear1'''   , '''fc1'''						)
						if "linear2" in name:
												lowercase :							str    		=name.replace('''linear2'''   , '''fc2'''						)
						if "norm1" in name:
												lowercase :							Tuple    		=name.replace('''norm1'''   , '''self_attn_layer_norm'''						)
						if "norm_cross" in name:
												lowercase :							Tuple    		=name.replace('''norm_cross'''   , '''encoder_attn_layer_norm'''						)
						if "norm2" in name:
												lowercase :							Dict    		=name.replace('''norm2'''   , '''final_layer_norm'''						)
						if "out_norm" in name:
												lowercase :							Tuple    		=name.replace('''out_norm'''   , '''model.decoder.layer_norm'''						)
						if "linears" in name:
												lowercase :							Tuple    		=name.replace('''linears'''   , '''lm_heads'''						)
						if "condition_provider.conditioners.description.output_proj" in name:
												lowercase :							Tuple    		=name.replace('''condition_provider.conditioners.description.output_proj'''   , '''enc_to_dec_proj'''						)
						return name






def 				_lowerCAmelCase			(   __magic_name__						:     OrderedDict   , __magic_name__						:     int						)				->  Tuple[Dict, Dict]:
						lowercase :							Any    		=list(state_dict.keys()						)
						lowercase :							Dict    		={}
						for key in keys:
												lowercase :							Optional[int]    		=state_dict.pop(__magic_name__						)
												lowercase :							str    		=rename_keys(__magic_name__						)
												if "in_proj_weight" in key:
																		# split fused qkv proj
																		lowercase :							Optional[int]    		=val[:hidden_size, :]
																		lowercase :							List[Any]    		=val[hidden_size : 2 * hidden_size, :]
																		lowercase :							Any    		=val[-hidden_size:, :]
												elif "enc_to_dec_proj" in key:
																		lowercase :							Union[str, Any]    		=val
												else:
																		lowercase :							List[str]    		=val
						return state_dict, enc_dec_proj_state_dict






def 				_lowerCAmelCase			(   __magic_name__						:     str						)				->  MusicgenDecoderConfig:
						if checkpoint == "small":
												# default config values
												lowercase :							Tuple    		=1024
												lowercase :							List[Any]    		=24
												lowercase :							Any    		=16
						elif checkpoint == "medium":
												lowercase :							Dict    		=1536
												lowercase :							Union[str, Any]    		=48
												lowercase :							str    		=24
						elif checkpoint == "large":
												lowercase :							Optional[Any]    		=2048
												lowercase :							str    		=48
												lowercase :							Union[str, Any]    		=32
						else:
												raise ValueError(f'''Checkpoint should be one of `[\'small\', \'medium\', \'large\']`, got {checkpoint}.'''						)
						lowercase :							Dict    		=MusicgenDecoderConfig(
						    hidden_size=__magic_name__   , ffn_dim=hidden_size * 4   , num_hidden_layers=__magic_name__   , num_attention_heads=__magic_name__   , )
						return config






@torch.no_grad()
def 				_lowerCAmelCase			(   __magic_name__						:     List[str]   , __magic_name__						:     List[Any]=None   , __magic_name__						:     List[Any]=None   , __magic_name__						:     Optional[Any]="cpu"						)				->  List[Any]:
						lowercase :							Optional[int]    		=MusicGen.get_pretrained(__magic_name__   , device=__magic_name__						)
						lowercase :							Optional[int]    		=decoder_config_from_checkpoint(__magic_name__						)

						lowercase :							List[str]    		=fairseq_model.lm.state_dict()
						lowercase		,	lowercase :							List[str]    		=rename_state_dict(
						    __magic_name__   , hidden_size=decoder_config.hidden_size						)

						lowercase :							List[Any]    		=TaEncoderModel.from_pretrained('''t5-base'''						)
						lowercase :							str    		=EncodecModel.from_pretrained('''facebook/encodec_32khz'''						)
						lowercase :							Tuple    		=MusicgenForCausalLM(__magic_name__						).eval()

						# load all decoder weights - expect that we'll be missing embeddings and enc-dec projection
						lowercase		,	lowercase :							List[Any]    		=decoder.load_state_dict(__magic_name__   , strict=__magic_name__						)

						for key in missing_keys.copy():
												if key.startswith(('''text_encoder''', '''audio_encoder''')						) or key in EXPECTED_MISSING_KEYS:
																		missing_keys.remove(__magic_name__						)

						if len(__magic_name__						) > 0:
												raise ValueError(f'''Missing key(s) in state_dict: {missing_keys}'''						)

						if len(__magic_name__						) > 0:
												raise ValueError(f'''Unexpected key(s) in state_dict: {unexpected_keys}'''						)

						# init the composite model
						lowercase :							Dict    		=MusicgenForConditionalGeneration(text_encoder=__magic_name__   , audio_encoder=__magic_name__   , decoder=__magic_name__						)

						# load the pre-trained enc-dec projection (from the decoder state dict)
						model.enc_to_dec_proj.load_state_dict(__magic_name__						)

						# check we can do a forward pass
						lowercase :							str    		=torch.arange(0   , 8   , dtype=torch.long						).reshape(2   , -1						)
						lowercase :							Optional[Any]    		=input_ids.reshape(2 * 4   , -1						)

						with torch.no_grad():
												lowercase :							Tuple    		=model(input_ids=__magic_name__   , decoder_input_ids=__magic_name__						).logits

						if logits.shape != (8, 1, 2048):
												raise ValueError('''Incorrect shape for logits'''						)

						# now construct the processor
						lowercase :							List[Any]    		=AutoTokenizer.from_pretrained('''t5-base'''						)
						lowercase :							List[Any]    		=AutoFeatureExtractor.from_pretrained('''facebook/encodec_32khz'''   , padding_side='''left'''						)

						lowercase :							Tuple    		=MusicgenProcessor(feature_extractor=__magic_name__   , tokenizer=__magic_name__						)

						# set the appropriate bos/pad token ids
						lowercase :							int    		=2048
						lowercase :							str    		=2048

						# set other default generation config params
						lowercase :							List[str]    		=int(30 * audio_encoder.config.frame_rate						)
						lowercase :							List[str]    		=True
						lowercase :							int    		=3.0

						if pytorch_dump_folder is not None:
												Path(__magic_name__						).mkdir(exist_ok=__magic_name__						)
												logger.info(f'''Saving model {checkpoint} to {pytorch_dump_folder}'''						)
												model.save_pretrained(__magic_name__						)
												processor.save_pretrained(__magic_name__						)

						if repo_id:
												logger.info(f'''Pushing model {checkpoint} to {repo_id}'''						)
												model.push_to_hub(__magic_name__						)
												processor.push_to_hub(__magic_name__						)


if __name__ == "__main__":
							UpperCamelCase_    						=       argparse.ArgumentParser()
							# Required parameters
							parser.add_argument(
							    """--checkpoint""",
							    default="""small""",
							    type=str,
							    help="""Checkpoint size of the MusicGen model you'd like to convert. Can be one of: `['small', 'medium', 'large']`.""",
							)
							parser.add_argument(
							    """--pytorch_dump_folder""",
							    required=True,
							    default=None,
							    type=str,
							    help="""Path to the output PyTorch model directory.""",
							)
							parser.add_argument(
							    """--push_to_hub""", default=None, type=str, help="""Where to upload the converted model on the 🤗 hub."""
							)
							parser.add_argument(
							    """--device""", default="""cpu""", type=str, help="""Torch device to run the conversion, either cpu or cuda."""
							)

							UpperCamelCase_    						=       parser.parse_args()
							convert_musicgen_checkpoint(args.checkpoint, args.pytorch_dump_folder, args.push_to_hub)








'''simple docstring'''





from __future__ import annotations

import random
import unittest

from transformers import TransfoXLConfig, is_tf_available
from transformers.testing_utils import require_tf, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
							import tensorflow as tf

							from transformers import (
							    TF_TRANSFO_XL_PRETRAINED_MODEL_ARCHIVE_LIST,
							    TFTransfoXLForSequenceClassification,
							    TFTransfoXLLMHeadModel,
							    TFTransfoXLModel,
							)





class 		__SCREAMING_SNAKE_CASE							:







				def __init__(       self     :		List[str]		,					UpperCAmelCase__     :		Dict		,					):
										'''simple docstring'''
										lowercase :							Any    		=parent
										lowercase :							Optional[int]    		=13
										lowercase :							Union[str, Any]    		=7
										lowercase :							str    		=30
										lowercase :							Optional[int]    		=self.seq_length + self.mem_len
										lowercase :							Dict    		=15
										lowercase :							List[str]    		=True
										lowercase :							Optional[int]    		=True
										lowercase :							Tuple    		=99
										lowercase :							str    		=[10, 50, 80]
										lowercase :							List[Any]    		=32
										lowercase :							Optional[int]    		=32
										lowercase :							int    		=4
										lowercase :							Any    		=8
										lowercase :							List[Any]    		=128
										lowercase :							List[str]    		=2
										lowercase :							Tuple    		=2
										lowercase :							int    		=None
										lowercase :							Optional[int]    		=1
										lowercase :							int    		=0
										lowercase :							List[str]    		=3
										lowercase :							str    		=self.vocab_size - 1
										lowercase :							Tuple    		=0.01







				def  lowerCamelCase_	(       self     :		List[Any]   ):
										'''simple docstring'''
										lowercase :							str    		=ids_tensor([self.batch_size, self.seq_length]		,					self.vocab_size   )
										lowercase :							Tuple    		=ids_tensor([self.batch_size, self.seq_length]		,					self.vocab_size   )

										lowercase :							str    		=None
										if self.use_labels:
																lowercase :							Tuple    		=ids_tensor([self.batch_size, self.seq_length]		,					self.vocab_size   )

										lowercase :							Union[str, Any]    		=TransfoXLConfig(
										    vocab_size=self.vocab_size		,					mem_len=self.mem_len		,					clamp_len=self.clamp_len		,					cutoffs=self.cutoffs		,					d_model=self.hidden_size		,					d_embed=self.d_embed		,					n_head=self.num_attention_heads		,					d_head=self.d_head		,					d_inner=self.d_inner		,					div_val=self.div_val		,					n_layer=self.num_hidden_layers		,					eos_token_id=self.eos_token_id		,					pad_token_id=self.vocab_size - 1		,					init_range=self.init_range		,					num_labels=self.num_labels		,					)

										return (config, input_ids_a, input_ids_a, lm_labels)







				def  lowerCamelCase_	(       self     :		Union[str, Any]   ):
										'''simple docstring'''
										random.seed(self.seed   )
										tf.random.set_seed(self.seed   )







				def  lowerCamelCase_	(       self     :		Union[str, Any]		,					UpperCAmelCase__     :		Optional[Any]		,					UpperCAmelCase__     :		Dict		,					UpperCAmelCase__     :		List[str]		,					UpperCAmelCase__     :		List[str]   ):
										'''simple docstring'''
										lowercase :							Tuple    		=TFTransfoXLModel(UpperCAmelCase__   )

										lowercase		,	lowercase :							Optional[Any]    		=model(UpperCAmelCase__   ).to_tuple()

										lowercase :							List[str]    		={'''input_ids''': input_ids_a, '''mems''': mems_a}

										lowercase		,	lowercase :							Any    		=model(UpperCAmelCase__   ).to_tuple()

										self.parent.assertEqual(hidden_states_a.shape		,					(self.batch_size, self.seq_length, self.hidden_size)   )
										self.parent.assertEqual(hidden_states_a.shape		,					(self.batch_size, self.seq_length, self.hidden_size)   )
										self.parent.assertListEqual(
										    [mem.shape for mem in mems_a]		,					[(self.mem_len, self.batch_size, self.hidden_size)] * self.num_hidden_layers		,					)
										self.parent.assertListEqual(
										    [mem.shape for mem in mems_a]		,					[(self.mem_len, self.batch_size, self.hidden_size)] * self.num_hidden_layers		,					)







				def  lowerCamelCase_	(       self     :		int		,					UpperCAmelCase__     :		List[str]		,					UpperCAmelCase__     :		str		,					UpperCAmelCase__     :		Dict		,					UpperCAmelCase__     :		int   ):
										'''simple docstring'''
										lowercase :							int    		=TFTransfoXLLMHeadModel(UpperCAmelCase__   )

										lowercase		,	lowercase :							Tuple    		=model(UpperCAmelCase__   ).to_tuple()

										lowercase :							Optional[Any]    		={'''input_ids''': input_ids_a, '''labels''': lm_labels}
										lowercase		,	lowercase :							Optional[int]    		=model(UpperCAmelCase__   ).to_tuple()

										lowercase		,	lowercase :							List[str]    		=model([input_ids_a, mems_a]   ).to_tuple()

										lowercase :							int    		={'''input_ids''': input_ids_a, '''mems''': mems_a, '''labels''': lm_labels}

										lowercase		,	lowercase :							str    		=model(UpperCAmelCase__   ).to_tuple()

										self.parent.assertEqual(lm_logits_a.shape		,					(self.batch_size, self.seq_length, self.vocab_size)   )
										self.parent.assertListEqual(
										    [mem.shape for mem in mems_a]		,					[(self.mem_len, self.batch_size, self.hidden_size)] * self.num_hidden_layers		,					)

										self.parent.assertEqual(lm_logits_a.shape		,					(self.batch_size, self.seq_length, self.vocab_size)   )
										self.parent.assertListEqual(
										    [mem.shape for mem in mems_a]		,					[(self.mem_len, self.batch_size, self.hidden_size)] * self.num_hidden_layers		,					)







				def  lowerCamelCase_	(       self     :		str		,					UpperCAmelCase__     :		int		,					UpperCAmelCase__     :		Any		,					UpperCAmelCase__     :		List[str]		,					UpperCAmelCase__     :		Union[str, Any]   ):
										'''simple docstring'''
										lowercase :							Optional[int]    		=TFTransfoXLForSequenceClassification(UpperCAmelCase__   )
										lowercase :							Union[str, Any]    		=model(UpperCAmelCase__   )
										self.parent.assertEqual(result.logits.shape		,					(self.batch_size, self.num_labels)   )







				def  lowerCamelCase_	(       self     :		List[str]   ):
										'''simple docstring'''
										lowercase :							List[Any]    		=self.prepare_config_and_inputs()
										((lowercase)		,	(lowercase)		,	(lowercase)		,	(lowercase)) :							Optional[Any]    		=config_and_inputs
										lowercase :							Union[str, Any]    		={'''input_ids''': input_ids_a}
										return config, inputs_dict





@require_tf
class 		__SCREAMING_SNAKE_CASE							(			lowercase__						,							lowercase__						,							unittest.TestCase      ):
				lowerCamelCase_		       =    (
				    (TFTransfoXLModel, TFTransfoXLLMHeadModel, TFTransfoXLForSequenceClassification) if is_tf_available() else ()
				)
				lowerCamelCase_		       =    () if is_tf_available() else ()
				lowerCamelCase_		       =    (
				    {
				        'feature-extraction': TFTransfoXLModel,
				        'text-classification': TFTransfoXLForSequenceClassification,
				        'text-generation': TFTransfoXLLMHeadModel,
				        'zero-shot': TFTransfoXLForSequenceClassification,
				    }
				    if is_tf_available()
				    else {}
				)
				# TODO: add this test when TFTransfoXLLMHead has a linear output layer implemented
				lowerCamelCase_		       =    False
				lowerCamelCase_		       =    False
				lowerCamelCase_		       =    False
				lowerCamelCase_		       =    False







				def  lowerCamelCase_	(       self     :		Dict		,					UpperCAmelCase__     :		List[Any]		,					UpperCAmelCase__     :		Optional[int]		,					UpperCAmelCase__     :		int		,					UpperCAmelCase__     :		Optional[int]		,					UpperCAmelCase__     :		int   ):
										'''simple docstring'''
										if pipeline_test_casse_name == "TextGenerationPipelineTests":
																# Get `ValueError: AttributeError: 'NoneType' object has no attribute 'new_ones'` or `AssertionError`.
																# `TransfoXLConfig` was never used in pipeline tests: cannot create a simple
																# tokenizer.
																return True

										return False







				def  lowerCamelCase_	(       self     :		int   ):
										'''simple docstring'''
										lowercase :							Any    		=TFTransfoXLModelTester(self   )
										lowercase :							Optional[int]    		=ConfigTester(self		,					config_class=UpperCAmelCase__		,					d_embed=37   )







				def  lowerCamelCase_	(       self     :		Any   ):
										'''simple docstring'''
										self.config_tester.run_common_tests()







				def  lowerCamelCase_	(       self     :		Optional[int]   ):
										'''simple docstring'''
										self.model_tester.set_seed()
										lowercase :							List[str]    		=self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_transfo_xl_model(*UpperCAmelCase__   )







				def  lowerCamelCase_	(       self     :		List[str]   ):
										'''simple docstring'''
										self.model_tester.set_seed()
										lowercase :							Any    		=self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_transfo_xl_lm_head(*UpperCAmelCase__   )







				def  lowerCamelCase_	(       self     :		List[Any]   ):
										'''simple docstring'''
										lowercase :							Dict    		=self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_transfo_xl_for_sequence_classification(*UpperCAmelCase__   )







				def  lowerCamelCase_	(       self     :		Optional[Any]   ):
										'''simple docstring'''
										lowercase		,	lowercase :							Tuple    		=self.model_tester.prepare_config_and_inputs_for_common()
										lowercase :							int    		=[TFTransfoXLForSequenceClassification]

										for model_class in self.all_model_classes:
																lowercase :							str    		=model_class(UpperCAmelCase__   )
																assert isinstance(model.get_input_embeddings()		,					tf.keras.layers.Layer   )
																if model_class in list_other_models_with_output_ebd:
																						lowercase :							Union[str, Any]    		=model.get_output_embeddings()
																						assert isinstance(UpperCAmelCase__		,					tf.keras.layers.Layer   )
																						lowercase :							Any    		=model.get_bias()
																						assert name is None
																else:
																						lowercase :							Optional[int]    		=model.get_output_embeddings()
																						assert x is None
																						lowercase :							Optional[int]    		=model.get_bias()
																						assert name is None







				def  lowerCamelCase_	(       self     :		Any   ):
										'''simple docstring'''
										# TODO JP: Make TransfoXL XLA compliant
										pass







				@slow
				def  lowerCamelCase_	(       self     :		Union[str, Any]   ):
										'''simple docstring'''
										for model_name in TF_TRANSFO_XL_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																lowercase :							int    		=TFTransfoXLModel.from_pretrained(UpperCAmelCase__   )
																self.assertIsNotNone(UpperCAmelCase__   )







				@unittest.skip(reason='''This model doesn\'t play well with fit() due to not returning a single loss.'''   )
				def  lowerCamelCase_	(       self     :		int   ):
										'''simple docstring'''
										pass





@require_tf
class 		__SCREAMING_SNAKE_CASE							(			unittest.TestCase      ):







				@unittest.skip('''Skip test until #12651 is resolved.'''   )
				@slow
				def  lowerCamelCase_	(       self     :		Tuple   ):
										'''simple docstring'''
										lowercase :							Optional[Any]    		=TFTransfoXLLMHeadModel.from_pretrained('''transfo-xl-wt103'''   )
										# fmt: off
										lowercase :							Tuple    		=tf.convert_to_tensor([[33,1297,2,1,1009,4,1109,11739,4762,358,5,25,245,22,1706,17,20098,5,3215,21,37,1110,3,13,1041,4,24,603,490,2,71477,20098,104447,2,20961,1,2604,4,1,329,3,6224,831,16002,2,8,603,78967,29546,23,803,20,25,416,5,8,232,4,277,6,1855,4601,3,29546,54,8,3609,5,57211,49,4,1,277,18,8,1755,15691,3,341,25,416,693,42573,71,17,401,94,31,17919,2,29546,7873,18,1,435,23,11011,755,5,5167,3,7983,98,84,2,29546,3267,8,3609,4,1,4865,1075,2,6087,71,6,346,8,5854,3,29546,824,1400,1868,2,19,160,2,311,8,5496,2,20920,17,25,15097,3,24,24,0]]		,					dtype=tf.intaa   )  # noqa: E231
										# fmt: on
										#  In 1991 , the remains of Russian Tsar Nicholas II and his family
										#  ( except for Alexei and Maria ) are discovered .
										#  The voice of Nicholas's young son , Tsarevich Alexei Nikolaevich , narrates the
										#  remainder of the story . 1883 Western Siberia ,
										#  a young Grigori Rasputin is asked by his father and a group of men to perform magic .
										#  Rasputin has a vision and denounces one of the men as a horse thief . Although his
										#  father initially slaps him for making such an accusation , Rasputin watches as the
										#  man is chased outside and beaten . Twenty years later , Rasputin sees a vision of
										#  the Virgin Mary , prompting him to become a priest . Rasputin quickly becomes famous ,
										#  with people , even a bishop , begging for his blessing . <eod> </s> <eos>

										# fmt: off
										lowercase :							Optional[int]    		=[33,1297,2,1,1009,4,1109,11739,4762,358,5,25,245,22,1706,17,20098,5,3215,21,37,1110,3,13,1041,4,24,603,490,2,71477,20098,104447,2,20961,1,2604,4,1,329,3,6224,831,16002,2,8,603,78967,29546,23,803,20,25,416,5,8,232,4,277,6,1855,4601,3,29546,54,8,3609,5,57211,49,4,1,277,18,8,1755,15691,3,341,25,416,693,42573,71,17,401,94,31,17919,2,29546,7873,18,1,435,23,11011,755,5,5167,3,7983,98,84,2,29546,3267,8,3609,4,1,4865,1075,2,6087,71,6,346,8,5854,3,29546,824,1400,1868,2,19,160,2,311,8,5496,2,20920,17,25,15097,3,24,24,0,33,1,1857,2,1,1009,4,1109,11739,4762,358,5,25,245,28,1110,3,13,1041,4,24,603,490,2,71477,20098,104447,2,20961,1,2604,4,1,329,3,0]  # noqa: E231
										# fmt: on
										#  In 1991, the remains of Russian Tsar Nicholas II and his family (
										#  except for Alexei and Maria ) are discovered. The voice of young son,
										#  Tsarevich Alexei Nikolaevich, narrates the remainder of the story.
										#  1883 Western Siberia, a young Grigori Rasputin is asked by his father
										#  and a group of men to perform magic. Rasputin has a vision and
										#  denounces one of the men as a horse thief. Although his father initially
										#  slaps him for making such an accusation, Rasputin watches as the man
										#  is chased outside and beaten. Twenty years later, Rasputin sees a vision
										#  of the Virgin Mary, prompting him to become a priest.
										#  Rasputin quickly becomes famous, with people, even a bishop, begging for
										#  his blessing. <unk> <unk> <eos> In the 1990s, the remains of Russian Tsar
										# Nicholas II and his family were discovered. The voice of <unk> young son,
										# Tsarevich Alexei Nikolaevich, narrates the remainder of the story.<eos>

										lowercase :							int    		=model.generate(UpperCAmelCase__		,					max_length=200		,					do_sample=UpperCAmelCase__   )
										self.assertListEqual(output_ids[0].numpy().tolist()		,					UpperCAmelCase__   )





