



import argparse
import json
import os

import torch
from torch import nn

from transformers import NllbMoeConfig, NllbMoeModel
from transformers.modeling_utils import dtype_byte_size
from transformers.utils import WEIGHTS_INDEX_NAME, WEIGHTS_NAME



def __lowerCAmelCase	(       SCREAMING_SNAKE_CASE_					):
			lowercase__   					=			[
			    "encoder.version",
			    "decoder.version",
			    "model.encoder.version",
			    "model.decoder.version",
			    "decoder.output_projection.weight",
			    "_float_tensor",
			    "encoder.embed_positions._float_tensor",
			    "decoder.embed_positions._float_tensor",
			]
			for k in ignore_keys:
						state_dict.pop(SCREAMING_SNAKE_CASE_  ,					SCREAMING_SNAKE_CASE_					)



def __lowerCAmelCase	(       SCREAMING_SNAKE_CASE_					):
			lowercase__					,    lowercase__   					=			emb.weight.shape
			lowercase__   					=			nn.Linear(SCREAMING_SNAKE_CASE_  ,					SCREAMING_SNAKE_CASE_  ,					bias=SCREAMING_SNAKE_CASE_					)
			lowercase__   					=			emb.weight.data
			return lin_layer



def __lowerCAmelCase	(       SCREAMING_SNAKE_CASE_  ,					SCREAMING_SNAKE_CASE_=None					):
			lowercase__   					=			{}
			for old_key in state_dict.keys():
						lowercase__   					=			old_key
						if "moe_layer.experts." in key:
									if expert_idx is not None:
												lowercase__   					=			key.replace("moe_layer.experts.0"  ,					f'''ffn.experts.expert_{expert_idx}'''					)
									else:
												lowercase__   					=			key.replace("moe_layer.experts."  ,					"ffn.experts.expert_"					)
						if "gate" in key:
									lowercase__   					=			key.replace(".moe_layer.gate.wg"  ,					".ffn.router.classifier"					)
						if "fc2" and "experts" not in key:
									lowercase__   					=			key.replace(".fc2."  ,					".ffn.fc2."					)
						if "fc1" and "experts" not in key:
									lowercase__   					=			key.replace(".fc1."  ,					".ffn.fc1."					)
						if ".encoder_attn." in key:
									lowercase__   					=			key.replace(".encoder_attn."  ,					".cross_attention."					)
						if "encoder_attn_layer_norm" in key:
									lowercase__   					=			key.replace("encoder_attn_layer_norm"  ,					"cross_attention_layer_norm"					)
						if "final_layer_norm" in key:
									lowercase__   					=			key.replace("final_layer_norm"  ,					"ff_layer_norm"					)
						lowercase__   					=			state_dict[old_key]
			return new_dict



def __lowerCAmelCase	(       SCREAMING_SNAKE_CASE_  ,					SCREAMING_SNAKE_CASE_  ,					SCREAMING_SNAKE_CASE_  ,					SCREAMING_SNAKE_CASE_  ,					SCREAMING_SNAKE_CASE_ = WEIGHTS_NAME					):
			lowercase__   					=			[]
			lowercase__   					=			0
			os.makedirs(SCREAMING_SNAKE_CASE_  ,					exist_ok=SCREAMING_SNAKE_CASE_					)

			for expert in range(SCREAMING_SNAKE_CASE_					):
						lowercase__   					=			switch_checkpoint_path + f'''-rank-{expert}.pt'''
						if os.path.isfile(SCREAMING_SNAKE_CASE_					):
									lowercase__   					=			torch.load(SCREAMING_SNAKE_CASE_					)["model"]
									remove_ignore_keys_(SCREAMING_SNAKE_CASE_					)
									lowercase__   					=			rename_fairseq_keys(SCREAMING_SNAKE_CASE_  ,					SCREAMING_SNAKE_CASE_					)
									lowercase__   					=			os.path.join(
									    SCREAMING_SNAKE_CASE_  ,					weights_name.replace(".bin"  ,					f'''-{len(SCREAMING_SNAKE_CASE_					)+1:05d}-of-???.bin'''					)					)
									torch.save(SCREAMING_SNAKE_CASE_  ,					SCREAMING_SNAKE_CASE_					)
									sharded_state_dicts.append(expert_state.keys()					)
									total_size += sum([value.numel() for key, value in expert_state.items()]					) * dtype_byte_size(
									    expert_state[list(SCREAMING_SNAKE_CASE_					)[0]].dtype					)

    # Add the last block
			lowercase__   					=			os.path.join(SCREAMING_SNAKE_CASE_  ,					weights_name.replace(".bin"  ,					f'''-{len(SCREAMING_SNAKE_CASE_					)+1:05d}-of-???.bin'''					)					)
			lowercase__   					=			torch.load(switch_checkpoint_path + "-shared.pt"					)["model"]
			remove_ignore_keys_(SCREAMING_SNAKE_CASE_					)
			lowercase__   					=			rename_fairseq_keys(SCREAMING_SNAKE_CASE_  ,					SCREAMING_SNAKE_CASE_					)
			lowercase__   					=			shared_weights["decoder.embed_tokens.weight"]
			sharded_state_dicts.append(shared_weights.keys()					)

			# If we only have the shared weights (dummy model/experts saved on the same file)
			if len(SCREAMING_SNAKE_CASE_					) == 1:
						lowercase__   					=			os.path.join(SCREAMING_SNAKE_CASE_  ,					SCREAMING_SNAKE_CASE_					)
						torch.save(SCREAMING_SNAKE_CASE_  ,					SCREAMING_SNAKE_CASE_					)
						return {weights_name: sharded_state_dicts[0]}, None
			else:
						torch.save(SCREAMING_SNAKE_CASE_  ,					SCREAMING_SNAKE_CASE_					)
			# Otherwise, let's build the index
			lowercase__   					=			{}
			for idx, shard in enumerate(SCREAMING_SNAKE_CASE_					):
						lowercase__   					=			weights_name.replace(".bin"  ,					f'''-{idx+1:05d}-of-{len(SCREAMING_SNAKE_CASE_					):05d}.bin'''					)
						lowercase__   					=			os.path.join(SCREAMING_SNAKE_CASE_  ,					weights_name.replace(".bin"  ,					f'''-{idx+1:05d}-of-???.bin'''					)					)
						os.rename(SCREAMING_SNAKE_CASE_  ,					os.path.join(SCREAMING_SNAKE_CASE_  ,					SCREAMING_SNAKE_CASE_					)					)
						for key in shard:
									lowercase__   					=			shard_file

    # Add the metadata
			lowercase__   					=			{"total_size": total_size}
			lowercase__   					=			{"metadata": metadata, "weight_map": weight_map}

			with open(os.path.join(SCREAMING_SNAKE_CASE_  ,					SCREAMING_SNAKE_CASE_					)  ,					"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_					)

			return metadata, index


if __name__ == "__main__":
							lowercase_						=       argparse.ArgumentParser()
							# Required parameters
							parser.add_argument(
							    """--nllb_moe_checkpoint_path""",
							    default="""/home/arthur_huggingface_co/fairseq/weights/checkpoints/model_moe_54b/checkpoint_2_300000""",
							    type=str,
							    required=False,
							    help="""Path to a directory containing a folder per layer. Follows the original Google format.""",
							)
							parser.add_argument("""--dtype""", default="""float32""", type=str, required=False, help="""dtype of the saved model""")
							parser.add_argument(
							    """--pytorch_dump_folder_path""",
							    default="""/home/arthur_huggingface_co/fairseq/weights/checkpoints/hf-converted-moe-54b""",
							    type=str,
							    required=False,
							    help="""Path to the output pytorch model.""",
							)
							lowercase_						=       parser.parse_args()
							lowercase_			,      lowercase_						=       shard_on_the_fly(
							    args.nllb_moe_checkpoint_path,
							    args.pytorch_dump_folder_path,
							    128,
							    args.dtype,
							)

							lowercase_						=       NllbMoeConfig.from_pretrained(
							    """facebook/nllb-200-3.3B""", encoder_sparse_step=4, decoder_sparse_step=4, num_experts=128
							)
							config.save_pretrained(args.pytorch_dump_folder_path)
							lowercase_						=       NllbMoeModel.from_pretrained(args.pytorch_dump_folder_path)
							print("""Done""")
							model.save_pretrained(args.pytorch_dump_folder_path)









from ...configuration_utils import PretrainedConfig
from ...utils import logging


lowercase_						=       logging.get_logger(__name__)

lowercase_						=       {
    """facebook/nllb-moe-54B""": """https://huggingface.co/facebook/nllb-moe-54b/resolve/main/config.json""",
}
class 			_snake_case     (   lowercase__):
				UpperCamelCase__			:   Tuple       						="""nllb-moe"""
				UpperCamelCase__			:   Dict       						=["""past_key_values"""]
				UpperCamelCase__			:   List[Any]       						={"""num_attention_heads""": """encoder_attention_heads""", """hidden_size""": """d_model"""}




				def __init__(  self   :  List[str],   __lowercase   :  str=12_8112,   __lowercase   :  List[Any]=1024,   __lowercase   :  int=12,   __lowercase   :  Union[str, Any]=4096,   __lowercase   :  Union[str, Any]=16,   __lowercase   :  Tuple=12,   __lowercase   :  int=4096,   __lowercase   :  Any=16,   __lowercase   :  List[str]=0.05,   __lowercase   :  Optional[int]=0.05,   __lowercase   :  int=True,   __lowercase   :  Optional[Any]=True,   __lowercase   :  Optional[Any]="relu",   __lowercase   :  str=1024,   __lowercase   :  Optional[Any]=0.1,   __lowercase   :  List[Any]=0.1,   __lowercase   :  Union[str, Any]=0.0,   __lowercase   :  Tuple=0.02,   __lowercase   :  Optional[Any]=2,   __lowercase   :  Optional[int]=True,   __lowercase   :  List[Any]=False,   __lowercase   :  Any="float32",   __lowercase   :  Union[str, Any]=False,   __lowercase   :  Optional[int]=128,   __lowercase   :  Union[str, Any]=64,   __lowercase   :  Optional[int]=4,   __lowercase   :  Dict=4,   __lowercase   :  Any=0.001,   __lowercase   :  Optional[int]=0.001,   __lowercase   :  int="all",   __lowercase   :  Any=False,   __lowercase   :  Optional[int]=False,   __lowercase   :  Dict=1.0,   __lowercase   :  int=0.2,   __lowercase   :  str=1,   __lowercase   :  Optional[Any]=0,   __lowercase   :  Optional[int]=2,   __lowercase   :  Any=False,   **__lowercase   :  List[str],   ):
							lowercase__   					=			vocab_size
							lowercase__   					=			max_position_embeddings
							lowercase__   					=			d_model
							lowercase__   					=			encoder_ffn_dim
							lowercase__   					=			encoder_layers
							lowercase__   					=			encoder_attention_heads
							lowercase__   					=			decoder_ffn_dim
							lowercase__   					=			decoder_layers
							lowercase__   					=			decoder_attention_heads
							lowercase__   					=			dropout
							lowercase__   					=			attention_dropout
							lowercase__   					=			activation_dropout
							lowercase__   					=			activation_function
							lowercase__   					=			init_std
							lowercase__   					=			encoder_layerdrop
							lowercase__   					=			decoder_layerdrop
							lowercase__   					=			use_cache
							lowercase__   					=			encoder_layers
							lowercase__   					=			scale_embedding  # scale factor will be sqrt(d_model) if True
							lowercase__   					=			router_z_loss_coef
							lowercase__   					=			router_aux_loss_coef
							lowercase__   					=			decoder_sparse_step
							lowercase__   					=			encoder_sparse_step
							lowercase__   					=			num_experts
							lowercase__   					=			expert_capacity
							lowercase__   					=			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}'''  )
							lowercase__   					=			router_dtype

							lowercase__   					=			router_ignore_padding_tokens
							lowercase__   					=			batch_prioritized_routing
							lowercase__   					=			second_expert_policy
							lowercase__   					=			normalize_router_prob_before_dropping
							lowercase__   					=			moe_eval_capacity_token_fraction
							lowercase__   					=			moe_token_dropout
							lowercase__   					=			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,   )




