






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

__lowerCamelCase   :					Optional[int]   							=					[
    '''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''',
]

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

def 	__UpperCAmelCase       ( __magic_name__     ,__magic_name__						)->    Optional[int]:



			"""simple docstring"""



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

def 	__UpperCAmelCase       ( __magic_name__						)->    Dict:



			"""simple docstring"""



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

def 	__UpperCAmelCase       ( __magic_name__     ,__magic_name__     ,__magic_name__     ,__magic_name__     ,__magic_name__						)->    Tuple:



			"""simple docstring"""



			if bloom_config_file == "":
						snake_case_				:					str   =       BloomConfig()
			else:
						snake_case_				:					Tuple   =       BloomConfig.from_json_file(__magic_name__						)

			if shard_model:
						snake_case_				:					str   =       os.listdir(__magic_name__						)
						snake_case_				:					List[str]   =       sorted(filter(lambda __magic_name__						: s.startswith("layer"						) and "model_00" in s     ,__magic_name__						)						)

						snake_case_				:					List[Any]   =       {"weight_map": {}, "metadata": {}}
						snake_case_				:					Union[str, Any]   =       0

						snake_case_				:					Tuple   =       None

						snake_case_				:					Optional[Any]   =       BloomConfig()

						for j, file in enumerate(__magic_name__						):
									print("Processing file: {}".format(__magic_name__						)						)
									snake_case_				:					List[Any]   =       None

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

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

												if tensors is None:
															snake_case_				:					Union[str, Any]   =       temp
												else:
															for key in tensors.keys():
																		if any(key.endswith(__magic_name__						) 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
																					snake_case_				:					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
																					snake_case_				:					List[Any]   =       torch.cat([tensors[key], temp[key]]     ,dim=__magic_name__						)

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

									for key in tensors.keys():
												snake_case_				:					Any   =       tensors[key]
												total_size += value.numel() * get_dtype_size(value.dtype						)
												if key not in index_dict["weight_map"]:
															snake_case_				:					int   =       "pytorch_model_{}-of-{}.bin".format(
															    str(j + 1						).zfill(5						)     ,str(len(__magic_name__						)						).zfill(5						)						)

						snake_case_				:					int   =       BloomConfig()
						snake_case_				:					Dict   =       pytorch_dump_folder_path + "/" + CONFIG_NAME
						snake_case_				:					Optional[Any]   =       total_size
						with open(__magic_name__     ,"w"     ,encoding="utf-8"						) as f:
									f.write(config.to_json_string()						)
						with open(os.path.join(__magic_name__     ,WEIGHTS_NAME + ".index.json"						)     ,"w"     ,encoding="utf-8"						) as f:
									snake_case_				:					Any   =       json.dumps(__magic_name__     ,indent=2     ,sort_keys=__magic_name__						) + "\n"
									f.write(__magic_name__						)
			else:
						snake_case_				:					Optional[Any]   =       BloomModel(__magic_name__						)

						snake_case_				:					str   =       os.listdir(__magic_name__						)
						snake_case_				:					Optional[int]   =       sorted(filter(lambda __magic_name__						: s.startswith("layer"						) and "model_00" in s     ,__magic_name__						)						)

						snake_case_				:					List[Any]   =       None
						for i, file in enumerate(__magic_name__						):
									snake_case_				:					Tuple   =       None
									for i in range(__magic_name__						):
												# load all TP files
												snake_case_				:					Tuple   =       file.replace("model_00"     ,F'''model_0{i}'''						)
												snake_case_				:					Dict   =       torch.load(os.path.join(__magic_name__     ,__magic_name__						)     ,map_location="cpu"						)

												# Rename keys in the transformers names
												snake_case_				:					Tuple   =       list(temp.keys()						)
												for key in keys:
															snake_case_				:					Tuple   =       temp.pop(__magic_name__						)

												if tensors is None:
															snake_case_				:					str   =       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(__magic_name__						) for end in WEIGHTS_TO_AVERAGE_ENDSWITH						):
																					tensors[key] += temp[key]
																		else:
																					# Some weights are RowParallelLinear in Megatron-Deepspeed, others are ColumnParallel
																					snake_case_				:					Tuple   =       1 if any(text in key for text in WEIGHTS_WITH_ROW_PARALLELISM_CONTAIN						) else 0
																					# We concatenate these weights accross TP ranks
																					snake_case_				:					Union[str, Any]   =       torch.cat([tensors[key], temp[key]]     ,dim=__magic_name__						)

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

									snake_case_				:					Optional[Any]   =       model.load_state_dict(__magic_name__     ,strict=__magic_name__						)
									assert not other_keys.unexpected_keys, F'''The keys {other_keys.unexpected_keys} are unexpected'''
									if missing_keys is None:
												snake_case_				:					Dict   =       set(other_keys.missing_keys						)
									else:
												snake_case_				:					Optional[int]   =       missing_keys.intersection(set(other_keys.missing_keys						)						)

						assert not missing_keys, F'''The keys {missing_keys} are missing'''

						# Save pytorch-model
						os.makedirs(__magic_name__     ,exist_ok=__magic_name__						)
						snake_case_				:					Dict   =       pytorch_dump_folder_path + "/" + WEIGHTS_NAME
						snake_case_				:					Optional[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:
									snake_case_				:					int   =       model.to(config.torch_dtype						)
						torch.save(model.state_dict()     ,__magic_name__						)
						print(F'''Save configuration file to {pytorch_config_dump_path}'''						)
						with open(__magic_name__     ,"w"     ,encoding="utf-8"						) as f:
									f.write(config.to_json_string()						)


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




import argparse
import json
import os

from tensorflow.core.protobuf.saved_model_pba import SavedModel


# All paths are set with the intent you should run this script from the root of the repo with the command
# python utils/check_copies.py
__lowerCamelCase   :					str   							=					'''.'''

# Internal TensorFlow ops that can be safely ignored (mostly specific to a saved model)
__lowerCamelCase   :					Tuple   							=					[
    '''Assert''',
    '''AssignVariableOp''',
    '''EmptyTensorList''',
    '''MergeV2Checkpoints''',
    '''ReadVariableOp''',
    '''ResourceGather''',
    '''RestoreV2''',
    '''SaveV2''',
    '''ShardedFilename''',
    '''StatefulPartitionedCall''',
    '''StaticRegexFullMatch''',
    '''VarHandleOp''',
]

def 	__UpperCAmelCase       ( __magic_name__     ,__magic_name__     ,__magic_name__						)->    List[str]:



			"""simple docstring"""



			snake_case_				:					Tuple   =       SavedModel()
			snake_case_				:					Dict   =       []

			with open(os.path.join(__magic_name__     ,"utils"     ,"tf_ops"     ,"onnx.json"						)						) as f:
						snake_case_				:					Dict   =       json.load(__magic_name__						)["opsets"]

			for i in range(1     ,opset + 1						):
						onnx_ops.extend(onnx_opsets[str(__magic_name__						)]						)

			with open(__magic_name__     ,"rb"						) as f:
						saved_model.ParseFromString(f.read()						)

			snake_case_				:					Tuple   =       set()

			# Iterate over every metagraph in case there is more than one (a saved model can contain multiple graphs)
			for meta_graph in saved_model.meta_graphs:
						# Add operations in the graph definition
						model_op_names.update(node.op for node in meta_graph.graph_def.node						)

						# Go through the functions in the graph definition
						for func in meta_graph.graph_def.library.function:
									# Add operations in each function
									model_op_names.update(node.op for node in func.node_def						)

    # Convert to list, sorted if you want
			snake_case_				:					str   =       sorted(__magic_name__						)
			snake_case_				:					Optional[int]   =       []

			for op in model_op_names:
						if op not in onnx_ops and op not in INTERNAL_OPS:
									incompatible_ops.append(__magic_name__						)

			if strict and len(__magic_name__						) > 0:
						raise Exception(F'''Found the following incompatible ops for the opset {opset}:\n''' + incompatible_ops						)
			elif len(__magic_name__						) > 0:
						print(F'''Found the following incompatible ops for the opset {opset}:'''						)
						print(*__magic_name__     ,sep="\n"						)
			else:
						print(F'''The saved model {saved_model_path} can properly be converted with ONNX.'''						)


if __name__ == "__main__":
		__lowerCamelCase   :					Any   							=					argparse.ArgumentParser()
		parser.add_argument('''--saved_model_path''', help='''Path of the saved model to check (the .pb file).''')
		parser.add_argument(
		    '''--opset''', default=12, type=int, help='''The ONNX opset against which the model has to be tested.'''
		)
		parser.add_argument(
		    '''--framework''', choices=['''onnx'''], default='''onnx''', help='''Frameworks against which to test the saved model.'''
		)
		parser.add_argument(
		    '''--strict''', action='''store_true''', help='''Whether make the checking strict (raise errors) or not (raise warnings)'''
		)
		__lowerCamelCase   :					Dict   							=					parser.parse_args()

		if args.framework == "onnx":
				onnx_compliancy(args.saved_model_path, args.strict, args.opset)




