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

SCREAMING_SNAKE_CASE__   :			Dict												=    [
    "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",
]

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




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




def  __magic_name__ (   __lowerCAmelCase	:						Dict							)			->   Any:
				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  __magic_name__ (   __lowerCAmelCase	:						Optional[Any]		,  __lowerCAmelCase	:						Union[str, Any]		,  __lowerCAmelCase	:						List[Any]		,  __lowerCAmelCase	:						List[Any]		,  __lowerCAmelCase	:						Optional[Any]							)			->   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__":
					SCREAMING_SNAKE_CASE__   :			List[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",
					)
					SCREAMING_SNAKE_CASE__   :			Optional[int]												=    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,
					)





from queue import PriorityQueue
from typing import Any

import numpy as np




def  __magic_name__ (   __lowerCAmelCase	:						dict		,  __lowerCAmelCase	:						str		,  __lowerCAmelCase	:						set		,  __lowerCAmelCase	:						set		,  __lowerCAmelCase	:						dict		,  __lowerCAmelCase	:						dict		,  __lowerCAmelCase	:						PriorityQueue		,  __lowerCAmelCase	:						dict		,  __lowerCAmelCase	:						float | int		,  )			->   float | int:
				for nxt, d in graph[v]:
								if nxt in visited_forward:
												continue
								__lowerCamelCase    					=	cst_fwd.get(__lowerCAmelCase		,  np.inf							)
								__lowerCamelCase    					=	cst_fwd[v] + d
								if new_cost_f < old_cost_f:
												queue.put((new_cost_f, nxt)							)
												__lowerCamelCase    					=	new_cost_f
												__lowerCamelCase    					=	v
								if nxt in visited_backward:
												if cst_fwd[v] + d + cst_bwd[nxt] < shortest_distance:
																__lowerCamelCase    					=	cst_fwd[v] + d + cst_bwd[nxt]
				return shortest_distance




def  __magic_name__ (   __lowerCAmelCase	:						str		,  __lowerCAmelCase	:						str		,  __lowerCAmelCase	:						dict		,  __lowerCAmelCase	:						dict							)			->   int:
				__lowerCamelCase    					=	-1

				__lowerCamelCase    					=	set()
				__lowerCamelCase    					=	set()
				__lowerCamelCase    					=	{source: 0}
				__lowerCamelCase    					=	{destination: 0}
				__lowerCamelCase    					=	{source: None}
				__lowerCamelCase    					=	{destination: None}
				__lowerCamelCase    					=	PriorityQueue()
				__lowerCamelCase    					=	PriorityQueue()

				__lowerCamelCase    					=	np.inf

				queue_forward.put((0, source)							)
				queue_backward.put((0, destination)							)

				if source == destination:
								return 0

				while not queue_forward.empty() and not queue_backward.empty():
								__lowerCamelCase					,    __lowerCamelCase    					=	queue_forward.get()
								visited_forward.add(__lowerCAmelCase							)

								__lowerCamelCase					,    __lowerCamelCase    					=	queue_backward.get()
								visited_backward.add(__lowerCAmelCase							)

								__lowerCamelCase    					=	pass_and_relaxation(
								    __lowerCAmelCase		,  __lowerCAmelCase		,  __lowerCAmelCase		,  __lowerCAmelCase		,  __lowerCAmelCase		,  __lowerCAmelCase		,  __lowerCAmelCase		,  __lowerCAmelCase		,  __lowerCAmelCase		,  )

								__lowerCamelCase    					=	pass_and_relaxation(
								    __lowerCAmelCase		,  __lowerCAmelCase		,  __lowerCAmelCase		,  __lowerCAmelCase		,  __lowerCAmelCase		,  __lowerCAmelCase		,  __lowerCAmelCase		,  __lowerCAmelCase		,  __lowerCAmelCase		,  )

								if cst_fwd[v_fwd] + cst_bwd[v_bwd] >= shortest_distance:
												break

				if shortest_distance != np.inf:
								__lowerCamelCase    					=	shortest_distance
				return shortest_path_distance


SCREAMING_SNAKE_CASE__   :			List[Any]												=    {
    "B": [["C", 1]],
    "C": [["D", 1]],
    "D": [["F", 1]],
    "E": [["B", 1], ["G", 2]],
    "F": [],
    "G": [["F", 1]],
}
SCREAMING_SNAKE_CASE__   :			Optional[int]												=    {
    "B": [["E", 1]],
    "C": [["B", 1]],
    "D": [["C", 1]],
    "F": [["D", 1], ["G", 1]],
    "E": [[None, np.inf]],
    "G": [["E", 2]],
}

if __name__ == "__main__":
					import doctest

					doctest.testmod()



