


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       	=  [
    '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       	=  [
    'mlp.dense_4h_to_h.weight',
    'self_attention.dense.weight',
]


def       UpperCAmelCase_					(   UpperCAmelCase__    ,				UpperCAmelCase__	):
					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*).*"""    ,				SCREAMING_SNAKE_CASE_	)[1]	)
					layer_number -= 3
					return F'''h.{layer_number}.''' + key


def       UpperCAmelCase_					(   UpperCAmelCase__	):





					if dtype == torch.bool:
										return 1 / 8
					lowercase_					=      re.search(r"""[^\d](\d+)$"""    ,				str(SCREAMING_SNAKE_CASE_	)	)
					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       UpperCAmelCase_					(   UpperCAmelCase__    ,				UpperCAmelCase__    ,				UpperCAmelCase__    ,				UpperCAmelCase__    ,				UpperCAmelCase__	):





					if bloom_config_file == "":
										lowercase_					=      BloomConfig()
					else:
										lowercase_					=      BloomConfig.from_json_file(SCREAMING_SNAKE_CASE_	)

					if shard_model:
										lowercase_					=      os.listdir(SCREAMING_SNAKE_CASE_	)
										lowercase_					=      sorted(filter(lambda UpperCAmelCase__	: s.startswith("""layer"""	) and "model_00" in s    ,				SCREAMING_SNAKE_CASE_	)	)

										lowercase_					=      {"""weight_map""": {}, """metadata""": {}}
										lowercase_					=      0

										lowercase_					=      None

										lowercase_					=      BloomConfig()

										for j, file in enumerate(SCREAMING_SNAKE_CASE_	):
															print("""Processing file: {}""".format(SCREAMING_SNAKE_CASE_	)	)
															lowercase_					=      None

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

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

																				if tensors is None:
																									lowercase_					=      temp
																				else:
																									for key in tensors.keys():
																														if any(key.endswith(SCREAMING_SNAKE_CASE_	) 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=SCREAMING_SNAKE_CASE_	)

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

										lowercase_					=      BloomConfig()
										lowercase_					=      pytorch_dump_folder_path + """/""" + CONFIG_NAME
										lowercase_					=      total_size
										with open(SCREAMING_SNAKE_CASE_    ,				"""w"""    ,				encoding="""utf-8"""	) as f:
															f.write(config.to_json_string()	)
										with open(os.path.join(SCREAMING_SNAKE_CASE_    ,				WEIGHTS_NAME + """.index.json"""	)    ,				"""w"""    ,				encoding="""utf-8"""	) as f:
															lowercase_					=      json.dumps(SCREAMING_SNAKE_CASE_    ,				indent=2    ,				sort_keys=SCREAMING_SNAKE_CASE_	) + """\n"""
															f.write(SCREAMING_SNAKE_CASE_	)
					else:
										lowercase_					=      BloomModel(SCREAMING_SNAKE_CASE_	)

										lowercase_					=      os.listdir(SCREAMING_SNAKE_CASE_	)
										lowercase_					=      sorted(filter(lambda UpperCAmelCase__	: s.startswith("""layer"""	) and "model_00" in s    ,				SCREAMING_SNAKE_CASE_	)	)

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

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

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

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

															lowercase_					=      model.load_state_dict(SCREAMING_SNAKE_CASE_    ,				strict=SCREAMING_SNAKE_CASE_	)
															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(SCREAMING_SNAKE_CASE_    ,				exist_ok=SCREAMING_SNAKE_CASE_	)
										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()    ,				SCREAMING_SNAKE_CASE_	)
										print(F'''Save configuration file to {pytorch_config_dump_path}'''	)
										with open(SCREAMING_SNAKE_CASE_    ,				"""w"""    ,				encoding="""utf-8"""	) as f:
															f.write(config.to_json_string()	)


if __name__ == "__main__":
							a       	=  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       	=  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'''


from sklearn.metrics import mean_squared_error

import datasets


_SCREAMING_SNAKE_CASE									=				"\\n@article{scikit-learn,\n title={Scikit-learn: Machine Learning in {P}ython},\n author={Pedregosa, F. and Varoquaux, G. and Gramfort, A. and Michel, V.\n         and Thirion, B. and Grisel, O. and Blondel, M. and Prettenhofer, P.\n         and Weiss, R. and Dubourg, V. and Vanderplas, J. and Passos, A. and\n         Cournapeau, D. and Brucher, M. and Perrot, M. and Duchesnay, E.},\n journal={Journal of Machine Learning Research},\n volume={12},\n pages={2825--2830},\n year={2011}\n}\n"

_SCREAMING_SNAKE_CASE									=				"\\nMean Squared Error(MSE) is the average of the square of difference between the predicted\nand actual values.\n"


_SCREAMING_SNAKE_CASE									=				"\nArgs:\n    predictions: array-like of shape (n_samples,) or (n_samples, n_outputs)\n        Estimated target values.\n    references: array-like of shape (n_samples,) or (n_samples, n_outputs)\n        Ground truth (correct) target values.\n    sample_weight: array-like of shape (n_samples,), default=None\n        Sample weights.\n    multioutput: {\"raw_values\", \"uniform_average\"} or array-like of shape (n_outputs,), default=\"uniform_average\"\n        Defines aggregating of multiple output values. Array-like value defines weights used to average errors.\n\n        \"raw_values\" : Returns a full set of errors in case of multioutput input.\n\n        \"uniform_average\" : Errors of all outputs are averaged with uniform weight.\n\n    squared : bool, default=True\n        If True returns MSE value, if False returns RMSE (Root Mean Squared Error) value.\n\nReturns:\n    mse : mean squared error.\nExamples:\n\n    >>> mse_metric = datasets.load_metric(\"mse\")\n    >>> predictions = [2.5, 0.0, 2, 8]\n    >>> references = [3, -0.5, 2, 7]\n    >>> results = mse_metric.compute(predictions=predictions, references=references)\n    >>> print(results)\n    {'mse': 0.375}\n    >>> rmse_result = mse_metric.compute(predictions=predictions, references=references, squared=False)\n    >>> print(rmse_result)\n    {'mse': 0.6123724356957945}\n\n    If you're using multi-dimensional lists, then set the config as follows :\n\n    >>> mse_metric = datasets.load_metric(\"mse\", \"multilist\")\n    >>> predictions = [[0.5, 1], [-1, 1], [7, -6]]\n    >>> references = [[0, 2], [-1, 2], [8, -5]]\n    >>> results = mse_metric.compute(predictions=predictions, references=references)\n    >>> print(results)\n    {'mse': 0.7083333333333334}\n    >>> results = mse_metric.compute(predictions=predictions, references=references, multioutput='raw_values')\n    >>> print(results) # doctest: +NORMALIZE_WHITESPACE\n    {'mse': array([0.41666667, 1. ])}\n"

@datasets.utils.file_utils.add_start_docstrings(_DESCRIPTION    ,_KWARGS_DESCRIPTION		)
class  lowerCAmelCase_		(		datasets.Metric		):




 def 				_snake_case (						self		)     ->				Dict:
     return datasets.MetricInfo(
         description=_DESCRIPTION      ,  citation=_CITATION      ,  inputs_description=_KWARGS_DESCRIPTION      ,  features=datasets.Features(self._get_feature_types()		)      ,  reference_urls=[
             "https://scikit-learn.org/stable/modules/generated/sklearn.metrics.mean_squared_error.html"
         ]      ,  )




 def 				_snake_case (						self		)     ->				Tuple:
     if self.config_name == "multilist":
         return {
             "predictions": datasets.Sequence(datasets.Value("float"		)		),
             "references": datasets.Sequence(datasets.Value("float"		)		),
         }
     else:
         return {
             "predictions": datasets.Value("float"		),
             "references": datasets.Value("float"		),
         }





 def 				_snake_case (						self      ,  _lowerCAmelCase      ,  _lowerCAmelCase      ,  _lowerCAmelCase=None      ,  _lowerCAmelCase="uniform_average"      ,  _lowerCAmelCase=True		)     ->				Union[str, Any]:
     _lowerCAmelCase						      =   mean_squared_error(
         _lowerCAmelCase      ,  _lowerCAmelCase      ,  sample_weight=_lowerCAmelCase      ,  multioutput=_lowerCAmelCase      ,  squared=_lowerCAmelCase		)

     return {"mse": mse}


