






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

A__	:			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',
]

A__	:			Dict			   = [
    'mlp.dense_4h_to_h.weight',
    'self_attention.dense.weight',
]






def   a			(						lowerCamelCase_		,   lowerCamelCase_				):






				'''simple docstring'''
				lowercase__    =			{
				    '''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
				lowercase__    =			int(re.match(r'''.*layer_(\d*).*'''		,   lowerCamelCase_				)[1]				)
				layer_number -= 3
				return F"""h.{layer_number}.""" + key






def   a			(						lowerCamelCase_				):






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






def   a			(						lowerCamelCase_		,   lowerCamelCase_		,   lowerCamelCase_		,   lowerCamelCase_		,   lowerCamelCase_				):






				'''simple docstring'''
				# Construct model
				if bloom_config_file == "":
								lowercase__    =			BloomConfig()
				else:
								lowercase__    =			BloomConfig.from_json_file(lowerCamelCase_				)

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

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

								lowercase__    =			None

								lowercase__    =			BloomConfig()

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

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

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

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

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

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

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

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

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

												lowercase__    =			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:
																lowercase__    =			set(other_keys.missing_keys				)
												else:
																lowercase__    =			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_				)
								lowercase__    =			pytorch_dump_folder_path + '''/''' + WEIGHTS_NAME
								lowercase__    =			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:
												lowercase__    =			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__":
							A__	:			Optional[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',
							)
							A__	:			Dict			   = 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 pathlib import PurePosixPath
from typing import Optional

import fsspec
from fsspec import AbstractFileSystem
from huggingface_hub.hf_api import DatasetInfo

from ..utils.file_utils import get_authentication_headers_for_url
from ..utils.hub import hf_hub_url
class       _UpperCAmelCase					(       A__	):

				"""simple docstring"""




				lowercase__											=       """"""
				lowercase__											=       """hf-legacy"""  # "hf://"" is reserved for hffs




				def __init__(     self						:    List[str],  lowerCamelCase						:    Optional[DatasetInfo] = None,  lowerCamelCase						:    Optional[str] = None,  **lowerCamelCase						:    str,  ):

								'''simple docstring'''
								super().__init__(self,  **lowerCamelCase  )
								lowercase__    =			repo_info
								lowercase__    =			token
								lowercase__    =			None




				def 				lowercase__   (     self						:    Optional[Any]  ):

								'''simple docstring'''
								if self.dir_cache is None:
												lowercase__    =			{}
												for hf_file in self.repo_info.siblings:
																# TODO(QL): add sizes
																lowercase__    =			{
																    '''name''': hf_file.rfilename,
																    '''size''': None,
																    '''type''': '''file''',
																}
																self.dir_cache.update(
																    {
																        str(lowerCamelCase  ): {'''name''': str(lowerCamelCase  ), '''size''': None, '''type''': '''directory'''}
																        for d in list(PurePosixPath(hf_file.rfilename  ).parents  )[:-1]
																    }  )




				def 				lowercase__   (     self						:    Any,  lowerCamelCase						:    str,  lowerCamelCase						:    str = "rb",  **lowerCamelCase						:    Any,  ):

								'''simple docstring'''
								if not isinstance(self.repo_info,  lowerCamelCase  ):
												raise NotImplementedError(F"""Open is only implemented for dataset repositories, but got {self.repo_info}"""  )
								lowercase__    =			hf_hub_url(self.repo_info.id,  lowerCamelCase,  revision=self.repo_info.sha  )
								return fsspec.open(
								    lowerCamelCase,  mode=lowerCamelCase,  headers=get_authentication_headers_for_url(lowerCamelCase,  use_auth_token=self.token  ),  client_kwargs={'''trust_env''': True},  ).open()




				def 				lowercase__   (     self						:    Dict,  lowerCamelCase						:    Any,  **lowerCamelCase						:    int  ):

								'''simple docstring'''
								self._get_dirs()
								lowercase__    =			self._strip_protocol(lowerCamelCase  )
								if path in self.dir_cache:
												return self.dir_cache[path]
								else:
												raise FileNotFoundError(lowerCamelCase  )




				def 				lowercase__   (     self						:    Optional[Any],  lowerCamelCase						:    List[str],  lowerCamelCase						:    Optional[int]=False,  **lowerCamelCase						:    str  ):

								'''simple docstring'''
								self._get_dirs()
								lowercase__    =			PurePosixPath(path.strip('''/'''  )  )
								lowercase__    =			{}
								for p, f in self.dir_cache.items():
												lowercase__    =			PurePosixPath(p.strip('''/'''  )  )
												lowercase__    =			p.parent
												if root == path:
																lowercase__    =			f
								lowercase__    =			list(paths.values()  )
								if detail:
												return out
								else:
												return sorted(f['''name'''] for f in out  )




