


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


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






def  __lowerCAmelCase (__lowerCAmelCase   ):
		if "emb" in name:
				_UpperCAmelCase		:       List[str]								=      name.replace("emb"      , "model.decoder.embed_tokens"   )
		if "transformer" in name:
				_UpperCAmelCase		:       List[str]								=      name.replace("transformer"      , "model.decoder"   )
		if "cross_attention" in name:
				_UpperCAmelCase		:       Union[str, Any]								=      name.replace("cross_attention"      , "encoder_attn"   )
		if "linear1" in name:
				_UpperCAmelCase		:       Tuple								=      name.replace("linear1"      , "fc1"   )
		if "linear2" in name:
				_UpperCAmelCase		:       Optional[Any]								=      name.replace("linear2"      , "fc2"   )
		if "norm1" in name:
				_UpperCAmelCase		:       List[Any]								=      name.replace("norm1"      , "self_attn_layer_norm"   )
		if "norm_cross" in name:
				_UpperCAmelCase		:       int								=      name.replace("norm_cross"      , "encoder_attn_layer_norm"   )
		if "norm2" in name:
				_UpperCAmelCase		:       Union[str, Any]								=      name.replace("norm2"      , "final_layer_norm"   )
		if "out_norm" in name:
				_UpperCAmelCase		:       Optional[int]								=      name.replace("out_norm"      , "model.decoder.layer_norm"   )
		if "linears" in name:
				_UpperCAmelCase		:       Any								=      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  __lowerCAmelCase (__lowerCAmelCase      , __lowerCAmelCase   ):
		_UpperCAmelCase		:       str								=      list(state_dict.keys()   )
		_UpperCAmelCase		:       Tuple								=      {}
		for key in keys:
				_UpperCAmelCase		:       List[Any]								=      state_dict.pop(__lowerCAmelCase   )
				_UpperCAmelCase		:       Tuple								=      rename_keys(__lowerCAmelCase   )
				if "in_proj_weight" in key:
						# split fused qkv proj
						_UpperCAmelCase		:       Any								=      val[:hidden_size, :]
						_UpperCAmelCase		:       Union[str, Any]								=      val[hidden_size : 2 * hidden_size, :]
						_UpperCAmelCase		:       Optional[int]								=      val[-hidden_size:, :]
				elif "enc_to_dec_proj" in key:
						_UpperCAmelCase		:       Optional[Any]								=      val
				else:
						_UpperCAmelCase		:       Dict								=      val
		return state_dict, enc_dec_proj_state_dict






def  __lowerCAmelCase (__lowerCAmelCase   ):
		if checkpoint == "small":
				# default config values
				_UpperCAmelCase		:       Any								=      1_024
				_UpperCAmelCase		:       Tuple								=      24
				_UpperCAmelCase		:       Union[str, Any]								=      16
		elif checkpoint == "medium":
				_UpperCAmelCase		:       List[Any]								=      1_536
				_UpperCAmelCase		:       str								=      48
				_UpperCAmelCase		:       str								=      24
		elif checkpoint == "large":
				_UpperCAmelCase		:       Dict								=      2_048
				_UpperCAmelCase		:       Optional[int]								=      48
				_UpperCAmelCase		:       Any								=      32
		else:
				raise ValueError(F"""Checkpoint should be one of `['small', 'medium', 'large']`, got {checkpoint}."""   )
		_UpperCAmelCase		:       Optional[int]								=      MusicgenDecoderConfig(
		    hidden_size=__lowerCAmelCase      , ffn_dim=hidden_size * 4      , num_hidden_layers=__lowerCAmelCase      , num_attention_heads=__lowerCAmelCase      , )
		return config






@torch.no_grad()
def  __lowerCAmelCase (__lowerCAmelCase      , __lowerCAmelCase=None      , __lowerCAmelCase=None      , __lowerCAmelCase="cpu"   ):
		_UpperCAmelCase		:       str								=      MusicGen.get_pretrained(__lowerCAmelCase      , device=__lowerCAmelCase   )
		_UpperCAmelCase		:       Union[str, Any]								=      decoder_config_from_checkpoint(__lowerCAmelCase   )

		_UpperCAmelCase		:       List[Any]								=      fairseq_model.lm.state_dict()
		_UpperCAmelCase					,    _UpperCAmelCase		:       Any								=      rename_state_dict(
		    __lowerCAmelCase      , hidden_size=decoder_config.hidden_size   )

		_UpperCAmelCase		:       Tuple								=      TaEncoderModel.from_pretrained("t5-base"   )
		_UpperCAmelCase		:       int								=      EncodecModel.from_pretrained("facebook/encodec_32khz"   )
		_UpperCAmelCase		:       Tuple								=      MusicgenForCausalLM(__lowerCAmelCase   ).eval()

		# load all decoder weights - expect that we'll be missing embeddings and enc-dec projection
		_UpperCAmelCase					,    _UpperCAmelCase		:       Any								=      decoder.load_state_dict(__lowerCAmelCase      , strict=__lowerCAmelCase   )

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

		if len(__lowerCAmelCase   ) > 0:
				raise ValueError(F"""Missing key(s) in state_dict: {missing_keys}"""   )

		if len(__lowerCAmelCase   ) > 0:
				raise ValueError(F"""Unexpected key(s) in state_dict: {unexpected_keys}"""   )

		# init the composite model
		_UpperCAmelCase		:       Tuple								=      MusicgenForConditionalGeneration(text_encoder=__lowerCAmelCase      , audio_encoder=__lowerCAmelCase      , decoder=__lowerCAmelCase   )

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

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

		with torch.no_grad():
				_UpperCAmelCase		:       Any								=      model(input_ids=__lowerCAmelCase      , decoder_input_ids=__lowerCAmelCase   ).logits

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

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

		_UpperCAmelCase		:       Optional[int]								=      MusicgenProcessor(feature_extractor=__lowerCAmelCase      , tokenizer=__lowerCAmelCase   )

		# set the appropriate bos/pad token ids
		_UpperCAmelCase		:       Dict								=      2_048
		_UpperCAmelCase		:       List[str]								=      2_048

		# set other default generation config params
		_UpperCAmelCase		:       Dict								=      int(30 * audio_encoder.config.frame_rate   )
		_UpperCAmelCase		:       Any								=      True
		_UpperCAmelCase		:       Dict								=      3.0

		if pytorch_dump_folder is not None:
				Path(__lowerCAmelCase   ).mkdir(exist_ok=__lowerCAmelCase   )
				logger.info(F"""Saving model {checkpoint} to {pytorch_dump_folder}"""   )
				model.save_pretrained(__lowerCAmelCase   )
				processor.save_pretrained(__lowerCAmelCase   )

		if repo_id:
				logger.info(F"""Pushing model {checkpoint} to {repo_id}"""   )
				model.push_to_hub(__lowerCAmelCase   )
				processor.push_to_hub(__lowerCAmelCase   )


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

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







'''simple docstring'''


import argparse
import gc
import json
import os
import re

import torch
from huggingface_hub import hf_hub_download

from transformers import AutoModelForCausalLM, AutoTokenizer, PreTrainedTokenizerFast, RwkvConfig
from transformers.modeling_utils import WEIGHTS_INDEX_NAME, shard_checkpoint


lowerCamelCase__    	=     {
    '169M': 12,
    '430M': 24,
    '1B5': 24,
    '3B': 32,
    '7B': 32,
    '14B': 40,
}

lowerCamelCase__    	=     {
    '169M': 768,
    '430M': 1_024,
    '1B5': 2_048,
    '3B': 2_560,
    '7B': 4_096,
    '14B': 5_120,
}






def  __lowerCAmelCase (__lowerCAmelCase   ):
		_UpperCAmelCase		:       List[str]								=      list(state_dict.keys()   )
		for name in state_dict_keys:
				_UpperCAmelCase		:       Optional[int]								=      state_dict.pop(__lowerCAmelCase   )
				# emb -> embedding
				if name.startswith("emb."   ):
						_UpperCAmelCase		:       Tuple								=      name.replace("emb."      , "embeddings."   )
				# ln_0 -> pre_ln (only present at block 0)
				if name.startswith("blocks.0.ln0"   ):
						_UpperCAmelCase		:       Optional[int]								=      name.replace("blocks.0.ln0"      , "blocks.0.pre_ln"   )
				# att -> attention
				_UpperCAmelCase		:       Union[str, Any]								=      re.sub(R"blocks\.(\d+)\.att"      , R"blocks.\1.attention"      , __lowerCAmelCase   )
				# ffn -> feed_forward
				_UpperCAmelCase		:       Dict								=      re.sub(R"blocks\.(\d+)\.ffn"      , R"blocks.\1.feed_forward"      , __lowerCAmelCase   )
				# time_mix_k -> time_mix_key and reshape
				if name.endswith(".time_mix_k"   ):
						_UpperCAmelCase		:       int								=      name.replace(".time_mix_k"      , ".time_mix_key"   )
				# time_mix_v -> time_mix_value and reshape
				if name.endswith(".time_mix_v"   ):
						_UpperCAmelCase		:       Union[str, Any]								=      name.replace(".time_mix_v"      , ".time_mix_value"   )
				# time_mix_r -> time_mix_key and reshape
				if name.endswith(".time_mix_r"   ):
						_UpperCAmelCase		:       int								=      name.replace(".time_mix_r"      , ".time_mix_receptance"   )

				if name != "head.weight":
						_UpperCAmelCase		:       List[str]								=      "rwkv." + name

				_UpperCAmelCase		:       Optional[Any]								=      weight
		return state_dict






def  __lowerCAmelCase (__lowerCAmelCase      , __lowerCAmelCase      , __lowerCAmelCase      , __lowerCAmelCase=None      , __lowerCAmelCase=None      , __lowerCAmelCase=False      , __lowerCAmelCase=None   ):
		# 1. If possible, build the tokenizer.
		if tokenizer_file is None:
				print("No `--tokenizer_file` provided, we will use the default tokenizer."   )
				_UpperCAmelCase		:       str								=      50_277
				_UpperCAmelCase		:       Optional[Any]								=      AutoTokenizer.from_pretrained("EleutherAI/gpt-neox-20b"   )
		else:
				_UpperCAmelCase		:       Tuple								=      PreTrainedTokenizerFast(tokenizer_file=__lowerCAmelCase   )
				_UpperCAmelCase		:       List[Any]								=      len(__lowerCAmelCase   )
		tokenizer.save_pretrained(__lowerCAmelCase   )

		# 2. Build the config
		_UpperCAmelCase		:       Optional[int]								=      list(NUM_HIDDEN_LAYERS_MAPPING.keys()   )
		if size is None:
				# Try to infer size from the checkpoint name
				for candidate in possible_sizes:
						if candidate in checkpoint_file:
								_UpperCAmelCase		:       Optional[Any]								=      candidate
								break
				if size is None:
						raise ValueError("Could not infer the size, please provide it with the `--size` argument."   )
		if size not in possible_sizes:
				raise ValueError(F"""`size` should be one of {possible_sizes}, got {size}."""   )

		_UpperCAmelCase		:       Any								=      RwkvConfig(
		    vocab_size=__lowerCAmelCase      , num_hidden_layers=NUM_HIDDEN_LAYERS_MAPPING[size]      , hidden_size=HIDEN_SIZE_MAPPING[size]      , )
		config.save_pretrained(__lowerCAmelCase   )

		# 3. Download model file then convert state_dict
		_UpperCAmelCase		:       str								=      hf_hub_download(__lowerCAmelCase      , __lowerCAmelCase   )
		_UpperCAmelCase		:       Optional[int]								=      torch.load(__lowerCAmelCase      , map_location="cpu"   )
		_UpperCAmelCase		:       Any								=      convert_state_dict(__lowerCAmelCase   )

		# 4. Split in shards and save
		_UpperCAmelCase					,    _UpperCAmelCase		:       List[str]								=      shard_checkpoint(__lowerCAmelCase   )
		for shard_file, shard in shards.items():
				torch.save(__lowerCAmelCase      , os.path.join(__lowerCAmelCase      , __lowerCAmelCase   )   )

		if index is not None:
				_UpperCAmelCase		:       int								=      os.path.join(__lowerCAmelCase      , __lowerCAmelCase   )
				# Save the index as well
				with open(__lowerCAmelCase      , "w"      , encoding="utf-8"   ) as f:
						_UpperCAmelCase		:       int								=      json.dumps(__lowerCAmelCase      , indent=2      , sort_keys=__lowerCAmelCase   ) + "\n"
						f.write(__lowerCAmelCase   )

				# 5. Clean up shards (for some reason the file PyTorch saves take the same space as the whole state_dict
				print(
				    "Cleaning up shards. This may error with an OOM error, it this is the case don't worry you still have converted the model."   )
				_UpperCAmelCase		:       Union[str, Any]								=      list(shards.keys()   )

				del state_dict
				del shards
				gc.collect()

				for shard_file in shard_files:
						_UpperCAmelCase		:       Union[str, Any]								=      torch.load(os.path.join(__lowerCAmelCase      , __lowerCAmelCase   )   )
						torch.save({k: v.cpu().clone() for k, v in state_dict.items()}      , os.path.join(__lowerCAmelCase      , __lowerCAmelCase   )   )

		del state_dict
		gc.collect()

		if push_to_hub:
				if model_name is None:
						raise ValueError("Please provide a `model_name` to push the model to the Hub."   )
				_UpperCAmelCase		:       Optional[Any]								=      AutoModelForCausalLM.from_pretrained(__lowerCAmelCase   )
				model.push_to_hub(__lowerCAmelCase      , max_shard_size="2GB"   )
				tokenizer.push_to_hub(__lowerCAmelCase   )


if __name__ == "__main__":
						lowerCamelCase__    	=     argparse.ArgumentParser()
						# Required parameters
						parser.add_argument(
						    '--repo_id', default=None, type=str, required=True, help='Repo ID from which to pull the checkpoint.'
						)
						parser.add_argument(
						    '--checkpoint_file', default=None, type=str, required=True, help='Name of the checkpoint file in the repo.'
						)
						parser.add_argument(
						    '--output_dir', default=None, type=str, required=True, help='Where to save the converted model.'
						)
						parser.add_argument(
						    '--tokenizer_file',
						    default=None,
						    type=str,
						    help='Path to the tokenizer file to use (if not provided, only the model is converted).',
						)
						parser.add_argument(
						    '--size',
						    default=None,
						    type=str,
						    help='Size of the model. Will be inferred from the `checkpoint_file` if not passed.',
						)
						parser.add_argument(
						    '--push_to_hub',
						    action='store_true',
						    help='Push to the Hub the converted model.',
						)
						parser.add_argument(
						    '--model_name',
						    default=None,
						    type=str,
						    help='Name of the pushed model on the Hub, including the username / organization.',
						)

						lowerCamelCase__    	=     parser.parse_args()
						convert_rmkv_checkpoint_to_hf_format(
						    args.repo_id,
						    args.checkpoint_file,
						    args.output_dir,
						    size=args.size,
						    tokenizer_file=args.tokenizer_file,
						    push_to_hub=args.push_to_hub,
						    model_name=args.model_name,
						)



