





import argparse
import json
import os
import re

import torch

from transformers import BloomConfig, BloomModel
from transformers.file_utils import CONFIG_NAME, WEIGHTS_NAME
from transformers.utils import logging


logging.set_verbosity_info()

SCREAMING_SNAKE_CASE_   :      Optional[Any]			=   [
    '''word_embeddings_layernorm.weight''',
    '''word_embeddings_layernorm.bias''',
    '''input_layernorm.weight''',
    '''input_layernorm.bias''',
    '''post_attention_layernorm.weight''',
    '''post_attention_layernorm.bias''',
    '''self_attention.dense.bias''',
    '''mlp.dense_4h_to_h.bias''',
    '''ln_f.weight''',
    '''ln_f.bias''',
]

SCREAMING_SNAKE_CASE_   :      List[Any]			=   [
    '''mlp.dense_4h_to_h.weight''',
    '''self_attention.dense.weight''',
]



def  SCREAMING_SNAKE_CASE				(      snake_case						,   snake_case							)		->    List[Any]:
		__lowercase											=     {
		    'word_embeddings.weight': 'word_embeddings.weight',
		    'word_embeddings.norm.weight': 'word_embeddings_layernorm.weight',
		    'word_embeddings.norm.bias': 'word_embeddings_layernorm.bias',
		    'weight': 'ln_f.weight',
		    'bias': 'ln_f.bias',
		}

		if key in layer_rename_map:
				return layer_rename_map[key]

		# Handle transformer blocks
		__lowercase											=     int(re.match(r'.*layer_(\d*).*'						,   snake_case							)[1]							)
		layer_number -= 3
		return F"h.{layer_number}." + key



def  SCREAMING_SNAKE_CASE				(      snake_case							)		->    Union[str, Any]:
		if dtype == torch.bool:
				return 1 / 8
		__lowercase											=     re.search(r'[^\d](\d+)$'						,   str(snake_case							)							)
		if bit_search is None:
				raise ValueError(F"`dtype` is not a valid dtype: {dtype}."							)
		__lowercase											=     int(bit_search.groups()[0]							)
		return bit_size // 8



def  SCREAMING_SNAKE_CASE				(      snake_case						,   snake_case						,   snake_case						,   snake_case						,   snake_case							)		->    Optional[int]:
		# Construct model
		if bloom_config_file == "":
				__lowercase											=     BloomConfig()
		else:
				__lowercase											=     BloomConfig.from_json_file(snake_case							)

		if shard_model:
				__lowercase											=     os.listdir(snake_case							)
				__lowercase											=     sorted(filter(lambda snake_case							: s.startswith('layer'							) and "model_00" in s						,   snake_case							)							)

				__lowercase											=     {'weight_map': {}, 'metadata': {}}
				__lowercase											=     0

				__lowercase											=     None

				__lowercase											=     BloomConfig()

				for j, file in enumerate(snake_case							):
						print('Processing file: {}'.format(snake_case							)							)
						__lowercase											=     None

						for i in range(snake_case							):
								# load all TP files
								__lowercase											=     file.replace('model_00'						,   F"model_0{i}"							)
								__lowercase											=     torch.load(os.path.join(snake_case						,   snake_case							)						,   map_location='cpu'							)

								# Rename keys in the transformers names
								__lowercase											=     list(temp.keys()							)
								for key in keys:
										__lowercase											=     temp.pop(snake_case							)

								if tensors is None:
										__lowercase											=     temp
								else:
										for key in tensors.keys():
												if any(key.endswith(snake_case							) for end in WEIGHTS_TO_AVERAGE_ENDSWITH							):
														# We average (sum and then divide) some weights accross TP ranks (see https://github.com/bigscience-workshop/Megatron-DeepSpeed/blob/olruwase/sync_layer_norms/megatron/training.py#L425)
														tensors[key] += temp[key]
												else:
														# Some weights are RowParallelLinear in Megatron-Deepspeed, others are ColumnParallel
														__lowercase											=     1 if any(text in key for text in WEIGHTS_WITH_ROW_PARALLELISM_CONTAIN							) else 0
														# We concatenate these weights accross TP ranks
														__lowercase											=     torch.cat([tensors[key], temp[key]]						,   dim=snake_case							)

            # Divide by the number of TP the weights we want to average
						for key in tensors.keys():
								if any(key.endswith(snake_case							) for end in WEIGHTS_TO_AVERAGE_ENDSWITH							):
										__lowercase											=     tensors[key] / pretraining_tp
						torch.save(
						    snake_case						,   os.path.join(
						        snake_case						,   'pytorch_model_{}-of-{}.bin'.format(str(j + 1							).zfill(5							)						,   str(len(snake_case							)							).zfill(5							)							)						,   )						,   )

						for key in tensors.keys():
								__lowercase											=     tensors[key]
								total_size += value.numel() * get_dtype_size(value.dtype							)
								if key not in index_dict["weight_map"]:
										__lowercase											=     'pytorch_model_{}-of-{}.bin'.format(
										    str(j + 1							).zfill(5							)						,   str(len(snake_case							)							).zfill(5							)							)

				__lowercase											=     BloomConfig()
				__lowercase											=     pytorch_dump_folder_path + '/' + CONFIG_NAME
				__lowercase											=     total_size
				with open(snake_case						,   'w'						,   encoding='utf-8'							) as f:
						f.write(config.to_json_string()							)
				with open(os.path.join(snake_case						,   WEIGHTS_NAME + '.index.json'							)						,   'w'						,   encoding='utf-8'							) as f:
						__lowercase											=     json.dumps(snake_case						,   indent=2						,   sort_keys=snake_case							) + '\n'
						f.write(snake_case							)
		else:
				__lowercase											=     BloomModel(snake_case							)

				__lowercase											=     os.listdir(snake_case							)
				__lowercase											=     sorted(filter(lambda snake_case							: s.startswith('layer'							) and "model_00" in s						,   snake_case							)							)

				__lowercase											=     None
				for i, file in enumerate(snake_case							):
						__lowercase											=     None
						for i in range(snake_case							):
								# load all TP files
								__lowercase											=     file.replace('model_00'						,   F"model_0{i}"							)
								__lowercase											=     torch.load(os.path.join(snake_case						,   snake_case							)						,   map_location='cpu'							)

								# Rename keys in the transformers names
								__lowercase											=     list(temp.keys()							)
								for key in keys:
										__lowercase											=     temp.pop(snake_case							)

								if tensors is None:
										__lowercase											=     temp
								else:
										for key in tensors.keys():
												# We average (sum and then divide) some weights accross TP ranks (see https://github.com/bigscience-workshop/Megatron-DeepSpeed/blob/olruwase/sync_layer_norms/megatron/training.py#L425)
												if any(key.endswith(snake_case							) for end in WEIGHTS_TO_AVERAGE_ENDSWITH							):
														tensors[key] += temp[key]
												else:
														# Some weights are RowParallelLinear in Megatron-Deepspeed, others are ColumnParallel
														__lowercase											=     1 if any(text in key for text in WEIGHTS_WITH_ROW_PARALLELISM_CONTAIN							) else 0
														# We concatenate these weights accross TP ranks
														__lowercase											=     torch.cat([tensors[key], temp[key]]						,   dim=snake_case							)

            # Divide by the number of TP the weights we want to average
						for key in tensors.keys():
								if any(key.endswith(snake_case							) for end in WEIGHTS_TO_AVERAGE_ENDSWITH							):
										__lowercase											=     tensors[key] / pretraining_tp

						__lowercase											=     model.load_state_dict(snake_case						,   strict=snake_case							)
						assert not other_keys.unexpected_keys, F"The keys {other_keys.unexpected_keys} are unexpected"
						if missing_keys is None:
								__lowercase											=     set(other_keys.missing_keys							)
						else:
								__lowercase											=     missing_keys.intersection(set(other_keys.missing_keys							)							)

				assert not missing_keys, F"The keys {missing_keys} are missing"

				# Save pytorch-model
				os.makedirs(snake_case						,   exist_ok=snake_case							)
				__lowercase											=     pytorch_dump_folder_path + '/' + WEIGHTS_NAME
				__lowercase											=     pytorch_dump_folder_path + '/' + CONFIG_NAME
				print(F"Save PyTorch model to {pytorch_weights_dump_path} with dtype {config.torch_dtype}"							)
				if config.torch_dtype is not None:
						__lowercase											=     model.to(config.torch_dtype							)
				torch.save(model.state_dict()						,   snake_case							)
				print(F"Save configuration file to {pytorch_config_dump_path}"							)
				with open(snake_case						,   'w'						,   encoding='utf-8'							) as f:
						f.write(config.to_json_string()							)


if __name__ == "__main__":
		SCREAMING_SNAKE_CASE_   :      int			=   argparse.ArgumentParser()
		# Required parameters
		parser.add_argument(
		    '''--bloom_checkpoint_path''',
		    default=None,
		    type=str,
		    required=True,
		    help='''Path to the Megatron-LM checkpoint path.''',
		)
		parser.add_argument(
		    '''--pytorch_dump_folder_path''', default=None, type=str, required=True, help='''Path to the output PyTorch model.'''
		)
		parser.add_argument(
		    '''--bloom_config_file''',
		    default='''''',
		    type=str,
		    help=(
		        '''An optional config json file corresponding to the pre-trained model. \n'''
		        '''This specifies the model architecture.'''
		    ),
		)
		parser.add_argument(
		    '''--shard_model''',
		    action='''store_true''',
		    help='''An optional setting to shard the output model \nThis enables sharding the converted checkpoint''',
		)
		parser.add_argument(
		    '''--pretraining_tp''',
		    default=4,
		    type=int,
		    help='''Pretraining TP rank that has been used when training the model in Megatron-LM \n''',
		)
		SCREAMING_SNAKE_CASE_   :      Tuple			=   parser.parse_args()
		convert_bloom_checkpoint_to_pytorch(
		    args.bloom_checkpoint_path,
		    args.bloom_config_file,
		    args.pytorch_dump_folder_path,
		    args.shard_model,
		    args.pretraining_tp,
		)











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()
SCREAMING_SNAKE_CASE_   :      Optional[Any]			=   logging.get_logger(__name__)


SCREAMING_SNAKE_CASE_   :      Any			=   ['''model.decoder.embed_positions.weights''']



def  SCREAMING_SNAKE_CASE				(      snake_case							)		->    int:
		if "emb" in name:
				__lowercase											=     name.replace('emb'						,   'model.decoder.embed_tokens'							)
		if "transformer" in name:
				__lowercase											=     name.replace('transformer'						,   'model.decoder'							)
		if "cross_attention" in name:
				__lowercase											=     name.replace('cross_attention'						,   'encoder_attn'							)
		if "linear1" in name:
				__lowercase											=     name.replace('linear1'						,   'fc1'							)
		if "linear2" in name:
				__lowercase											=     name.replace('linear2'						,   'fc2'							)
		if "norm1" in name:
				__lowercase											=     name.replace('norm1'						,   'self_attn_layer_norm'							)
		if "norm_cross" in name:
				__lowercase											=     name.replace('norm_cross'						,   'encoder_attn_layer_norm'							)
		if "norm2" in name:
				__lowercase											=     name.replace('norm2'						,   'final_layer_norm'							)
		if "out_norm" in name:
				__lowercase											=     name.replace('out_norm'						,   'model.decoder.layer_norm'							)
		if "linears" in name:
				__lowercase											=     name.replace('linears'						,   'lm_heads'							)
		if "condition_provider.conditioners.description.output_proj" in name:
				__lowercase											=     name.replace('condition_provider.conditioners.description.output_proj'						,   'enc_to_dec_proj'							)
		return name



def  SCREAMING_SNAKE_CASE				(      snake_case						,   snake_case							)		->    Tuple[Dict, Dict]:
		__lowercase											=     list(state_dict.keys()							)
		__lowercase											=     {}
		for key in keys:
				__lowercase											=     state_dict.pop(snake_case							)
				__lowercase											=     rename_keys(snake_case							)
				if "in_proj_weight" in key:
						# split fused qkv proj
						__lowercase											=     val[:hidden_size, :]
						__lowercase											=     val[hidden_size : 2 * hidden_size, :]
						__lowercase											=     val[-hidden_size:, :]
				elif "enc_to_dec_proj" in key:
						__lowercase											=     val
				else:
						__lowercase											=     val
		return state_dict, enc_dec_proj_state_dict



def  SCREAMING_SNAKE_CASE				(      snake_case							)		->    MusicgenDecoderConfig:
		if checkpoint == "small":
				# default config values
				__lowercase											=     1_024
				__lowercase											=     24
				__lowercase											=     16
		elif checkpoint == "medium":
				__lowercase											=     1_536
				__lowercase											=     48
				__lowercase											=     24
		elif checkpoint == "large":
				__lowercase											=     2_048
				__lowercase											=     48
				__lowercase											=     32
		else:
				raise ValueError(F"Checkpoint should be one of `['small', 'medium', 'large']`, got {checkpoint}."							)
		__lowercase											=     MusicgenDecoderConfig(
		    hidden_size=snake_case						,   ffn_dim=hidden_size * 4						,   num_hidden_layers=snake_case						,   num_attention_heads=snake_case						,   )
		return config



@torch.no_grad()
def  SCREAMING_SNAKE_CASE				(      snake_case						,   snake_case=None						,   snake_case=None						,   snake_case="cpu"							)		->    Optional[int]:
		__lowercase											=     MusicGen.get_pretrained(snake_case						,   device=snake_case							)
		__lowercase											=     decoder_config_from_checkpoint(snake_case							)

		__lowercase											=     fairseq_model.lm.state_dict()
		__lowercase						,      __lowercase											=     rename_state_dict(
		    snake_case						,   hidden_size=decoder_config.hidden_size							)

		__lowercase											=     TaEncoderModel.from_pretrained('t5-base'							)
		__lowercase											=     EncodecModel.from_pretrained('facebook/encodec_32khz'							)
		__lowercase											=     MusicgenForCausalLM(snake_case							).eval()

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

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

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

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

		# init the composite model
		__lowercase											=     MusicgenForConditionalGeneration(text_encoder=snake_case						,   audio_encoder=snake_case						,   decoder=snake_case							)

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

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

		with torch.no_grad():
				__lowercase											=     model(input_ids=snake_case						,   decoder_input_ids=snake_case							).logits

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

		# now construct the processor
		__lowercase											=     AutoTokenizer.from_pretrained('t5-base'							)
		__lowercase											=     AutoFeatureExtractor.from_pretrained('facebook/encodec_32khz'						,   padding_side='left'							)

		__lowercase											=     MusicgenProcessor(feature_extractor=snake_case						,   tokenizer=snake_case							)

		# set the appropriate bos/pad token ids
		__lowercase											=     2_048
		__lowercase											=     2_048

		# set other default generation config params
		__lowercase											=     int(30 * audio_encoder.config.frame_rate							)
		__lowercase											=     True
		__lowercase											=     3.0

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

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


if __name__ == "__main__":
		SCREAMING_SNAKE_CASE_   :      Optional[int]			=   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.'''
		)

		SCREAMING_SNAKE_CASE_   :      Any			=   parser.parse_args()
		convert_musicgen_checkpoint(args.checkpoint, args.pytorch_dump_folder, args.push_to_hub)




