





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

__lowercase       :Tuple       =  [
    "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",
]

__lowercase       :str       =  [
    "mlp.dense_4h_to_h.weight",
    "self_attention.dense.weight",
]







def       UpperCAmelCase							(     _lowerCamelCase      :  Optional[int]			,   _lowerCamelCase      :  Optional[int]     ):





						'''simple docstring'''


						SCREAMING_SNAKE_CASE__     :				List[str]   =					{
						    """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
						SCREAMING_SNAKE_CASE__     :				Dict   =					int(re.match(r".*layer_(\d*).*"			,   UpperCAmelCase__     )[1]     )
						layer_number -= 3
						return f"""h.{layer_number}.""" + key







def       UpperCAmelCase							(     _lowerCamelCase      :  Any     ):





						'''simple docstring'''


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







def       UpperCAmelCase							(     _lowerCamelCase      :  Dict			,   _lowerCamelCase      :  Optional[Any]			,   _lowerCamelCase      :  Optional[int]			,   _lowerCamelCase      :  Union[str, Any]			,   _lowerCamelCase      :  Optional[int]     ):





						'''simple docstring'''


						if bloom_config_file == "":
												SCREAMING_SNAKE_CASE__     :				Union[str, Any]   =					BloomConfig()
						else:
												SCREAMING_SNAKE_CASE__     :				int   =					BloomConfig.from_json_file(UpperCAmelCase__     )

						if shard_model:
												SCREAMING_SNAKE_CASE__     :				Optional[Any]   =					os.listdir(UpperCAmelCase__     )
												SCREAMING_SNAKE_CASE__     :				Union[str, Any]   =					sorted(filter(lambda _lowerCamelCase     : s.startswith("layer"     ) and "model_00" in s			,   UpperCAmelCase__     )     )

												SCREAMING_SNAKE_CASE__     :				Union[str, Any]   =					{"""weight_map""": {}, """metadata""": {}}
												SCREAMING_SNAKE_CASE__     :				Optional[int]   =					0

												SCREAMING_SNAKE_CASE__     :				int   =					None

												SCREAMING_SNAKE_CASE__     :				str   =					BloomConfig()

												for j, file in enumerate(UpperCAmelCase__     ):
																		print("Processing file: {}".format(UpperCAmelCase__     )     )
																		SCREAMING_SNAKE_CASE__     :				Any   =					None

																		for i in range(UpperCAmelCase__     ):
																								# load all TP files
																								SCREAMING_SNAKE_CASE__     :				Any   =					file.replace("model_00"			,   f"""model_0{i}"""     )
																								SCREAMING_SNAKE_CASE__     :				Any   =					torch.load(os.path.join(UpperCAmelCase__			,   UpperCAmelCase__     )			,   map_location="cpu"     )

																								# Rename keys in the transformers names
																								SCREAMING_SNAKE_CASE__     :				int   =					list(temp.keys()     )
																								for key in keys:
																														SCREAMING_SNAKE_CASE__     :				Any   =					temp.pop(UpperCAmelCase__     )

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

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

												SCREAMING_SNAKE_CASE__     :				int   =					os.listdir(UpperCAmelCase__     )
												SCREAMING_SNAKE_CASE__     :				Optional[int]   =					sorted(filter(lambda _lowerCamelCase     : s.startswith("layer"     ) and "model_00" in s			,   UpperCAmelCase__     )     )

												SCREAMING_SNAKE_CASE__     :				Union[str, Any]   =					None
												for i, file in enumerate(UpperCAmelCase__     ):
																		SCREAMING_SNAKE_CASE__     :				Tuple   =					None
																		for i in range(UpperCAmelCase__     ):
																								# load all TP files
																								SCREAMING_SNAKE_CASE__     :				List[str]   =					file.replace("model_00"			,   f"""model_0{i}"""     )
																								SCREAMING_SNAKE_CASE__     :				List[Any]   =					torch.load(os.path.join(UpperCAmelCase__			,   UpperCAmelCase__     )			,   map_location="cpu"     )

																								# Rename keys in the transformers names
																								SCREAMING_SNAKE_CASE__     :				int   =					list(temp.keys()     )
																								for key in keys:
																														SCREAMING_SNAKE_CASE__     :				int   =					temp.pop(UpperCAmelCase__     )

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

																		SCREAMING_SNAKE_CASE__     :				Optional[int]   =					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:
																								SCREAMING_SNAKE_CASE__     :				Union[str, Any]   =					set(other_keys.missing_keys     )
																		else:
																								SCREAMING_SNAKE_CASE__     :				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(UpperCAmelCase__			,   exist_ok=UpperCAmelCase__     )
												SCREAMING_SNAKE_CASE__     :				Union[str, Any]   =					pytorch_dump_folder_path + """/""" + WEIGHTS_NAME
												SCREAMING_SNAKE_CASE__     :				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:
																		SCREAMING_SNAKE_CASE__     :				Union[str, Any]   =					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__":
				__lowercase       :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",
				)
				__lowercase       :List[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,
				)






import numpy



class    _a     :






							"""simple docstring"""

							def __init__(      self				:     Optional[int]							,	a				:     numpy.ndarray							,	a				:     numpy.ndarray    )	->None:
													SCREAMING_SNAKE_CASE__     :				Any   =					input_array

													# Random initial weights are assigned where first argument is the
													# number of nodes in previous layer and second argument is the
													# number of nodes in the next layer.

													# Random initial weights are assigned.
													# self.input_array.shape[1] is used to represent number of nodes in input layer.
													# First hidden layer consists of 4 nodes.
													SCREAMING_SNAKE_CASE__     :				int   =					numpy.random.rand(
													    self.input_array.shape[1]							,	4    )

													# Random initial values for the first hidden layer.
													# First hidden layer has 4 nodes.
													# Second hidden layer has 3 nodes.
													SCREAMING_SNAKE_CASE__     :				Dict   =					numpy.random.rand(
													    4							,	3    )

													# Random initial values for the second hidden layer.
													# Second hidden layer has 3 nodes.
													# Output layer has 1 node.
													SCREAMING_SNAKE_CASE__     :				List[Any]   =					numpy.random.rand(3							,	1    )

													# Real output values provided.
													SCREAMING_SNAKE_CASE__     :				str   =					output_array

													# Predicted output values by the neural network.
													# Predicted_output array initially consists of zeroes.
													SCREAMING_SNAKE_CASE__     :				Tuple   =					numpy.zeros(output_array.shape    )

							def      A_  (      self				:     Union[str, Any]    )	->numpy.ndarray:
													SCREAMING_SNAKE_CASE__     :				List[Any]   =					sigmoid(
													    numpy.dot(self.input_array							,	self.input_layer_and_first_hidden_layer_weights    )    )

													# layer_between_first_hidden_layer_and_second_hidden_layer is the layer
													# connecting the first hidden set of nodes with the second hidden set of nodes.
													SCREAMING_SNAKE_CASE__     :				Optional[int]   =					sigmoid(
													    numpy.dot(
													        self.layer_between_input_and_first_hidden_layer							,	self.first_hidden_layer_and_second_hidden_layer_weights							,	)    )

													# layer_between_second_hidden_layer_and_output is the layer connecting
													# second hidden layer with the output node.
													SCREAMING_SNAKE_CASE__     :				int   =					sigmoid(
													    numpy.dot(
													        self.layer_between_first_hidden_layer_and_second_hidden_layer							,	self.second_hidden_layer_and_output_layer_weights							,	)    )

													return self.layer_between_second_hidden_layer_and_output

							def      A_  (      self				:     int    )	->None:
													SCREAMING_SNAKE_CASE__     :				Optional[int]   =					numpy.dot(
													    self.layer_between_first_hidden_layer_and_second_hidden_layer.T							,	2
													    * (self.output_array - self.predicted_output)
													    * sigmoid_derivative(self.predicted_output    )							,	)
													SCREAMING_SNAKE_CASE__     :				Union[str, Any]   =					numpy.dot(
													    self.layer_between_input_and_first_hidden_layer.T							,	numpy.dot(
													        2
													        * (self.output_array - self.predicted_output)
													        * sigmoid_derivative(self.predicted_output    )							,	self.second_hidden_layer_and_output_layer_weights.T							,	)
													    * sigmoid_derivative(
													        self.layer_between_first_hidden_layer_and_second_hidden_layer    )							,	)
													SCREAMING_SNAKE_CASE__     :				int   =					numpy.dot(
													    self.input_array.T							,	numpy.dot(
													        numpy.dot(
													            2
													            * (self.output_array - self.predicted_output)
													            * sigmoid_derivative(self.predicted_output    )							,	self.second_hidden_layer_and_output_layer_weights.T							,	)
													        * sigmoid_derivative(
													            self.layer_between_first_hidden_layer_and_second_hidden_layer    )							,	self.first_hidden_layer_and_second_hidden_layer_weights.T							,	)
													    * sigmoid_derivative(self.layer_between_input_and_first_hidden_layer    )							,	)

													self.input_layer_and_first_hidden_layer_weights += (
													    updated_input_layer_and_first_hidden_layer_weights
													)
													self.first_hidden_layer_and_second_hidden_layer_weights += (
													    updated_first_hidden_layer_and_second_hidden_layer_weights
													)
													self.second_hidden_layer_and_output_layer_weights += (
													    updated_second_hidden_layer_and_output_layer_weights
													)

							def      A_  (      self				:     int							,	a				:     numpy.ndarray							,	a				:     int							,	a				:     bool    )	->None:
													for iteration in range(1							,	iterations + 1    ):
																			SCREAMING_SNAKE_CASE__     :				Dict   =					self.feedforward()
																			self.back_propagation()
																			if give_loss:
																									SCREAMING_SNAKE_CASE__     :				int   =					numpy.mean(numpy.square(output - self.feedforward()    )    )
																									print(f"""Iteration {iteration} Loss: {loss}"""    )






							def      A_  (      self				:     Tuple							,	a				:     numpy.ndarray    )	->int:
													SCREAMING_SNAKE_CASE__     :				Optional[int]   =					input_arr

													SCREAMING_SNAKE_CASE__     :				Dict   =					sigmoid(
													    numpy.dot(self.array							,	self.input_layer_and_first_hidden_layer_weights    )    )

													SCREAMING_SNAKE_CASE__     :				Any   =					sigmoid(
													    numpy.dot(
													        self.layer_between_input_and_first_hidden_layer							,	self.first_hidden_layer_and_second_hidden_layer_weights							,	)    )

													SCREAMING_SNAKE_CASE__     :				Optional[int]   =					sigmoid(
													    numpy.dot(
													        self.layer_between_first_hidden_layer_and_second_hidden_layer							,	self.second_hidden_layer_and_output_layer_weights							,	)    )

													return int(self.layer_between_second_hidden_layer_and_output > 0.6    )







def       UpperCAmelCase							(     _lowerCamelCase      :  numpy.ndarray     ):





						'''simple docstring'''


						return 1 / (1 + numpy.exp(-value     ))







def       UpperCAmelCase							(     _lowerCamelCase      :  numpy.ndarray     ):





						'''simple docstring'''


						return (value) * (1 - (value))







def       UpperCAmelCase							(     ):





						'''simple docstring'''


						SCREAMING_SNAKE_CASE__     :				str   =					numpy.array(
						    (
						        [0, 0, 0],
						        [0, 0, 1],
						        [0, 1, 0],
						        [0, 1, 1],
						        [1, 0, 0],
						        [1, 0, 1],
						        [1, 1, 0],
						        [1, 1, 1],
						    )			,   dtype=numpy.floataa			,   )

						# True output values for the given input values.
						SCREAMING_SNAKE_CASE__     :				Any   =					numpy.array(([0], [1], [1], [0], [1], [0], [0], [1])			,   dtype=numpy.floataa     )

						# Calling neural network class.
						SCREAMING_SNAKE_CASE__     :				List[Any]   =					TwoHiddenLayerNeuralNetwork(
						    input_array=_lowerCamelCase			,   output_array=_lowerCamelCase     )

						# Calling training function.
						# Set give_loss to True if you want to see loss in every iteration.
						neural_network.train(output=_lowerCamelCase			,   iterations=10			,   give_loss=_lowerCamelCase     )

						return neural_network.predict(numpy.array(([1, 1, 1])			,   dtype=numpy.floataa     )     )


if __name__ == "__main__":
				example()