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



def 			__SCREAMING_SNAKE_CASE						(      __UpperCamelCase					: int				,  __UpperCamelCase					: List[str]			)							->      int:







		"""simple docstring"""





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



def 			__SCREAMING_SNAKE_CASE						(      __UpperCamelCase					: Optional[int]			)							->      List[Any]:







		"""simple docstring"""





		if dtype == torch.bool:
				return 1 / 8
		SCREAMING_SNAKE_CASE__    							=   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__    							=   int(bit_search.groups()[0]			)
		return bit_size // 8



def 			__SCREAMING_SNAKE_CASE						(      __UpperCamelCase					: Tuple				,  __UpperCamelCase					: Optional[int]				,  __UpperCamelCase					: Optional[int]				,  __UpperCamelCase					: int				,  __UpperCamelCase					: Optional[Any]			)							->      Dict:







		"""simple docstring"""





		if bloom_config_file == "":
				SCREAMING_SNAKE_CASE__    							=   BloomConfig()
		else:
				SCREAMING_SNAKE_CASE__    							=   BloomConfig.from_json_file(__UpperCamelCase			)

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

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

				SCREAMING_SNAKE_CASE__    							=   None

				SCREAMING_SNAKE_CASE__    							=   BloomConfig()

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

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

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

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

				SCREAMING_SNAKE_CASE__    							=   BloomConfig()
				SCREAMING_SNAKE_CASE__    							=   pytorch_dump_folder_path + """/""" + CONFIG_NAME
				SCREAMING_SNAKE_CASE__    							=   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__    							=   json.dumps(__UpperCamelCase				,  indent=2				,  sort_keys=__UpperCamelCase			) + """\n"""
						f.write(__UpperCamelCase			)
		else:
				SCREAMING_SNAKE_CASE__    							=   BloomModel(__UpperCamelCase			)

				SCREAMING_SNAKE_CASE__    							=   os.listdir(__UpperCamelCase			)
				SCREAMING_SNAKE_CASE__    							=   sorted(filter(lambda __UpperCamelCase			: s.startswith("""layer"""			) and "model_00" in s				,  __UpperCamelCase			)			)

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

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

								if tensors is None:
										SCREAMING_SNAKE_CASE__    							=   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__    							=   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__    							=   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__    							=   tensors[key] / pretraining_tp

						SCREAMING_SNAKE_CASE__    							=   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__    							=   set(other_keys.missing_keys			)
						else:
								SCREAMING_SNAKE_CASE__    							=   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__    							=   pytorch_dump_folder_path + """/""" + WEIGHTS_NAME
				SCREAMING_SNAKE_CASE__    							=   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__    							=   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__":
			__lowerCamelCase    :				str							     =	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    :				Dict							     =	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 ...configuration_utils import PretrainedConfig
from ...utils import logging


__lowerCamelCase    :				Optional[int]							     =	logging.get_logger(__name__)

__lowerCamelCase    :				str							     =	{
    '''facebook/nllb-moe-54B''': '''https://huggingface.co/facebook/nllb-moe-54b/resolve/main/config.json''',
}







class __snake_case       (		lowerCamelCase_	):
   lowerCAmelCase_														=							"nllb-moe"
   lowerCAmelCase_														=							["past_key_values"]
   lowerCAmelCase_														=							{"num_attention_heads": "encoder_attention_heads", "hidden_size": "d_model"}




   def __init__(					self    :       int ,			_lowercase    :       Tuple=12_81_12 ,			_lowercase    :       List[Any]=10_24 ,			_lowercase    :       Any=12 ,			_lowercase    :       List[Any]=40_96 ,			_lowercase    :       str=16 ,			_lowercase    :       str=12 ,			_lowercase    :       Optional[int]=40_96 ,			_lowercase    :       List[Any]=16 ,			_lowercase    :       str=0.05 ,			_lowercase    :       Tuple=0.05 ,			_lowercase    :       str=True ,			_lowercase    :       List[str]=True ,			_lowercase    :       Optional[Any]="relu" ,			_lowercase    :       str=10_24 ,			_lowercase    :       Tuple=0.1 ,			_lowercase    :       int=0.1 ,			_lowercase    :       Dict=0.0 ,			_lowercase    :       List[str]=0.02 ,			_lowercase    :       int=2 ,			_lowercase    :       Optional[Any]=True ,			_lowercase    :       List[Any]=False ,			_lowercase    :       List[str]="float32" ,			_lowercase    :       Optional[Any]=False ,			_lowercase    :       str=1_28 ,			_lowercase    :       int=64 ,			_lowercase    :       Optional[int]=4 ,			_lowercase    :       List[str]=4 ,			_lowercase    :       Union[str, Any]=0.0_01 ,			_lowercase    :       List[Any]=0.0_01 ,			_lowercase    :       List[str]="all" ,			_lowercase    :       Optional[Any]=False ,			_lowercase    :       int=False ,			_lowercase    :       Tuple=1.0 ,			_lowercase    :       Optional[int]=0.2 ,			_lowercase    :       Optional[int]=1 ,			_lowercase    :       List[Any]=0 ,			_lowercase    :       List[Any]=2 ,			_lowercase    :       int=False ,			**_lowercase    :       Union[str, Any] ,			):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__    							=   vocab_size
     SCREAMING_SNAKE_CASE__    							=   max_position_embeddings
     SCREAMING_SNAKE_CASE__    							=   d_model
     SCREAMING_SNAKE_CASE__    							=   encoder_ffn_dim
     SCREAMING_SNAKE_CASE__    							=   encoder_layers
     SCREAMING_SNAKE_CASE__    							=   encoder_attention_heads
     SCREAMING_SNAKE_CASE__    							=   decoder_ffn_dim
     SCREAMING_SNAKE_CASE__    							=   decoder_layers
     SCREAMING_SNAKE_CASE__    							=   decoder_attention_heads
     SCREAMING_SNAKE_CASE__    							=   dropout
     SCREAMING_SNAKE_CASE__    							=   attention_dropout
     SCREAMING_SNAKE_CASE__    							=   activation_dropout
     SCREAMING_SNAKE_CASE__    							=   activation_function
     SCREAMING_SNAKE_CASE__    							=   init_std
     SCREAMING_SNAKE_CASE__    							=   encoder_layerdrop
     SCREAMING_SNAKE_CASE__    							=   decoder_layerdrop
     SCREAMING_SNAKE_CASE__    							=   use_cache
     SCREAMING_SNAKE_CASE__    							=   encoder_layers
     SCREAMING_SNAKE_CASE__    							=   scale_embedding  # scale factor will be sqrt(d_model) if True
     SCREAMING_SNAKE_CASE__    							=   router_z_loss_coef
     SCREAMING_SNAKE_CASE__    							=   router_aux_loss_coef
     SCREAMING_SNAKE_CASE__    							=   decoder_sparse_step
     SCREAMING_SNAKE_CASE__    							=   encoder_sparse_step
     SCREAMING_SNAKE_CASE__    							=   num_experts
     SCREAMING_SNAKE_CASE__    							=   expert_capacity
     SCREAMING_SNAKE_CASE__    							=   router_bias
     if router_dtype not in ["float32", "float16", "bfloat16"]:
       raise ValueError(f"""`router_dtype` must be one of 'float32', 'float16' or 'bfloat16', got {router_dtype}"""   )
     SCREAMING_SNAKE_CASE__    							=   router_dtype

     SCREAMING_SNAKE_CASE__    							=   router_ignore_padding_tokens
     SCREAMING_SNAKE_CASE__    							=   batch_prioritized_routing
     SCREAMING_SNAKE_CASE__    							=   second_expert_policy
     SCREAMING_SNAKE_CASE__    							=   normalize_router_prob_before_dropping
     SCREAMING_SNAKE_CASE__    							=   moe_eval_capacity_token_fraction
     SCREAMING_SNAKE_CASE__    							=   moe_token_dropout
     SCREAMING_SNAKE_CASE__    							=   output_router_logits
     super().__init__(
         pad_token_id=_lowercase ,			bos_token_id=_lowercase ,			eos_token_id=_lowercase ,			is_encoder_decoder=_lowercase ,			decoder_start_token_id=_lowercase ,			**_lowercase ,			)
