





import argparse

import torch

from transformers import (
    EncodecConfig,
    EncodecFeatureExtractor,
    EncodecModel,
    logging,
)


# checkpoints downloaded from:
# https://dl.fbaipublicfiles.com/encodec/v0/encodec_24khz-d7cc33bc.th
# https://huggingface.co/facebook/musicgen-small/resolve/main/compression_state_dict.bin
# https://dl.fbaipublicfiles.com/encodec/v0/encodec_48khz-7e698e3e.th


logging.set_verbosity_info()
UpperCamelCase__   : Dict					  = logging.get_logger("""transformers.models.encodec""")

UpperCamelCase__   : Optional[int]					  = {
    """quantizer.vq.layers.*._codebook.inited""": """quantizer.layers.*.codebook.inited""",
    """quantizer.vq.layers.*._codebook.cluster_size""": """quantizer.layers.*.codebook.cluster_size""",
    """quantizer.vq.layers.*._codebook.embed""": """quantizer.layers.*.codebook.embed""",
    """quantizer.vq.layers.*._codebook.embed_avg""": """quantizer.layers.*.codebook.embed_avg""",
}
UpperCamelCase__   : List[Any]					  = {
    """encoder.model.0.conv.conv""": """encoder.layers.0.conv""",
    """encoder.model.1.block.1.conv.conv""": """encoder.layers.1.block.1.conv""",
    """encoder.model.1.block.3.conv.conv""": """encoder.layers.1.block.3.conv""",
    """encoder.model.1.shortcut.conv.conv""": """encoder.layers.1.shortcut.conv""",
    """encoder.model.3.conv.conv""": """encoder.layers.3.conv""",
    """encoder.model.4.block.1.conv.conv""": """encoder.layers.4.block.1.conv""",
    """encoder.model.4.block.3.conv.conv""": """encoder.layers.4.block.3.conv""",
    """encoder.model.4.shortcut.conv.conv""": """encoder.layers.4.shortcut.conv""",
    """encoder.model.6.conv.conv""": """encoder.layers.6.conv""",
    """encoder.model.7.block.1.conv.conv""": """encoder.layers.7.block.1.conv""",
    """encoder.model.7.block.3.conv.conv""": """encoder.layers.7.block.3.conv""",
    """encoder.model.7.shortcut.conv.conv""": """encoder.layers.7.shortcut.conv""",
    """encoder.model.9.conv.conv""": """encoder.layers.9.conv""",
    """encoder.model.10.block.1.conv.conv""": """encoder.layers.10.block.1.conv""",
    """encoder.model.10.block.3.conv.conv""": """encoder.layers.10.block.3.conv""",
    """encoder.model.10.shortcut.conv.conv""": """encoder.layers.10.shortcut.conv""",
    """encoder.model.12.conv.conv""": """encoder.layers.12.conv""",
    """encoder.model.13.lstm""": """encoder.layers.13.lstm""",
    """encoder.model.15.conv.conv""": """encoder.layers.15.conv""",
}
UpperCamelCase__   : Union[str, Any]					  = {
    """encoder.model.0.conv.norm""": """encoder.layers.0.norm""",
    """encoder.model.1.block.1.conv.norm""": """encoder.layers.1.block.1.norm""",
    """encoder.model.1.block.3.conv.norm""": """encoder.layers.1.block.3.norm""",
    """encoder.model.1.shortcut.conv.norm""": """encoder.layers.1.shortcut.norm""",
    """encoder.model.3.conv.norm""": """encoder.layers.3.norm""",
    """encoder.model.4.block.1.conv.norm""": """encoder.layers.4.block.1.norm""",
    """encoder.model.4.block.3.conv.norm""": """encoder.layers.4.block.3.norm""",
    """encoder.model.4.shortcut.conv.norm""": """encoder.layers.4.shortcut.norm""",
    """encoder.model.6.conv.norm""": """encoder.layers.6.norm""",
    """encoder.model.7.block.1.conv.norm""": """encoder.layers.7.block.1.norm""",
    """encoder.model.7.block.3.conv.norm""": """encoder.layers.7.block.3.norm""",
    """encoder.model.7.shortcut.conv.norm""": """encoder.layers.7.shortcut.norm""",
    """encoder.model.9.conv.norm""": """encoder.layers.9.norm""",
    """encoder.model.10.block.1.conv.norm""": """encoder.layers.10.block.1.norm""",
    """encoder.model.10.block.3.conv.norm""": """encoder.layers.10.block.3.norm""",
    """encoder.model.10.shortcut.conv.norm""": """encoder.layers.10.shortcut.norm""",
    """encoder.model.12.conv.norm""": """encoder.layers.12.norm""",
    """encoder.model.15.conv.norm""": """encoder.layers.15.norm""",
}
UpperCamelCase__   : int					  = {
    """decoder.model.0.conv.conv""": """decoder.layers.0.conv""",
    """decoder.model.1.lstm""": """decoder.layers.1.lstm""",
    """decoder.model.3.convtr.convtr""": """decoder.layers.3.conv""",
    """decoder.model.4.block.1.conv.conv""": """decoder.layers.4.block.1.conv""",
    """decoder.model.4.block.3.conv.conv""": """decoder.layers.4.block.3.conv""",
    """decoder.model.4.shortcut.conv.conv""": """decoder.layers.4.shortcut.conv""",
    """decoder.model.6.convtr.convtr""": """decoder.layers.6.conv""",
    """decoder.model.7.block.1.conv.conv""": """decoder.layers.7.block.1.conv""",
    """decoder.model.7.block.3.conv.conv""": """decoder.layers.7.block.3.conv""",
    """decoder.model.7.shortcut.conv.conv""": """decoder.layers.7.shortcut.conv""",
    """decoder.model.9.convtr.convtr""": """decoder.layers.9.conv""",
    """decoder.model.10.block.1.conv.conv""": """decoder.layers.10.block.1.conv""",
    """decoder.model.10.block.3.conv.conv""": """decoder.layers.10.block.3.conv""",
    """decoder.model.10.shortcut.conv.conv""": """decoder.layers.10.shortcut.conv""",
    """decoder.model.12.convtr.convtr""": """decoder.layers.12.conv""",
    """decoder.model.13.block.1.conv.conv""": """decoder.layers.13.block.1.conv""",
    """decoder.model.13.block.3.conv.conv""": """decoder.layers.13.block.3.conv""",
    """decoder.model.13.shortcut.conv.conv""": """decoder.layers.13.shortcut.conv""",
    """decoder.model.15.conv.conv""": """decoder.layers.15.conv""",
}
UpperCamelCase__   : Tuple					  = {
    """decoder.model.0.conv.norm""": """decoder.layers.0.norm""",
    """decoder.model.3.convtr.norm""": """decoder.layers.3.norm""",
    """decoder.model.4.block.1.conv.norm""": """decoder.layers.4.block.1.norm""",
    """decoder.model.4.block.3.conv.norm""": """decoder.layers.4.block.3.norm""",
    """decoder.model.4.shortcut.conv.norm""": """decoder.layers.4.shortcut.norm""",
    """decoder.model.6.convtr.norm""": """decoder.layers.6.norm""",
    """decoder.model.7.block.1.conv.norm""": """decoder.layers.7.block.1.norm""",
    """decoder.model.7.block.3.conv.norm""": """decoder.layers.7.block.3.norm""",
    """decoder.model.7.shortcut.conv.norm""": """decoder.layers.7.shortcut.norm""",
    """decoder.model.9.convtr.norm""": """decoder.layers.9.norm""",
    """decoder.model.10.block.1.conv.norm""": """decoder.layers.10.block.1.norm""",
    """decoder.model.10.block.3.conv.norm""": """decoder.layers.10.block.3.norm""",
    """decoder.model.10.shortcut.conv.norm""": """decoder.layers.10.shortcut.norm""",
    """decoder.model.12.convtr.norm""": """decoder.layers.12.norm""",
    """decoder.model.13.block.1.conv.norm""": """decoder.layers.13.block.1.norm""",
    """decoder.model.13.block.3.conv.norm""": """decoder.layers.13.block.3.norm""",
    """decoder.model.13.shortcut.conv.norm""": """decoder.layers.13.shortcut.norm""",
    """decoder.model.15.conv.norm""": """decoder.layers.15.norm""",
}
UpperCamelCase__   : Optional[Any]					  = {
    **MAPPING_QUANTIZER,
    **MAPPING_ENCODER,
    **MAPPING_DECODER,
}
UpperCamelCase__   : Union[str, Any]					  = {
    **MAPPING_QUANTIZER,
    **MAPPING_ENCODER,
    **MAPPING_ENCODER_48K,
    **MAPPING_DECODER,
    **MAPPING_DECODER_48K,
}
UpperCamelCase__   : List[str]					  = []
UpperCamelCase__   : List[Any]					  = []




def    A_(			A				,      A				,      A				,      A				,      A   ):
	for attribute in key.split("""."""   ):
		UpperCAmelCase_								=     getattr(A				,      A   )

	if weight_type is not None:
		UpperCAmelCase_								=     getattr(A				,      A   ).shape
	else:
		UpperCAmelCase_								=     hf_pointer.shape

	if hf_shape != value.shape:
		raise ValueError(
		    f"""Shape of hf {key + '.' + weight_type if weight_type is not None else ''} is {hf_shape}, but should be"""
		    f""" {value.shape} for {full_name}"""   )

	if weight_type == "weight":
		UpperCAmelCase_								=     value
	elif weight_type == "weight_g":
		UpperCAmelCase_								=     value
	elif weight_type == "weight_v":
		UpperCAmelCase_								=     value
	elif weight_type == "bias":
		UpperCAmelCase_								=     value
	elif weight_type == "running_mean":
		UpperCAmelCase_								=     value
	elif weight_type == "running_var":
		UpperCAmelCase_								=     value
	elif weight_type == "num_batches_tracked":
		UpperCAmelCase_								=     value
	elif weight_type == "weight_ih_l0":
		UpperCAmelCase_								=     value
	elif weight_type == "weight_hh_l0":
		UpperCAmelCase_								=     value
	elif weight_type == "bias_ih_l0":
		UpperCAmelCase_								=     value
	elif weight_type == "bias_hh_l0":
		UpperCAmelCase_								=     value
	elif weight_type == "weight_ih_l1":
		UpperCAmelCase_								=     value
	elif weight_type == "weight_hh_l1":
		UpperCAmelCase_								=     value
	elif weight_type == "bias_ih_l1":
		UpperCAmelCase_								=     value
	elif weight_type == "bias_hh_l1":
		UpperCAmelCase_								=     value
	else:
		UpperCAmelCase_								=     value

	logger.info(f"""{key + ('.' + weight_type if weight_type is not None else '')} was initialized from {full_name}."""   )




def    A_(			A				,      A   ):
	for key in ignore_keys:
		if key.endswith(""".*"""   ):
			if name.startswith(key[:-1]   ):
				return True
		elif ".*." in key:
			UpperCAmelCase_						,	UpperCAmelCase_								=     key.split(""".*."""   )
			if prefix in name and suffix in name:
				return True
		elif key in name:
			return True
	return False




def    A_(			A				,      A				,      A   ):
	UpperCAmelCase_								=     []

	if model_name == "encodec_24khz" or "encodec_32khz":
		UpperCAmelCase_								=     MAPPING_24K
	elif model_name == "encodec_48khz":
		UpperCAmelCase_								=     MAPPING_48K
	else:
		raise ValueError(f"""Unsupported model: {model_name}"""   )

	for name, value in orig_dict.items():
		if should_ignore(A				,      A   ):
			logger.info(f"""{name} was ignored"""   )
			continue

		UpperCAmelCase_								=     False
		for key, mapped_key in MAPPING.items():
			if "*" in key:
				UpperCAmelCase_						,	UpperCAmelCase_								=     key.split(""".*."""   )
				if prefix in name and suffix in name:
					UpperCAmelCase_								=     suffix

			if key in name:
				# HACK otherwise .embed gets initialized with .embed_avg too
				if key.endswith("""embed"""   ) and name.endswith("""embed_avg"""   ):
					continue

				UpperCAmelCase_								=     True
				if "*" in mapped_key:
					UpperCAmelCase_								=     name.split(A   )[0].split("""."""   )[-2]
					UpperCAmelCase_								=     mapped_key.replace("""*"""				,      A   )
				if "weight_g" in name:
					UpperCAmelCase_								=     """weight_g"""
				elif "weight_v" in name:
					UpperCAmelCase_								=     """weight_v"""
				elif "weight_ih_l0" in name:
					UpperCAmelCase_								=     """weight_ih_l0"""
				elif "weight_hh_l0" in name:
					UpperCAmelCase_								=     """weight_hh_l0"""
				elif "bias_ih_l0" in name:
					UpperCAmelCase_								=     """bias_ih_l0"""
				elif "bias_hh_l0" in name:
					UpperCAmelCase_								=     """bias_hh_l0"""
				elif "weight_ih_l1" in name:
					UpperCAmelCase_								=     """weight_ih_l1"""
				elif "weight_hh_l1" in name:
					UpperCAmelCase_								=     """weight_hh_l1"""
				elif "bias_ih_l1" in name:
					UpperCAmelCase_								=     """bias_ih_l1"""
				elif "bias_hh_l1" in name:
					UpperCAmelCase_								=     """bias_hh_l1"""
				elif "bias" in name:
					UpperCAmelCase_								=     """bias"""
				elif "weight" in name:
					UpperCAmelCase_								=     """weight"""
				elif "running_mean" in name:
					UpperCAmelCase_								=     """running_mean"""
				elif "running_var" in name:
					UpperCAmelCase_								=     """running_var"""
				elif "num_batches_tracked" in name:
					UpperCAmelCase_								=     """num_batches_tracked"""
				else:
					UpperCAmelCase_								=     None
				set_recursively(A				,      A				,      A				,      A				,      A   )
			continue
		if not is_used:
			unused_weights.append(A   )

	logger.warning(f"""Unused weights: {unused_weights}"""   )




@torch.no_grad()
def    A_(			A				,      A				,      A				,      A=None				,      A=None				,      ):
	if config_path is not None:
		UpperCAmelCase_								=     EncodecConfig.from_pretrained(A   )
	else:
		UpperCAmelCase_								=     EncodecConfig()

	if model_name == "encodec_24khz":
		pass  # config is already correct
	elif model_name == "encodec_32khz":
		UpperCAmelCase_								=     [8, 5, 4, 4]
		UpperCAmelCase_								=     [2.2]
		UpperCAmelCase_								=     64
		UpperCAmelCase_								=     32000
		UpperCAmelCase_								=     2048
		UpperCAmelCase_								=     False
		UpperCAmelCase_								=     False
		UpperCAmelCase_								=     False
	elif model_name == "encodec_48khz":
		UpperCAmelCase_								=     [8, 5, 4, 2]
		UpperCAmelCase_								=     [3.0, 6.0, 12.0, 24.0]
		UpperCAmelCase_								=     48000
		UpperCAmelCase_								=     2
		UpperCAmelCase_								=     False
		UpperCAmelCase_								=     """time_group_norm"""
		UpperCAmelCase_								=     True
		UpperCAmelCase_								=     1.0
		UpperCAmelCase_								=     0.01
	else:
		raise ValueError(f"""Unknown model name: {model_name}"""   )

	UpperCAmelCase_								=     EncodecModel(A   )

	UpperCAmelCase_								=     EncodecFeatureExtractor(
	    feature_size=config.audio_channels				,      sampling_rate=config.sampling_rate				,      chunk_length_s=config.chunk_length_s				,      overlap=config.overlap				,      )
	feature_extractor.save_pretrained(A   )

	UpperCAmelCase_								=     torch.load(A   )
	if "best_state" in original_checkpoint:
		# we might have a training state saved, in which case discard the yaml results and just retain the weights
		UpperCAmelCase_								=     original_checkpoint["""best_state"""]
	recursively_load_weights(A				,      A				,      A   )
	model.save_pretrained(A   )

	if repo_id:
		print("""Pushing to the hub..."""   )
		feature_extractor.push_to_hub(A   )
		model.push_to_hub(A   )


if __name__ == "__main__":
				UpperCamelCase__   : Dict					  = argparse.ArgumentParser()
				parser.add_argument(
				    """--model""",
				    default="""encodec_24khz""",
				    type=str,
				    help="""The model to convert. Should be one of 'encodec_24khz', 'encodec_32khz', 'encodec_48khz'.""",
				)
				parser.add_argument("""--checkpoint_path""", required=True, default=None, type=str, help="""Path to original checkpoint""")
				parser.add_argument("""--config_path""", default=None, type=str, help="""Path to hf config.json of model to convert""")
				parser.add_argument(
				    """--pytorch_dump_folder_path""", required=True, default=None, type=str, help="""Path to the output PyTorch model."""
				)
				parser.add_argument(
				    """--push_to_hub""", default=None, type=str, help="""Where to upload the converted model on the 🤗 hub."""
				)

				UpperCamelCase__   : Tuple					  = parser.parse_args()
				convert_checkpoint(
				    args.model,
				    args.checkpoint_path,
				    args.pytorch_dump_folder_path,
				    args.config_path,
				    args.push_to_hub,
				)









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


UpperCamelCase__   : int					  = {
    """169M""": 12,
    """430M""": 24,
    """1B5""": 24,
    """3B""": 32,
    """7B""": 32,
    """14B""": 40,
}

UpperCamelCase__   : Union[str, Any]					  = {
    """169M""": 7_68,
    """430M""": 10_24,
    """1B5""": 20_48,
    """3B""": 25_60,
    """7B""": 40_96,
    """14B""": 51_20,
}




def    A_(			A   ):
	UpperCAmelCase_								=     list(state_dict.keys()   )
	for name in state_dict_keys:
		UpperCAmelCase_								=     state_dict.pop(A   )
		# emb -> embedding
		if name.startswith("""emb."""   ):
			UpperCAmelCase_								=     name.replace("""emb."""				,      """embeddings."""   )
		# ln_0 -> pre_ln (only present at block 0)
		if name.startswith("""blocks.0.ln0"""   ):
			UpperCAmelCase_								=     name.replace("""blocks.0.ln0"""				,      """blocks.0.pre_ln"""   )
		# att -> attention
		UpperCAmelCase_								=     re.sub(R"""blocks\.(\d+)\.att"""				,      R"""blocks.\1.attention"""				,      A   )
		# ffn -> feed_forward
		UpperCAmelCase_								=     re.sub(R"""blocks\.(\d+)\.ffn"""				,      R"""blocks.\1.feed_forward"""				,      A   )
		# time_mix_k -> time_mix_key and reshape
		if name.endswith(""".time_mix_k"""   ):
			UpperCAmelCase_								=     name.replace(""".time_mix_k"""				,      """.time_mix_key"""   )
		# time_mix_v -> time_mix_value and reshape
		if name.endswith(""".time_mix_v"""   ):
			UpperCAmelCase_								=     name.replace(""".time_mix_v"""				,      """.time_mix_value"""   )
		# time_mix_r -> time_mix_key and reshape
		if name.endswith(""".time_mix_r"""   ):
			UpperCAmelCase_								=     name.replace(""".time_mix_r"""				,      """.time_mix_receptance"""   )

		if name != "head.weight":
			UpperCAmelCase_								=     """rwkv.""" + name

		UpperCAmelCase_								=     weight
	return state_dict




def    A_(			A				,      A				,      A				,      A=None				,      A=None				,      A=False				,      A=None   ):
	# 1. If possible, build the tokenizer.
	if tokenizer_file is None:
		print("""No `--tokenizer_file` provided, we will use the default tokenizer."""   )
		UpperCAmelCase_								=     50277
		UpperCAmelCase_								=     AutoTokenizer.from_pretrained("""EleutherAI/gpt-neox-20b"""   )
	else:
		UpperCAmelCase_								=     PreTrainedTokenizerFast(tokenizer_file=A   )
		UpperCAmelCase_								=     len(A   )
	tokenizer.save_pretrained(A   )

	# 2. Build the config
	UpperCAmelCase_								=     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_								=     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_								=     RwkvConfig(
	    vocab_size=A				,      num_hidden_layers=NUM_HIDDEN_LAYERS_MAPPING[size]				,      hidden_size=HIDEN_SIZE_MAPPING[size]				,      )
	config.save_pretrained(A   )

	# 3. Download model file then convert state_dict
	UpperCAmelCase_								=     hf_hub_download(A				,      A   )
	UpperCAmelCase_								=     torch.load(A				,      map_location="""cpu"""   )
	UpperCAmelCase_								=     convert_state_dict(A   )

	# 4. Split in shards and save
	UpperCAmelCase_						,	UpperCAmelCase_								=     shard_checkpoint(A   )
	for shard_file, shard in shards.items():
		torch.save(A				,      os.path.join(A				,      A   )   )

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

		# 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_								=     list(shards.keys()   )

		del state_dict
		del shards
		gc.collect()

		for shard_file in shard_files:
			UpperCAmelCase_								=     torch.load(os.path.join(A				,      A   )   )
			torch.save({k: v.cpu().clone() for k, v in state_dict.items()}				,      os.path.join(A				,      A   )   )

	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_								=     AutoModelForCausalLM.from_pretrained(A   )
		model.push_to_hub(A				,      max_shard_size="""2GB"""   )
		tokenizer.push_to_hub(A   )


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

				UpperCamelCase__   : List[str]					  = 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,
				)


