



"""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()
__lowerCAmelCase							:       Optional[int]          =       logging.get_logger('''transformers.models.encodec''')

__lowerCAmelCase							:       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''',
}
__lowerCAmelCase							:       Optional[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''',
}
__lowerCAmelCase							:       Dict          =       {
    '''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''',
}
__lowerCAmelCase							:       Any          =       {
    '''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''',
}
__lowerCAmelCase							:       List[str]          =       {
    '''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''',
}
__lowerCAmelCase							:       str          =       {
    **MAPPING_QUANTIZER,
    **MAPPING_ENCODER,
    **MAPPING_DECODER,
}
__lowerCAmelCase							:       Optional[Any]          =       {
    **MAPPING_QUANTIZER,
    **MAPPING_ENCODER,
    **MAPPING_ENCODER_48K,
    **MAPPING_DECODER,
    **MAPPING_DECODER_48K,
}
__lowerCAmelCase							:       List[Any]          =       []
__lowerCAmelCase							:       List[str]          =       []
def 						__lowerCAmelCase    (		__UpperCamelCase   :      Tuple       ,			__UpperCamelCase   :      List[Any]       ,			__UpperCamelCase   :      str       ,			__UpperCamelCase   :      int       ,			__UpperCamelCase   :      Tuple     ):


						'''simple docstring'''






						for attribute in key.split("""."""     ):
												snake_case_ :			str            =      getattr(__UpperCamelCase       ,			__UpperCamelCase     )

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

						logger.info(F'{key + ("." + weight_type if weight_type is not None else "")} was initialized from {full_name}.'     )
def 						__lowerCAmelCase    (		__UpperCamelCase   :      List[str]       ,			__UpperCamelCase   :      Tuple     ):


						'''simple docstring'''






						for key in ignore_keys:
												if key.endswith(""".*"""     ):
																		if name.startswith(key[:-1]     ):
																								return True
												elif ".*." in key:
																		snake_case_      ,		snake_case_ :			Any            =      key.split(""".*."""     )
																		if prefix in name and suffix in name:
																								return True
												elif key in name:
																		return True
						return False
def 						__lowerCAmelCase    (		__UpperCamelCase   :      List[Any]       ,			__UpperCamelCase   :      List[str]       ,			__UpperCamelCase   :      Tuple     ):


						'''simple docstring'''






						snake_case_ :			int            =      []

						if model_name == "encodec_24khz" or "encodec_32khz":
												snake_case_ :			str            =      MAPPING_24K
						elif model_name == "encodec_48khz":
												snake_case_ :			Optional[int]            =      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

												snake_case_ :			Dict            =      False
												for key, mapped_key in MAPPING.items():
																		if "*" in key:
																								snake_case_      ,		snake_case_ :			str            =      key.split(""".*."""     )
																								if prefix in name and suffix in name:
																														snake_case_ :			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

																								snake_case_ :			Tuple            =      True
																								if "*" in mapped_key:
																														snake_case_ :			List[str]            =      name.split(__UpperCamelCase     )[0].split("""."""     )[-2]
																														snake_case_ :			List[str]            =      mapped_key.replace("""*"""       ,			__UpperCamelCase     )
																								if "weight_g" in name:
																														snake_case_ :			int            =      """weight_g"""
																								elif "weight_v" in name:
																														snake_case_ :			Union[str, Any]            =      """weight_v"""
																								elif "weight_ih_l0" in name:
																														snake_case_ :			Optional[int]            =      """weight_ih_l0"""
																								elif "weight_hh_l0" in name:
																														snake_case_ :			Optional[int]            =      """weight_hh_l0"""
																								elif "bias_ih_l0" in name:
																														snake_case_ :			List[str]            =      """bias_ih_l0"""
																								elif "bias_hh_l0" in name:
																														snake_case_ :			Optional[Any]            =      """bias_hh_l0"""
																								elif "weight_ih_l1" in name:
																														snake_case_ :			List[str]            =      """weight_ih_l1"""
																								elif "weight_hh_l1" in name:
																														snake_case_ :			str            =      """weight_hh_l1"""
																								elif "bias_ih_l1" in name:
																														snake_case_ :			Any            =      """bias_ih_l1"""
																								elif "bias_hh_l1" in name:
																														snake_case_ :			List[Any]            =      """bias_hh_l1"""
																								elif "bias" in name:
																														snake_case_ :			Union[str, Any]            =      """bias"""
																								elif "weight" in name:
																														snake_case_ :			List[Any]            =      """weight"""
																								elif "running_mean" in name:
																														snake_case_ :			Optional[int]            =      """running_mean"""
																								elif "running_var" in name:
																														snake_case_ :			int            =      """running_var"""
																								elif "num_batches_tracked" in name:
																														snake_case_ :			List[str]            =      """num_batches_tracked"""
																								else:
																														snake_case_ :			Tuple            =      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 						__lowerCAmelCase    (		__UpperCamelCase   :      int       ,			__UpperCamelCase   :      Tuple       ,			__UpperCamelCase   :      Optional[Any]       ,			__UpperCamelCase   :      Optional[Any]=None       ,			__UpperCamelCase   :      List[str]=None       ,			):


						'''simple docstring'''






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

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

						snake_case_ :			Optional[Any]            =      EncodecModel(__UpperCamelCase     )

						snake_case_ :			Tuple            =      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     )

						snake_case_ :			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
												snake_case_ :			Optional[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__":
					__lowerCAmelCase							:       List[str]          =       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.'''
					)

					__lowerCAmelCase							:       List[str]          =       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 math import factorial

SCREAMING_SNAKE_CASE     =       {str(digit): factorial(digit) for digit in range(1_0)}







def snake_case_     (					lowercase__		):
  if not isinstance(lowercase__   ,       lowercase__		):
    raise TypeError("Parameter number must be int"		)

  if number < 0:
    raise ValueError("Parameter number must be greater than or equal to 0"		)

  # Converts number in string to iterate on its digits and adds its factorial.
  return sum(DIGIT_FACTORIAL[digit] for digit in str(lowercase__		)		)







def snake_case_     (					lowercase__ = 6_0   ,       lowercase__ = 1_0_0_0_0_0_0		):

  if not isinstance(lowercase__   ,       lowercase__		) or not isinstance(lowercase__   ,       lowercase__		):
    raise TypeError("Parameters chain_length and number_limit must be int"		)

  if chain_length <= 0 or number_limit <= 0:
    raise ValueError(
        "Parameters chain_length and number_limit must be greater than 0"		)

  # the counter for the chains with the exact desired length
  UpperCAmelCase__      :					Union[str, Any]				       =							0
  # the cached sizes of the previous chains
  UpperCAmelCase__      :					dict[int, int]				       =							{}

  for start_chain_element in range(1   ,       lowercase__		):
    # The temporary set will contain the elements of the chain
    UpperCAmelCase__      :					Any				       =							set()
    UpperCAmelCase__      :					int				       =							0

    # Stop computing the chain when you find a cached size, a repeating item or the
    # length is greater then the desired one.
    UpperCAmelCase__      :					List[str]				       =							start_chain_element
    while (
        chain_element not in chain_sets_lengths
        and chain_element not in chain_set
        and chain_set_length <= chain_length
    ):
      chain_set.add(lowercase__		)
      chain_set_length += 1
      UpperCAmelCase__      :					List[str]				       =							digit_factorial_sum(lowercase__		)

    if chain_element in chain_sets_lengths:
      chain_set_length += chain_sets_lengths[chain_element]

    UpperCAmelCase__      :					List[str]				       =							chain_set_length

    # If chain contains the exact amount of elements increase the counter
    if chain_set_length == chain_length:
      chains_counter += 1

  return chains_counter


if __name__ == "__main__":
      import doctest

      doctest.testmod()
      print(F'{solution()}')



