


'''simple docstring'''



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()

_lowercase    :    int			     =  [
    """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""",
]

_lowercase    :    Any			     =  [
    """mlp.dense_4h_to_h.weight""",
    """self_attention.dense.weight""",
]


def    lowerCamelCase__   ( A     :    Dict   ,     A     :    Optional[int]  ):




						'''simple docstring'''






						UpperCAmelCase         =			{
						    '''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
						UpperCAmelCase         =			int(re.match(R'''.*layer_(\d*).*'''   ,     A  )[1]  )
						layer_number -= 3
						return f"""h.{layer_number}.""" + key


def    lowerCamelCase__   ( A     :    Dict  ):




						'''simple docstring'''






						if dtype == torch.bool:
												return 1 / 8
						UpperCAmelCase         =			re.search(R'''[^\d](\d+)$'''   ,     str(A  )  )
						if bit_search is None:
												raise ValueError(f"""`dtype` is not a valid dtype: {dtype}."""  )
						UpperCAmelCase         =			int(bit_search.groups()[0]  )
						return bit_size // 8


def    lowerCamelCase__   ( A     :    Union[str, Any]   ,     A     :    str   ,     A     :    Any   ,     A     :    List[Any]   ,     A     :    str  ):




						'''simple docstring'''






						if bloom_config_file == "":
												UpperCAmelCase         =			BloomConfig()
						else:
												UpperCAmelCase         =			BloomConfig.from_json_file(A  )

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

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

												UpperCAmelCase         =			None

												UpperCAmelCase         =			BloomConfig()

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

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

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

																								if tensors is None:
																														UpperCAmelCase         =			temp
																								else:
																														for key in tensors.keys():
																																				if any(key.endswith(A  ) 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
																																										UpperCAmelCase         =			1 if any(text in key for text in WEIGHTS_WITH_ROW_PARALLELISM_CONTAIN  ) else 0
																																										# We concatenate these weights accross TP ranks
																																										UpperCAmelCase         =			torch.cat([tensors[key], temp[key]]   ,     dim=A  )

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

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

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

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

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

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

																								if tensors is None:
																														UpperCAmelCase         =			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(A  ) for end in WEIGHTS_TO_AVERAGE_ENDSWITH  ):
																																										tensors[key] += temp[key]
																																				else:
																																										# Some weights are RowParallelLinear in Megatron-Deepspeed, others are ColumnParallel
																																										UpperCAmelCase         =			1 if any(text in key for text in WEIGHTS_WITH_ROW_PARALLELISM_CONTAIN  ) else 0
																																										# We concatenate these weights accross TP ranks
																																										UpperCAmelCase         =			torch.cat([tensors[key], temp[key]]   ,     dim=A  )

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

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

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

												# Save pytorch-model
												os.makedirs(A   ,     exist_ok=A  )
												UpperCAmelCase         =			pytorch_dump_folder_path + '''/''' + WEIGHTS_NAME
												UpperCAmelCase         =			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:
																		UpperCAmelCase         =			model.to(config.torch_dtype  )
												torch.save(model.state_dict()   ,     A  )
												print(f"""Save configuration file to {pytorch_config_dump_path}"""  )
												with open(A   ,     '''w'''   ,     encoding='''utf-8'''  ) as f:
																		f.write(config.to_json_string()  )


if __name__ == "__main__":
	_lowercase    :    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""",
	)
	_lowercase    :    List[Any]			     =  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,
	)









'''simple docstring'''



import math
from typing import Optional

import numpy as np

from ...configuration_utils import PretrainedConfig
from ...utils import logging


_lowercase    :    Union[str, Any]			     =  logging.get_logger(__name__)

_lowercase    :    List[str]			     =  {
    """facebook/encodec_24khz""": """https://huggingface.co/facebook/encodec_24khz/resolve/main/config.json""",
    """facebook/encodec_48khz""": """https://huggingface.co/facebook/encodec_48khz/resolve/main/config.json""",
}






class  UpperCamelCase__(     lowerCAmelCase ):
				__magic_name__       :   List[str]							= "encodec"




				def __init__(	self    :				List[str]   ,					lowerCAmelCase    :				int=[1.5, 3.0, 6.0, 12.0, 24.0]   ,					lowerCAmelCase    :				Tuple=24000   ,					lowerCAmelCase    :				List[Any]=1   ,					lowerCAmelCase    :				Union[str, Any]=False   ,					lowerCAmelCase    :				str=None   ,					lowerCAmelCase    :				Dict=None   ,					lowerCAmelCase    :				str=128   ,					lowerCAmelCase    :				Any=32   ,					lowerCAmelCase    :				Any=1   ,					lowerCAmelCase    :				List[Any]=[8, 5, 4, 2]   ,					lowerCAmelCase    :				Union[str, Any]="weight_norm"   ,					lowerCAmelCase    :				str=7   ,					lowerCAmelCase    :				Optional[int]=7   ,					lowerCAmelCase    :				Any=3   ,					lowerCAmelCase    :				Tuple=2   ,					lowerCAmelCase    :				Optional[Any]=True   ,					lowerCAmelCase    :				List[str]="reflect"   ,					lowerCAmelCase    :				Optional[int]=2   ,					lowerCAmelCase    :				List[Any]=2   ,					lowerCAmelCase    :				Union[str, Any]=1.0   ,					lowerCAmelCase    :				Optional[Any]=1024   ,					lowerCAmelCase    :				Optional[Any]=None   ,					lowerCAmelCase    :				str=True   ,					**lowerCAmelCase    :				str   ,					)->  List[str]:

										"""simple docstring"""




										UpperCAmelCase         =			target_bandwidths
										UpperCAmelCase         =			sampling_rate
										UpperCAmelCase         =			audio_channels
										UpperCAmelCase         =			normalize
										UpperCAmelCase         =			chunk_length_s
										UpperCAmelCase         =			overlap
										UpperCAmelCase         =			hidden_size
										UpperCAmelCase         =			num_filters
										UpperCAmelCase         =			num_residual_layers
										UpperCAmelCase         =			upsampling_ratios
										UpperCAmelCase         =			norm_type
										UpperCAmelCase         =			kernel_size
										UpperCAmelCase         =			last_kernel_size
										UpperCAmelCase         =			residual_kernel_size
										UpperCAmelCase         =			dilation_growth_rate
										UpperCAmelCase         =			use_causal_conv
										UpperCAmelCase         =			pad_mode
										UpperCAmelCase         =			compress
										UpperCAmelCase         =			num_lstm_layers
										UpperCAmelCase         =			trim_right_ratio
										UpperCAmelCase         =			codebook_size
										UpperCAmelCase         =			codebook_dim if codebook_dim is not None else hidden_size
										UpperCAmelCase         =			use_conv_shortcut

										if self.norm_type not in ["weight_norm", "time_group_norm"]:
																raise ValueError(
																    F"""self.norm_type must be one of `\"weight_norm\"`, `\"time_group_norm\"`), got {self.norm_type}"""   )

										super().__init__(**lowerCAmelCase   )




				@property
				def   a__(	self    :				str   )->  Optional[int]:

										"""simple docstring"""




										if self.chunk_length_s is None:
																return None
										else:
																return int(self.chunk_length_s * self.sampling_rate   )




				@property
				def   a__(	self    :				List[str]   )->  Optional[int]:

										"""simple docstring"""




										if self.chunk_length_s is None or self.overlap is None:
																return None
										else:
																return max(1   ,					int((1.0 - self.overlap) * self.chunk_length   )   )




				@property
				def   a__(	self    :				List[Any]   )->  int:

										"""simple docstring"""




										UpperCAmelCase         =			np.prod(self.upsampling_ratios   )
										return math.ceil(self.sampling_rate / hop_length   )





				@property
				def   a__(	self    :				int   )->  int:

										"""simple docstring"""




										return int(1000 * self.target_bandwidths[-1] // (self.frame_rate * 10)   )





