
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 		UpperCamelCase	(      lowercase_   )		->	Optional[Any]:


  '''simple docstring'''






  lowercase__  :   Optional[Any]    		=      [
      """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(lowercase_ , lowercase_   )

def 		UpperCamelCase	(      lowercase_   )		->	Any:


  '''simple docstring'''






  lowercase__		,    lowercase__  :   Dict    		=      emb.weight.shape
  lowercase__  :   Optional[Any]    		=      nn.Linear(lowercase_ , lowercase_ , bias=lowercase_   )
  lowercase__  :   int    		=      emb.weight.data
  return lin_layer

def 		UpperCamelCase	(      lowercase_ , lowercase_=None   )		->	List[str]:


  '''simple docstring'''






  lowercase__  :   Optional[int]    		=      {}
  for old_key in state_dict.keys():
    lowercase__  :   int    		=      old_key
    if "moe_layer.experts." in key:
      if expert_idx is not None:
        lowercase__  :   List[str]    		=      key.replace("""moe_layer.experts.0""" , F'ffn.experts.expert_{expert_idx}'   )
      else:
        lowercase__  :   Optional[Any]    		=      key.replace("""moe_layer.experts.""" , """ffn.experts.expert_"""   )
    if "gate" in key:
      lowercase__  :   Tuple    		=      key.replace(""".moe_layer.gate.wg""" , """.ffn.router.classifier"""   )
    if "fc2" and "experts" not in key:
      lowercase__  :   Union[str, Any]    		=      key.replace(""".fc2.""" , """.ffn.fc2."""   )
    if "fc1" and "experts" not in key:
      lowercase__  :   Dict    		=      key.replace(""".fc1.""" , """.ffn.fc1."""   )
    if ".encoder_attn." in key:
      lowercase__  :   Tuple    		=      key.replace(""".encoder_attn.""" , """.cross_attention."""   )
    if "encoder_attn_layer_norm" in key:
      lowercase__  :   Dict    		=      key.replace("""encoder_attn_layer_norm""" , """cross_attention_layer_norm"""   )
    if "final_layer_norm" in key:
      lowercase__  :   Union[str, Any]    		=      key.replace("""final_layer_norm""" , """ff_layer_norm"""   )
    lowercase__  :   Union[str, Any]    		=      state_dict[old_key]
  return new_dict

def 		UpperCamelCase	(      lowercase_ , lowercase_ , lowercase_ , lowercase_ , lowercase_ = WEIGHTS_NAME   )		->	List[Any]:


  '''simple docstring'''






  lowercase__  :   Any    		=      []
  lowercase__  :   int    		=      0
  os.makedirs(lowercase_ , exist_ok=lowercase_   )

  for expert in range(lowercase_   ):
    lowercase__  :   str    		=      switch_checkpoint_path + F'-rank-{expert}.pt'
    if os.path.isfile(lowercase_   ):
      lowercase__  :   Any    		=      torch.load(lowercase_   )["""model"""]
      remove_ignore_keys_(lowercase_   )
      lowercase__  :   List[str]    		=      rename_fairseq_keys(lowercase_ , lowercase_   )
      lowercase__  :   List[Any]    		=      os.path.join(
          lowercase_ , weights_name.replace(""".bin""" , F'-{len(lowercase_   )+1:05d}-of-???.bin'   )   )
      torch.save(lowercase_ , lowercase_   )
      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(lowercase_   )[0]].dtype   )

    # Add the last block
  lowercase__  :   Optional[int]    		=      os.path.join(lowercase_ , weights_name.replace(""".bin""" , F'-{len(lowercase_   )+1:05d}-of-???.bin'   )   )
  lowercase__  :   Any    		=      torch.load(switch_checkpoint_path + """-shared.pt"""   )["""model"""]
  remove_ignore_keys_(lowercase_   )
  lowercase__  :   Optional[int]    		=      rename_fairseq_keys(lowercase_ , lowercase_   )
  lowercase__  :   Tuple    		=      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(lowercase_   ) == 1:
    lowercase__  :   Tuple    		=      os.path.join(lowercase_ , lowercase_   )
    torch.save(lowercase_ , lowercase_   )
    return {weights_name: sharded_state_dicts[0]}, None
  else:
    torch.save(lowercase_ , lowercase_   )
  # Otherwise, let's build the index
  lowercase__  :   Optional[Any]    		=      {}
  for idx, shard in enumerate(lowercase_   ):
    lowercase__  :   List[Any]    		=      weights_name.replace(""".bin""" , F'-{idx+1:05d}-of-{len(lowercase_   ):05d}.bin'   )
    lowercase__  :   Dict    		=      os.path.join(lowercase_ , weights_name.replace(""".bin""" , F'-{idx+1:05d}-of-???.bin'   )   )
    os.rename(lowercase_ , os.path.join(lowercase_ , lowercase_   )   )
    for key in shard:
      lowercase__  :   Optional[Any]    		=      shard_file

    # Add the metadata
  lowercase__  :   Any    		=      {"""total_size""": total_size}
  lowercase__  :   Optional[int]    		=      {"""metadata""": metadata, """weight_map""": weight_map}

  with open(os.path.join(lowercase_ , lowercase_   ) , """w""" , encoding="""utf-8"""   ) as f:
    lowercase__  :   int    		=      json.dumps(lowercase_ , indent=2 , sort_keys=lowercase_   ) + """\n"""
    f.write(lowercase_   )

  return metadata, index


if __name__ == "__main__":
     lowerCamelCase__    :		Any          =   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.""",
     )
     lowerCamelCase__    :		Optional[int]          =   parser.parse_args()
     lowerCamelCase__      ,						lowerCamelCase__    :		List[Any]          =   shard_on_the_fly(
         args.nllb_moe_checkpoint_path,
         args.pytorch_dump_folder_path,
         1_2_8,
         args.dtype,
     )

     lowerCamelCase__    :		Any          =   NllbMoeConfig.from_pretrained(
         """facebook/nllb-200-3.3B""", encoder_sparse_step=4, decoder_sparse_step=4, num_experts=1_2_8
     )
     config.save_pretrained(args.pytorch_dump_folder_path)
     lowerCamelCase__    :		Dict          =   NllbMoeModel.from_pretrained(args.pytorch_dump_folder_path)
     print("""Done""")
     model.save_pretrained(args.pytorch_dump_folder_path)








from math import factorial

lowerCamelCase__    :		dict[str, int]          =   {str(digit): factorial(digit) for digit in range(1_0)}

def 		UpperCamelCase	(      lowercase_   )		->	int:


  '''simple docstring'''






  if not isinstance(lowercase_ , lowercase_   ):
    raise TypeError("""Parameter number must be int"""   )

  if number < 0:
    raise ValueError("""Parameter number must be greater than or equal to 0"""   )

  # Converts number in string to iterate on its digits and adds its factorial.
  return sum(DIGIT_FACTORIAL[digit] for digit in str(lowercase_   )   )

def 		UpperCamelCase	(      lowercase_ = 60 , lowercase_ = 1_00_00_00   )		->	int:


  '''simple docstring'''






  if not isinstance(lowercase_ , lowercase_   ) or not isinstance(lowercase_ , lowercase_   ):
    raise TypeError("""Parameters chain_length and number_limit must be int"""   )

  if chain_length <= 0 or number_limit <= 0:
    raise ValueError(
        """Parameters chain_length and number_limit must be greater than 0"""   )

  # the counter for the chains with the exact desired length
  lowercase__  :   List[Any]    		=      0
  # the cached sizes of the previous chains
  lowercase__  :   dict[int, int]    		=      {}

  for start_chain_element in range(1 , lowercase_   ):
    # The temporary set will contain the elements of the chain
    lowercase__  :   Optional[Any]    		=      set()
    lowercase__  :   Union[str, Any]    		=      0

    # Stop computing the chain when you find a cached size, a repeating item or the
    # length is greater then the desired one.
    lowercase__  :   Union[str, Any]    		=      start_chain_element
    while (
        chain_element not in chain_sets_lengths
        and chain_element not in chain_set
        and chain_set_length <= chain_length
    ):
      chain_set.add(lowercase_   )
      chain_set_length += 1
      lowercase__  :   str    		=      digit_factorial_sum(lowercase_   )

    if chain_element in chain_sets_lengths:
      chain_set_length += chain_sets_lengths[chain_element]

    lowercase__  :   str    		=      chain_set_length

    # If chain contains the exact amount of elements increase the counter
    if chain_set_length == chain_length:
      chains_counter += 1

  return chains_counter


if __name__ == "__main__":
     import doctest

     doctest.testmod()
     print(f'''{solution()}''')






