


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

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


def  _snake_case       (					A			,   A			,   A			,   A			,   A	)						->   List[Any]:
							for attribute in key.split('''.'''	):
														lowerCAmelCase__	       =     getattr(A			,   A	)

							if weight_type is not None:
														lowerCAmelCase__	       =     getattr(A			,   A	).shape
							else:
														lowerCAmelCase__	       =     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":
														lowerCAmelCase__	       =     value
							elif weight_type == "weight_g":
														lowerCAmelCase__	       =     value
							elif weight_type == "weight_v":
														lowerCAmelCase__	       =     value
							elif weight_type == "bias":
														lowerCAmelCase__	       =     value
							elif weight_type == "running_mean":
														lowerCAmelCase__	       =     value
							elif weight_type == "running_var":
														lowerCAmelCase__	       =     value
							elif weight_type == "num_batches_tracked":
														lowerCAmelCase__	       =     value
							elif weight_type == "weight_ih_l0":
														lowerCAmelCase__	       =     value
							elif weight_type == "weight_hh_l0":
														lowerCAmelCase__	       =     value
							elif weight_type == "bias_ih_l0":
														lowerCAmelCase__	       =     value
							elif weight_type == "bias_hh_l0":
														lowerCAmelCase__	       =     value
							elif weight_type == "weight_ih_l1":
														lowerCAmelCase__	       =     value
							elif weight_type == "weight_hh_l1":
														lowerCAmelCase__	       =     value
							elif weight_type == "bias_ih_l1":
														lowerCAmelCase__	       =     value
							elif weight_type == "bias_hh_l1":
														lowerCAmelCase__	       =     value
							else:
														lowerCAmelCase__	       =     value

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


def  _snake_case       (					A			,   A	)						->   Optional[Any]:
							for key in ignore_keys:
														if key.endswith('''.*'''	):
																					if name.startswith(key[:-1]	):
																												return True
														elif ".*." in key:
																					lowerCAmelCase__       , lowerCAmelCase__	       =     key.split('''.*.'''	)
																					if prefix in name and suffix in name:
																												return True
														elif key in name:
																					return True
							return False


def  _snake_case       (					A			,   A			,   A	)						->   List[str]:
							lowerCAmelCase__	       =     []

							if model_name == "encodec_24khz" or "encodec_32khz":
														lowerCAmelCase__	       =     MAPPING_24K
							elif model_name == "encodec_48khz":
														lowerCAmelCase__	       =     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

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

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

																												lowerCAmelCase__	       =     True
																												if "*" in mapped_key:
																																			lowerCAmelCase__	       =     name.split(A	)[0].split('''.'''	)[-2]
																																			lowerCAmelCase__	       =     mapped_key.replace('''*'''			,   A	)
																												if "weight_g" in name:
																																			lowerCAmelCase__	       =     '''weight_g'''
																												elif "weight_v" in name:
																																			lowerCAmelCase__	       =     '''weight_v'''
																												elif "weight_ih_l0" in name:
																																			lowerCAmelCase__	       =     '''weight_ih_l0'''
																												elif "weight_hh_l0" in name:
																																			lowerCAmelCase__	       =     '''weight_hh_l0'''
																												elif "bias_ih_l0" in name:
																																			lowerCAmelCase__	       =     '''bias_ih_l0'''
																												elif "bias_hh_l0" in name:
																																			lowerCAmelCase__	       =     '''bias_hh_l0'''
																												elif "weight_ih_l1" in name:
																																			lowerCAmelCase__	       =     '''weight_ih_l1'''
																												elif "weight_hh_l1" in name:
																																			lowerCAmelCase__	       =     '''weight_hh_l1'''
																												elif "bias_ih_l1" in name:
																																			lowerCAmelCase__	       =     '''bias_ih_l1'''
																												elif "bias_hh_l1" in name:
																																			lowerCAmelCase__	       =     '''bias_hh_l1'''
																												elif "bias" in name:
																																			lowerCAmelCase__	       =     '''bias'''
																												elif "weight" in name:
																																			lowerCAmelCase__	       =     '''weight'''
																												elif "running_mean" in name:
																																			lowerCAmelCase__	       =     '''running_mean'''
																												elif "running_var" in name:
																																			lowerCAmelCase__	       =     '''running_var'''
																												elif "num_batches_tracked" in name:
																																			lowerCAmelCase__	       =     '''num_batches_tracked'''
																												else:
																																			lowerCAmelCase__	       =     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  _snake_case       (					A			,   A			,   A			,   A=None			,   A=None			,   )						->   Any:
							if config_path is not None:
														lowerCAmelCase__	       =     EncodecConfig.from_pretrained(A	)
							else:
														lowerCAmelCase__	       =     EncodecConfig()

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

							lowerCAmelCase__	       =     EncodecModel(A	)

							lowerCAmelCase__	       =     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	)

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



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

__UpperCAmelCase       =							[
    '''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''',
]

__UpperCAmelCase       =							[
    '''mlp.dense_4h_to_h.weight''',
    '''self_attention.dense.weight''',
]


def  _snake_case       (					A			,   A	)						->   Optional[Any]:
							lowerCAmelCase__	       =     {
							    '''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
							lowerCAmelCase__	       =     int(re.match(R'''.*layer_(\d*).*'''			,   A	)[1]	)
							layer_number -= 3
							return F"""h.{layer_number}.""" + key


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


def  _snake_case       (					A			,   A			,   A			,   A			,   A	)						->   Dict:
							# Construct model
							if bloom_config_file == "":
														lowerCAmelCase__	       =     BloomConfig()
							else:
														lowerCAmelCase__	       =     BloomConfig.from_json_file(A	)

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

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

														lowerCAmelCase__	       =     None

														lowerCAmelCase__	       =     BloomConfig()

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

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

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

																												if tensors is None:
																																			lowerCAmelCase__	       =     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
																																																	lowerCAmelCase__	       =     1 if any(text in key for text in WEIGHTS_WITH_ROW_PARALLELISM_CONTAIN	) else 0
																																																	# We concatenate these weights accross TP ranks
																																																	lowerCAmelCase__	       =     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	):
																																			lowerCAmelCase__	       =     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():
																												lowerCAmelCase__	       =     tensors[key]
																												total_size += value.numel() * get_dtype_size(value.dtype	)
																												if key not in index_dict["weight_map"]:
																																			lowerCAmelCase__	       =     '''pytorch_model_{}-of-{}.bin'''.format(
																																			    str(j + 1	).zfill(5	)			,   str(len(A	)	).zfill(5	)	)

														lowerCAmelCase__	       =     BloomConfig()
														lowerCAmelCase__	       =     pytorch_dump_folder_path + '''/''' + CONFIG_NAME
														lowerCAmelCase__	       =     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:
																					lowerCAmelCase__	       =     json.dumps(A			,   indent=2			,   sort_keys=A	) + '''\n'''
																					f.write(A	)
							else:
														lowerCAmelCase__	       =     BloomModel(A	)

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

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

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

																												if tensors is None:
																																			lowerCAmelCase__	       =     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
																																																	lowerCAmelCase__	       =     1 if any(text in key for text in WEIGHTS_WITH_ROW_PARALLELISM_CONTAIN	) else 0
																																																	# We concatenate these weights accross TP ranks
																																																	lowerCAmelCase__	       =     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	):
																																			lowerCAmelCase__	       =     tensors[key] / pretraining_tp

																					lowerCAmelCase__	       =     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:
																												lowerCAmelCase__	       =     set(other_keys.missing_keys	)
																					else:
																												lowerCAmelCase__	       =     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	)
														lowerCAmelCase__	       =     pytorch_dump_folder_path + '''/''' + WEIGHTS_NAME
														lowerCAmelCase__	       =     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:
																					lowerCAmelCase__	       =     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__":
		__UpperCAmelCase       =							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''',
		)
		__UpperCAmelCase       =							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,
		)