


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






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






def       lowerCamelCase					(a_						:Tuple				,				a_						:Optional[int]				,				a_						:Dict				,				a_						:List[Any]				,				a_						:int)     ->							Tuple:
							# Construct model
							if bloom_config_file == "":
														lowercase		:str       							=				BloomConfig()
							else:
														lowercase		:List[str]       							=				BloomConfig.from_json_file(a_)

							if shard_model:
														lowercase		:Any       							=				os.listdir(a_)
														lowercase		:Optional[int]       							=				sorted(filter(lambda a_: s.startswith('''layer''') and "model_00" in s				,				a_))

														lowercase		:Optional[Any]       							=				{'''weight_map''': {}, '''metadata''': {}}
														lowercase		:Dict       							=				0

														lowercase		:List[Any]       							=				None

														lowercase		:List[str]       							=				BloomConfig()

														for j, file in enumerate(a_):
																					print('''Processing file: {}'''.format(a_))
																					lowercase		:Union[str, Any]       							=				None

																					for i in range(a_):
																												# load all TP files
																												lowercase		:Optional[int]       							=				file.replace('''model_00'''				,				F"""model_0{i}""")
																												lowercase		:Union[str, Any]       							=				torch.load(os.path.join(a_				,				a_)				,				map_location='''cpu''')

																												# Rename keys in the transformers names
																												lowercase		:Any       							=				list(temp.keys())
																												for key in keys:
																																			lowercase		:Dict       							=				temp.pop(a_)

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

														lowercase		:List[Any]       							=				BloomConfig()
														lowercase		:List[Any]       							=				pytorch_dump_folder_path + '''/''' + CONFIG_NAME
														lowercase		:Optional[int]       							=				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:
																					lowercase		:List[str]       							=				json.dumps(a_				,				indent=2				,				sort_keys=a_) + '''\n'''
																					f.write(a_)
							else:
														lowercase		:int       							=				BloomModel(a_)

														lowercase		:Dict       							=				os.listdir(a_)
														lowercase		:str       							=				sorted(filter(lambda a_: s.startswith('''layer''') and "model_00" in s				,				a_))

														lowercase		:Optional[int]       							=				None
														for i, file in enumerate(a_):
																					lowercase		:Dict       							=				None
																					for i in range(a_):
																												# load all TP files
																												lowercase		:Union[str, Any]       							=				file.replace('''model_00'''				,				F"""model_0{i}""")
																												lowercase		:Dict       							=				torch.load(os.path.join(a_				,				a_)				,				map_location='''cpu''')

																												# Rename keys in the transformers names
																												lowercase		:str       							=				list(temp.keys())
																												for key in keys:
																																			lowercase		:Union[str, Any]       							=				temp.pop(a_)

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

																					lowercase		:Tuple       							=				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:
																												lowercase		:Any       							=				set(other_keys.missing_keys)
																					else:
																												lowercase		:List[str]       							=				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_)
														lowercase		:Dict       							=				pytorch_dump_folder_path + '''/''' + WEIGHTS_NAME
														lowercase		:Any       							=				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		:str       							=				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,
	)





"""simple docstring"""



# Copyright 2023 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import re

from ..utils import cached_file


# docstyle-ignore
UpperCAmelCase      		=			'''
Human: <<task>>

Assistant: '''


UpperCAmelCase      		=			'''huggingface-tools/default-prompts'''
UpperCAmelCase      		=			{'''chat''': '''chat_prompt_template.txt''', '''run''': '''run_prompt_template.txt'''}






def       lowerCamelCase					(a_						:int				,				a_						:str				,				a_						:Dict="run")     ->							Optional[Any]:
       if prompt_or_repo_id is None:
              lowercase		:Tuple       							=				DEFAULT_PROMPTS_REPO

       # prompt is considered a repo ID when it does not contain any kind of space
       if re.search('''\\s'''				,				a_) is not None:
              return prompt_or_repo_id

       lowercase		:List[str]       							=				cached_file(
           a_				,				PROMPT_FILES[mode]				,				repo_type='''dataset'''				,				user_agent={'''agent''': agent_name})
       with open(a_				,				'''r'''				,				encoding='''utf-8''') as f:
              return f.read()

