






from collections import deque

from .hash_table import HashTable


class     __SCREAMING_SNAKE_CASE					(   __lowercase):






						def __init__(	self     ,						*_UpperCamelCase     ,						**_UpperCamelCase   ):






												"""simple docstring"""
												super().__init__(*_UpperCamelCase     ,						**_UpperCamelCase   )






						def       UpperCamelCase__      (	self     ,						_UpperCamelCase     ,						_UpperCamelCase   ):






												"""simple docstring"""
												lowerCAmelCase__    						=						deque([]   ) if self.values[key] is None else self.values[key]
												self.values[key].appendleft(_UpperCamelCase   )
												lowerCAmelCase__    						=						self.values[key]






						def       UpperCamelCase__      (	self   ):






												"""simple docstring"""
												return (
												    sum(self.charge_factor - len(_UpperCamelCase   ) for slot in self.values   )
												    / self.size_table
												    * self.charge_factor
												)







						def       UpperCamelCase__      (	self     ,						_UpperCamelCase     ,						_UpperCamelCase=None   ):






												"""simple docstring"""
												if not (
												    len(self.values[key]   ) == self.charge_factor and self.values.count(_UpperCamelCase   ) == 0
												):
																		return key
												return super()._collision_resolution(_UpperCamelCase     ,						_UpperCamelCase   )













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

__snake_case    :    str         =     [
    """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""",
]

__snake_case    :    int         =     [
    """mlp.dense_4h_to_h.weight""",
    """self_attention.dense.weight""",
]







def 	_UpperCamelCase   (    UpperCamelCase_       :					Dict    ,       UpperCamelCase_       :					Dict		)  ->   Dict:

						"""simple docstring"""



						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*).*'    ,       UpperCamelCase_		)[1]		)
						layer_number -= 3
						return F"h.{layer_number}." + key







def 	_UpperCamelCase   (    UpperCamelCase_       :					Any		)  ->   Optional[Any]:

						"""simple docstring"""



						if dtype == torch.bool:
												return 1 / 8
						lowerCAmelCase__    						=						re.search(r'[^\d](\d+)$'    ,       str(UpperCamelCase_		)		)
						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 	_UpperCamelCase   (    UpperCamelCase_       :					List[Any]    ,       UpperCamelCase_       :					Any    ,       UpperCamelCase_       :					Optional[Any]    ,       UpperCamelCase_       :					Dict    ,       UpperCamelCase_       :					Tuple		)  ->   Union[str, Any]:

						"""simple docstring"""



						if bloom_config_file == "":
												lowerCAmelCase__    						=						BloomConfig()
						else:
												lowerCAmelCase__    						=						BloomConfig.from_json_file(UpperCamelCase_		)

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

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

												lowerCAmelCase__    						=						None

												lowerCAmelCase__    						=						BloomConfig()

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

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

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

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

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

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

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

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

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

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

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

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


if __name__ == "__main__":
					__snake_case    :    Union[str, Any]         =     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""",
					)
					__snake_case    :    str         =     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,
					)





