



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_	       =			[
    '''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_	       =			[
    '''mlp.dense_4h_to_h.weight''',
    '''self_attention.dense.weight''',
]






def 					__lowerCAmelCase		(  __lowerCamelCase	:    Any				,   __lowerCamelCase	:    List[str]	)  ->						List[str]:
					__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*).*"""				,   __lowerCamelCase	)[1]	)
					layer_number -= 3
					return f"""h.{layer_number}.""" + key






def 					__lowerCAmelCase		(  __lowerCamelCase	:    Any	)  ->						Optional[int]:
					if dtype == torch.bool:
										return 1 / 8
					__lowerCAmelCase	   =re.search(r"""[^\d](\d+)$"""				,   str(__lowerCamelCase	)	)
					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 					__lowerCAmelCase		(  __lowerCamelCase	:    Any				,   __lowerCamelCase	:    Optional[int]				,   __lowerCamelCase	:    str				,   __lowerCamelCase	:    Union[str, Any]				,   __lowerCamelCase	:    List[str]	)  ->						Union[str, Any]:
					# Construct model
					if bloom_config_file == "":
										__lowerCAmelCase	   =BloomConfig()
					else:
										__lowerCAmelCase	   =BloomConfig.from_json_file(__lowerCamelCase	)

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

										__lowerCAmelCase	   ={"""weight_map""": {}, """metadata""": {}}
										__lowerCAmelCase	   =0

										__lowerCAmelCase	   =None

										__lowerCAmelCase	   =BloomConfig()

										for j, file in enumerate(__lowerCamelCase	):
															print("""Processing file: {}""".format(__lowerCamelCase	)	)
															__lowerCAmelCase	   =None

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

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

																				if tensors is None:
																									__lowerCAmelCase	   =temp
																				else:
																									for key in tensors.keys():
																														if any(key.endswith(__lowerCamelCase	) 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=__lowerCamelCase	)

            # Divide by the number of TP the weights we want to average
															for key in tensors.keys():
																				if any(key.endswith(__lowerCamelCase	) for end in WEIGHTS_TO_AVERAGE_ENDSWITH	):
																									__lowerCAmelCase	   =tensors[key] / pretraining_tp
															torch.save(
															    __lowerCamelCase				,   os.path.join(
															        __lowerCamelCase				,   """pytorch_model_{}-of-{}.bin""".format(str(j + 1	).zfill(5	)				,   str(len(__lowerCamelCase	)	).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(__lowerCamelCase	)	).zfill(5	)	)

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

										__lowerCAmelCase	   =os.listdir(__lowerCamelCase	)
										__lowerCAmelCase	   =sorted(filter(lambda __lowerCamelCase	: s.startswith("""layer"""	) and "model_00" in s				,   __lowerCamelCase	)	)

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

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

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

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

															__lowerCAmelCase	   =model.load_state_dict(__lowerCamelCase				,   strict=__lowerCamelCase	)
															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(__lowerCamelCase				,   exist_ok=__lowerCamelCase	)
										__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()				,   __lowerCamelCase	)
										print(f"""Save configuration file to {pytorch_config_dump_path}"""	)
										with open(__lowerCamelCase				,   """w"""				,   encoding="""utf-8"""	) as f:
															f.write(config.to_json_string()	)


if __name__ == "__main__":
				lowercase_	       =			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_	       =			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,
				)









import os
from typing import Dict, List, Tuple, TypeVar, Union


lowercase_	       =			TypeVar('''T''')

lowercase_	       =			Union[List[T], Tuple[T, ...]]
lowercase_	       =			Union[T, List[T], Dict[str, T]]
lowercase_	       =			Union[str, bytes, os.PathLike]




