


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


__A	=['model.decoder.embed_positions.weights']







def     _UpperCamelCase      (		UpperCamelCase__					):
			if "emb" in name:
						UpperCAmelCase__  :					Union[str, Any]        =     name.replace("""emb""" ,      """model.decoder.embed_tokens"""					)
			if "transformer" in name:
						UpperCAmelCase__  :					Optional[int]        =     name.replace("""transformer""" ,      """model.decoder"""					)
			if "cross_attention" in name:
						UpperCAmelCase__  :					List[Any]        =     name.replace("""cross_attention""" ,      """encoder_attn"""					)
			if "linear1" in name:
						UpperCAmelCase__  :					Any        =     name.replace("""linear1""" ,      """fc1"""					)
			if "linear2" in name:
						UpperCAmelCase__  :					Dict        =     name.replace("""linear2""" ,      """fc2"""					)
			if "norm1" in name:
						UpperCAmelCase__  :					List[Any]        =     name.replace("""norm1""" ,      """self_attn_layer_norm"""					)
			if "norm_cross" in name:
						UpperCAmelCase__  :					Union[str, Any]        =     name.replace("""norm_cross""" ,      """encoder_attn_layer_norm"""					)
			if "norm2" in name:
						UpperCAmelCase__  :					int        =     name.replace("""norm2""" ,      """final_layer_norm"""					)
			if "out_norm" in name:
						UpperCAmelCase__  :					Dict        =     name.replace("""out_norm""" ,      """model.decoder.layer_norm"""					)
			if "linears" in name:
						UpperCAmelCase__  :					int        =     name.replace("""linears""" ,      """lm_heads"""					)
			if "condition_provider.conditioners.description.output_proj" in name:
						UpperCAmelCase__  :					Any        =     name.replace("""condition_provider.conditioners.description.output_proj""" ,      """enc_to_dec_proj"""					)
			return name







def     _UpperCamelCase      (		UpperCamelCase__ ,      UpperCamelCase__					):
			UpperCAmelCase__  :					Optional[Any]        =     list(state_dict.keys()					)
			UpperCAmelCase__  :					str        =     {}
			for key in keys:
						UpperCAmelCase__  :					str        =     state_dict.pop(UpperCamelCase__					)
						UpperCAmelCase__  :					List[Any]        =     rename_keys(UpperCamelCase__					)
						if "in_proj_weight" in key:
									# split fused qkv proj
									UpperCAmelCase__  :					int        =     val[:hidden_size, :]
									UpperCAmelCase__  :					Any        =     val[hidden_size : 2 * hidden_size, :]
									UpperCAmelCase__  :					Optional[int]        =     val[-hidden_size:, :]
						elif "enc_to_dec_proj" in key:
									UpperCAmelCase__  :					int        =     val
						else:
									UpperCAmelCase__  :					int        =     val
			return state_dict, enc_dec_proj_state_dict







def     _UpperCamelCase      (		UpperCamelCase__					):
			if checkpoint == "small":
						# default config values
						UpperCAmelCase__  :					Optional[Any]        =     1_0_2_4
						UpperCAmelCase__  :					List[str]        =     2_4
						UpperCAmelCase__  :					Any        =     1_6
			elif checkpoint == "medium":
						UpperCAmelCase__  :					Any        =     1_5_3_6
						UpperCAmelCase__  :					Tuple        =     4_8
						UpperCAmelCase__  :					str        =     2_4
			elif checkpoint == "large":
						UpperCAmelCase__  :					List[Any]        =     2_0_4_8
						UpperCAmelCase__  :					str        =     4_8
						UpperCAmelCase__  :					Optional[Any]        =     3_2
			else:
						raise ValueError(f'''Checkpoint should be one of `[\'small\', \'medium\', \'large\']`, got {checkpoint}.'''					)
			UpperCAmelCase__  :					Dict        =     MusicgenDecoderConfig(
			    hidden_size=UpperCamelCase__ ,      ffn_dim=hidden_size * 4 ,      num_hidden_layers=UpperCamelCase__ ,      num_attention_heads=UpperCamelCase__ ,      )
			return config







@torch.no_grad()
def     _UpperCamelCase      (		UpperCamelCase__ ,      UpperCamelCase__=None ,      UpperCamelCase__=None ,      UpperCamelCase__="cpu"					):
			UpperCAmelCase__  :					Any        =     MusicGen.get_pretrained(UpperCamelCase__ ,      device=UpperCamelCase__					)
			UpperCAmelCase__  :					List[Any]        =     decoder_config_from_checkpoint(UpperCamelCase__					)

			UpperCAmelCase__  :					Optional[Any]        =     fairseq_model.lm.state_dict()
			UpperCAmelCase__   ,       UpperCAmelCase__  :					Dict        =     rename_state_dict(
			    UpperCamelCase__ ,      hidden_size=decoder_config.hidden_size					)

			UpperCAmelCase__  :					Any        =     TaEncoderModel.from_pretrained("""t5-base"""					)
			UpperCAmelCase__  :					List[str]        =     EncodecModel.from_pretrained("""facebook/encodec_32khz"""					)
			UpperCAmelCase__  :					Tuple        =     MusicgenForCausalLM(UpperCamelCase__					).eval()

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

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

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

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

			# init the composite model
			UpperCAmelCase__  :					Dict        =     MusicgenForConditionalGeneration(text_encoder=UpperCamelCase__ ,      audio_encoder=UpperCamelCase__ ,      decoder=UpperCamelCase__					)

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

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

			with torch.no_grad():
						UpperCAmelCase__  :					Any        =     model(input_ids=UpperCamelCase__ ,      decoder_input_ids=UpperCamelCase__					).logits

			if logits.shape != (8, 1, 2_0_4_8):
						raise ValueError("""Incorrect shape for logits"""					)

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

			UpperCAmelCase__  :					List[Any]        =     MusicgenProcessor(feature_extractor=UpperCamelCase__ ,      tokenizer=UpperCamelCase__					)

			# set the appropriate bos/pad token ids
			UpperCAmelCase__  :					Union[str, Any]        =     2_0_4_8
			UpperCAmelCase__  :					Union[str, Any]        =     2_0_4_8

			# set other default generation config params
			UpperCAmelCase__  :					List[Any]        =     int(3_0 * audio_encoder.config.frame_rate					)
			UpperCAmelCase__  :					str        =     True
			UpperCAmelCase__  :					Union[str, Any]        =     3.0

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

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


if __name__ == "__main__":
		__A	=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.'
		)

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



'''simple docstring'''



from collections import Counter
from pathlib import Path
from typing import Optional, Tuple

import yaml


class 		_snake_case					(     yaml.SafeLoader	):




					def 				snake_case__						(      self							,							_lowerCamelCase):
								UpperCAmelCase__  :					Dict        =     [self.constructed_objects[key_node] for key_node, _ in node.value]
								UpperCAmelCase__  :					List[str]        =     [tuple(_lowerCamelCase) if isinstance(_lowerCamelCase							,							_lowerCamelCase) else key for key in keys]
								UpperCAmelCase__  :					List[str]        =     Counter(_lowerCamelCase)
								UpperCAmelCase__  :					Any        =     [key for key in counter if counter[key] > 1]
								if duplicate_keys:
											raise TypeError(f'''Got duplicate yaml keys: {duplicate_keys}''')




					def 				snake_case__						(      self							,							_lowerCamelCase							,							_lowerCamelCase=False):
								UpperCAmelCase__  :					List[str]        =     super().construct_mapping(_lowerCamelCase							,							deep=_lowerCamelCase)
								self._check_no_duplicates_on_constructed_node(_lowerCamelCase)
								return mapping







def     _UpperCamelCase      (		UpperCamelCase__					):
			UpperCAmelCase__  :					Tuple        =     list(readme_content.splitlines()					)
			if full_content and full_content[0] == "---" and "---" in full_content[1:]:
						UpperCAmelCase__  :					Dict        =     full_content[1:].index("""---"""					) + 1
						UpperCAmelCase__  :					List[Any]        =     """\n""".join(full_content[1:sep_idx]					)
						return yamlblock, "\n".join(full_content[sep_idx + 1 :]					)

			return None, "\n".join(UpperCamelCase__					)


class 		_snake_case					(     a__	):
					# class attributes
					lowerCAmelCase      :str   =					{'''train_eval_index'''}  # train-eval-index in the YAML metadata




					@classmethod
					def 				snake_case__						(      cls							,							_lowerCamelCase):
								with open(_lowerCamelCase							,							encoding="""utf-8""") as readme_file:
											UpperCAmelCase__   ,       UpperCAmelCase__  :					Any        =     _split_yaml_from_readme(readme_file.read())
								if yaml_string is not None:
											return cls.from_yaml_string(_lowerCamelCase)
								else:
											return cls()




					def 				snake_case__						(      self							,							_lowerCamelCase):
								if path.exists():
											with open(_lowerCamelCase							,							encoding="""utf-8""") as readme_file:
														UpperCAmelCase__  :					int        =     readme_file.read()
								else:
											UpperCAmelCase__  :					List[str]        =     None
								UpperCAmelCase__  :					Union[str, Any]        =     self._to_readme(_lowerCamelCase)
								with open(_lowerCamelCase							,							"""w"""							,							encoding="""utf-8""") as readme_file:
											readme_file.write(_lowerCamelCase)




					def 				snake_case__						(      self							,							_lowerCamelCase = None):
								if readme_content is not None:
											UpperCAmelCase__   ,       UpperCAmelCase__  :					Optional[Any]        =     _split_yaml_from_readme(_lowerCamelCase)
											UpperCAmelCase__  :					str        =     """---\n""" + self.to_yaml_string() + """---\n""" + content
								else:
											UpperCAmelCase__  :					Optional[int]        =     """---\n""" + self.to_yaml_string() + """---\n"""
								return full_content




					@classmethod
					def 				snake_case__						(      cls							,							_lowerCamelCase):
								UpperCAmelCase__  :					List[str]        =     yaml.load(_lowerCamelCase							,							Loader=_NoDuplicateSafeLoader) or {}

								# Convert the YAML keys to DatasetMetadata fields
								UpperCAmelCase__  :					List[str]        =     {
								    (key.replace("""-"""							,							"""_""") if key.replace("""-"""							,							"""_""") in cls._FIELDS_WITH_DASHES else key): value
								    for key, value in metadata_dict.items()
								}
								return cls(**_lowerCamelCase)




					def 				snake_case__						(      self):
								return yaml.safe_dump(
								    {
								        (key.replace("""_"""							,							"""-""") if key in self._FIELDS_WITH_DASHES else key): value
								        for key, value in self.items()
								    }							,							sort_keys=_lowerCamelCase							,							allow_unicode=_lowerCamelCase							,							encoding="""utf-8"""							,							).decode("""utf-8""")

__A	={
    'image-classification': [],
    'translation': [],
    'image-segmentation': [],
    'fill-mask': [],
    'automatic-speech-recognition': [],
    'token-classification': [],
    'sentence-similarity': [],
    'audio-classification': [],
    'question-answering': [],
    'summarization': [],
    'zero-shot-classification': [],
    'table-to-text': [],
    'feature-extraction': [],
    'other': [],
    'multiple-choice': [],
    'text-classification': [],
    'text-to-image': [],
    'text2text-generation': [],
    'zero-shot-image-classification': [],
    'tabular-classification': [],
    'tabular-regression': [],
    'image-to-image': [],
    'tabular-to-text': [],
    'unconditional-image-generation': [],
    'text-retrieval': [],
    'text-to-speech': [],
    'object-detection': [],
    'audio-to-audio': [],
    'text-generation': [],
    'conversational': [],
    'table-question-answering': [],
    'visual-question-answering': [],
    'image-to-text': [],
    'reinforcement-learning': [],
    'voice-activity-detection': [],
    'time-series-forecasting': [],
    'document-question-answering': [],
}


if __name__ == "__main__":
		from argparse import ArgumentParser

		__A	=ArgumentParser(usage='Validate the yaml metadata block of a README.md file.')
		ap.add_argument('readme_filepath')
		__A	=ap.parse_args()

		__A	=Path(args.readme_filepath)
		__A	=DatasetMetadata.from_readme(readme_filepath)
		print(dataset_metadata)
		dataset_metadata.to_readme(readme_filepath)