



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 			_lowerCAmelCase       (					__lowerCAmelCase  ,  __lowerCAmelCase      )   ->					Optional[Any]:


				"""simple docstring"""

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



def 			_lowerCAmelCase       (					__lowerCAmelCase      )   ->					str:


				"""simple docstring"""

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



def 			_lowerCAmelCase       (					__lowerCAmelCase  ,  __lowerCAmelCase  ,  __lowerCAmelCase  ,  __lowerCAmelCase  ,  __lowerCAmelCase      )   ->					List[Any]:


				"""simple docstring"""

				if bloom_config_file == "":
								snake_case__	:      Optional[Any]     		=	BloomConfig()
				else:
								snake_case__	:      Tuple     		=	BloomConfig.from_json_file(lowerCamelCase_      )

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

								snake_case__	:      Dict     		=	{'''weight_map''': {}, '''metadata''': {}}
								snake_case__	:      Dict     		=	0

								snake_case__	:      List[Any]     		=	None

								snake_case__	:      Tuple     		=	BloomConfig()

								for j, file in enumerate(lowerCamelCase_      ):
												print('''Processing file: {}'''.format(lowerCamelCase_      )      )
												snake_case__	:      List[str]     		=	None

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

																# Rename keys in the transformers names
																snake_case__	:      int     		=	list(temp.keys()      )
																for key in keys:
																				snake_case__	:      Optional[Any]     		=	temp.pop(lowerCamelCase_      )

																if tensors is None:
																				snake_case__	:      Tuple     		=	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
																												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__	:      str     		=	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      ):
																				snake_case__	:      Tuple     		=	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():
																snake_case__	:      Dict     		=	tensors[key]
																total_size += value.numel() * get_dtype_size(value.dtype      )
																if key not in index_dict["weight_map"]:
																				snake_case__	:      List[str]     		=	'''pytorch_model_{}-of-{}.bin'''.format(
																				    str(j + 1      ).zfill(5      )  ,  str(len(lowerCamelCase_      )      ).zfill(5      )      )

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

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

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

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

																if tensors is None:
																				snake_case__	:      Optional[Any]     		=	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
																												snake_case__	:      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__	:      Union[str, Any]     		=	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      ):
																				snake_case__	:      Optional[int]     		=	tensors[key] / pretraining_tp

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









from math import asin, atan, cos, radians, sin, sqrt, tan

A__	:			Optional[int]			   = 637_8137.0
A__	:			List[str]			   = 635_6752.31_4245
A__	:			Union[str, Any]			   = 6_37_81_37






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






				'''simple docstring'''
				lowercase__    =			(AXIS_A - AXIS_B) / AXIS_A
				lowercase__    =			atan((1 - flattening) * tan(radians(lowerCamelCase_				)				)				)
				lowercase__    =			atan((1 - flattening) * tan(radians(lowerCamelCase_				)				)				)
				lowercase__    =			radians(lowerCamelCase_				)
				lowercase__    =			radians(lowerCamelCase_				)
				# Equation
				lowercase__    =			sin((phi_a - phi_a) / 2				)
				lowercase__    =			sin((lambda_a - lambda_a) / 2				)
				# Square both values
				sin_sq_phi *= sin_sq_phi
				sin_sq_lambda *= sin_sq_lambda
				lowercase__    =			sqrt(sin_sq_phi + (cos(lowerCamelCase_				) * cos(lowerCamelCase_				) * sin_sq_lambda)				)
				return 2 * RADIUS * asin(lowerCamelCase_				)


if __name__ == "__main__":
							import doctest

							doctest.testmod()




