


'''simple docstring'''



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

__A	={
    """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""",
}
__A	={
    """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""",
}
__A	={
    """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""",
}
__A	={
    """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""",
}
__A	={
    """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""",
}
__A	={
    **MAPPING_QUANTIZER,
    **MAPPING_ENCODER,
    **MAPPING_DECODER,
}
__A	={
    **MAPPING_QUANTIZER,
    **MAPPING_ENCODER,
    **MAPPING_ENCODER_48K,
    **MAPPING_DECODER,
    **MAPPING_DECODER_48K,
}
__A	=[]
__A	=[]







def     _UpperCamelCase      (		UpperCamelCase__ ,      UpperCamelCase__ ,      UpperCamelCase__ ,      UpperCamelCase__ ,      UpperCamelCase__					):

			'''simple docstring'''





			for attribute in key.split("""."""					):
						UpperCAmelCase__  :					int        =     getattr(UpperCamelCase__ ,      UpperCamelCase__					)

			if weight_type is not None:
						UpperCAmelCase__  :					Any        =     getattr(UpperCamelCase__ ,      UpperCamelCase__					).shape
			else:
						UpperCAmelCase__  :					List[str]        =     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__  :					Any        =     value
			elif weight_type == "weight_g":
						UpperCAmelCase__  :					str        =     value
			elif weight_type == "weight_v":
						UpperCAmelCase__  :					List[Any]        =     value
			elif weight_type == "bias":
						UpperCAmelCase__  :					int        =     value
			elif weight_type == "running_mean":
						UpperCAmelCase__  :					Any        =     value
			elif weight_type == "running_var":
						UpperCAmelCase__  :					str        =     value
			elif weight_type == "num_batches_tracked":
						UpperCAmelCase__  :					int        =     value
			elif weight_type == "weight_ih_l0":
						UpperCAmelCase__  :					int        =     value
			elif weight_type == "weight_hh_l0":
						UpperCAmelCase__  :					Optional[Any]        =     value
			elif weight_type == "bias_ih_l0":
						UpperCAmelCase__  :					Optional[Any]        =     value
			elif weight_type == "bias_hh_l0":
						UpperCAmelCase__  :					List[str]        =     value
			elif weight_type == "weight_ih_l1":
						UpperCAmelCase__  :					Tuple        =     value
			elif weight_type == "weight_hh_l1":
						UpperCAmelCase__  :					Tuple        =     value
			elif weight_type == "bias_ih_l1":
						UpperCAmelCase__  :					Optional[Any]        =     value
			elif weight_type == "bias_hh_l1":
						UpperCAmelCase__  :					Dict        =     value
			else:
						UpperCAmelCase__  :					str        =     value

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







def     _UpperCamelCase      (		UpperCamelCase__ ,      UpperCamelCase__					):

			'''simple docstring'''





			for key in ignore_keys:
						if key.endswith(""".*"""					):
									if name.startswith(key[:-1]					):
												return True
						elif ".*." in key:
									UpperCAmelCase__  :					List[str]        =     key.split(""".*."""					)
									if prefix in name and suffix in name:
												return True
						elif key in name:
									return True
			return False







def     _UpperCamelCase      (		UpperCamelCase__ ,      UpperCamelCase__ ,      UpperCamelCase__					):

			'''simple docstring'''





			UpperCAmelCase__  :					int        =     []

			if model_name == "encodec_24khz" or "encodec_32khz":
						UpperCAmelCase__  :					Optional[Any]        =     MAPPING_24K
			elif model_name == "encodec_48khz":
						UpperCAmelCase__  :					str        =     MAPPING_48K
			else:
						raise ValueError(f'''Unsupported model: {model_name}'''					)

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

						UpperCAmelCase__  :					Tuple        =     False
						for key, mapped_key in MAPPING.items():
									if "*" in key:
												UpperCAmelCase__  :					List[str]        =     key.split(""".*."""					)
												if prefix in name and suffix in name:
															UpperCAmelCase__  :					str        =     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__  :					Optional[int]        =     True
												if "*" in mapped_key:
															UpperCAmelCase__  :					int        =     name.split(UpperCamelCase__					)[0].split("""."""					)[-2]
															UpperCAmelCase__  :					Optional[int]        =     mapped_key.replace("""*""" ,      UpperCamelCase__					)
												if "weight_g" in name:
															UpperCAmelCase__  :					Optional[int]        =     "weight_g"
												elif "weight_v" in name:
															UpperCAmelCase__  :					Any        =     "weight_v"
												elif "weight_ih_l0" in name:
															UpperCAmelCase__  :					Dict        =     "weight_ih_l0"
												elif "weight_hh_l0" in name:
															UpperCAmelCase__  :					Dict        =     "weight_hh_l0"
												elif "bias_ih_l0" in name:
															UpperCAmelCase__  :					Any        =     "bias_ih_l0"
												elif "bias_hh_l0" in name:
															UpperCAmelCase__  :					Union[str, Any]        =     "bias_hh_l0"
												elif "weight_ih_l1" in name:
															UpperCAmelCase__  :					Optional[Any]        =     "weight_ih_l1"
												elif "weight_hh_l1" in name:
															UpperCAmelCase__  :					str        =     "weight_hh_l1"
												elif "bias_ih_l1" in name:
															UpperCAmelCase__  :					Optional[Any]        =     "bias_ih_l1"
												elif "bias_hh_l1" in name:
															UpperCAmelCase__  :					int        =     "bias_hh_l1"
												elif "bias" in name:
															UpperCAmelCase__  :					Union[str, Any]        =     "bias"
												elif "weight" in name:
															UpperCAmelCase__  :					Optional[int]        =     "weight"
												elif "running_mean" in name:
															UpperCAmelCase__  :					List[str]        =     "running_mean"
												elif "running_var" in name:
															UpperCAmelCase__  :					int        =     "running_var"
												elif "num_batches_tracked" in name:
															UpperCAmelCase__  :					Union[str, Any]        =     "num_batches_tracked"
												else:
															UpperCAmelCase__  :					Union[str, Any]        =     None
												set_recursively(UpperCamelCase__ ,      UpperCamelCase__ ,      UpperCamelCase__ ,      UpperCamelCase__ ,      UpperCamelCase__					)
									continue
						if not is_used:
									unused_weights.append(UpperCamelCase__					)

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







@torch.no_grad()
def     _UpperCamelCase      (		UpperCamelCase__ ,      UpperCamelCase__ ,      UpperCamelCase__ ,      UpperCamelCase__=None ,      UpperCamelCase__=None ,      ):

			'''simple docstring'''





			if config_path is not None:
						UpperCAmelCase__  :					int        =     EncodecConfig.from_pretrained(UpperCamelCase__					)
			else:
						UpperCAmelCase__  :					Optional[Any]        =     EncodecConfig()

			if model_name == "encodec_24khz":
						pass  # config is already correct
			elif model_name == "encodec_32khz":
						UpperCAmelCase__  :					Any        =     [8, 5, 4, 4]
						UpperCAmelCase__  :					Union[str, Any]        =     [2.2]
						UpperCAmelCase__  :					Union[str, Any]        =     6_4
						UpperCAmelCase__  :					Optional[int]        =     3_2_0_0_0
						UpperCAmelCase__  :					int        =     2_0_4_8
						UpperCAmelCase__  :					str        =     False
						UpperCAmelCase__  :					List[str]        =     False
						UpperCAmelCase__  :					List[str]        =     False
			elif model_name == "encodec_48khz":
						UpperCAmelCase__  :					Any        =     [8, 5, 4, 2]
						UpperCAmelCase__  :					List[Any]        =     [3.0, 6.0, 12.0, 24.0]
						UpperCAmelCase__  :					Optional[int]        =     4_8_0_0_0
						UpperCAmelCase__  :					Any        =     2
						UpperCAmelCase__  :					Tuple        =     False
						UpperCAmelCase__  :					int        =     "time_group_norm"
						UpperCAmelCase__  :					Union[str, Any]        =     True
						UpperCAmelCase__  :					Any        =     1.0
						UpperCAmelCase__  :					Dict        =     0.01
			else:
						raise ValueError(f'''Unknown model name: {model_name}'''					)

			UpperCAmelCase__  :					Optional[int]        =     EncodecModel(UpperCamelCase__					)

			UpperCAmelCase__  :					Union[str, Any]        =     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(UpperCamelCase__					)

			UpperCAmelCase__  :					Union[str, Any]        =     torch.load(UpperCamelCase__					)
			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__  :					int        =     original_checkpoint["best_state"]
			recursively_load_weights(UpperCamelCase__ ,      UpperCamelCase__ ,      UpperCamelCase__					)
			model.save_pretrained(UpperCamelCase__					)

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


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

		__A	=parser.parse_args()
		convert_checkpoint(
		    args.model,
		    args.checkpoint_path,
		    args.pytorch_dump_folder_path,
		    args.config_path,
		    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)