





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__   (							snake_case_     :  Optional[Any]	)				->    Dict:
   __snake_case								=       [
       '''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(snake_case_	,				snake_case_	)






def lowerCamelCase__   (							snake_case_     :  Union[str, Any]	)				->    Union[str, Any]:
   __snake_case	,     __snake_case								=       emb.weight.shape
   __snake_case								=       nn.Linear(snake_case_	,				snake_case_	,				bias=snake_case_	)
   __snake_case								=       emb.weight.data
   return lin_layer






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






def lowerCamelCase__   (							snake_case_     :  Union[str, Any]	,				snake_case_     :  str	,				snake_case_     :  int	,				snake_case_     :  Dict	,				snake_case_     :  str = WEIGHTS_NAME	)				->    str:
   __snake_case								=       []
   __snake_case								=       0
   os.makedirs(snake_case_	,				exist_ok=snake_case_	)

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

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

    # Add the metadata
   __snake_case								=       {'''total_size''': total_size}
   __snake_case								=       {'''metadata''': metadata, '''weight_map''': weight_map}

   with open(os.path.join(snake_case_	,				snake_case_	)	,				'''w'''	,				encoding='''utf-8'''	) as f:
      __snake_case								=       json.dumps(snake_case_	,				indent=2	,				sort_keys=snake_case_	) + '''\n'''
      f.write(snake_case_	)

   return metadata, index


if __name__ == "__main__":
    snake_case_           =      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.',
    )
    snake_case_           =      parser.parse_args()
    snake_case_      ,		snake_case_           =      shard_on_the_fly(
        args.nllb_moe_checkpoint_path,
        args.pytorch_dump_folder_path,
        128,
        args.dtype,
    )

    snake_case_           =      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)
    snake_case_           =      NllbMoeModel.from_pretrained(args.pytorch_dump_folder_path)
    print('Done')
    model.save_pretrained(args.pytorch_dump_folder_path)








import re
import string

import numpy as np

import datasets


snake_case_           =      '\nReturns the rate at which the input predicted strings exactly match their references, ignoring any strings input as part of the regexes_to_ignore list.\n'

snake_case_           =      '\nArgs:\n    predictions: List of predicted texts.\n    references: List of reference texts.\n    regexes_to_ignore: List, defaults to None. Regex expressions of characters to\n        ignore when calculating the exact matches. Note: these regexes are removed\n        from the input data before the changes based on the options below (e.g. ignore_case,\n        ignore_punctuation, ignore_numbers) are applied.\n    ignore_case: Boolean, defaults to False. If true, turns everything\n        to lowercase so that capitalization differences are ignored.\n    ignore_punctuation: Boolean, defaults to False. If true, removes all punctuation before\n        comparing predictions and references.\n    ignore_numbers: Boolean, defaults to False. If true, removes all punctuation before\n        comparing predictions and references.\nReturns:\n    exact_match: Dictionary containing exact_match rate. Possible values are between 0.0 and 100.0, inclusive.\nExamples:\n    >>> exact_match = datasets.load_metric("exact_match")\n    >>> refs = ["the cat", "theater", "YELLING", "agent007"]\n    >>> preds = ["cat?", "theater", "yelling", "agent"]\n    >>> results = exact_match.compute(references=refs, predictions=preds)\n    >>> print(round(results["exact_match"], 1))\n    25.0\n\n    >>> exact_match = datasets.load_metric("exact_match")\n    >>> refs = ["the cat", "theater", "YELLING", "agent007"]\n    >>> preds = ["cat?", "theater", "yelling", "agent"]\n    >>> results = exact_match.compute(references=refs, predictions=preds, regexes_to_ignore=["the ", "yell"], ignore_case=True, ignore_punctuation=True)\n    >>> print(round(results["exact_match"], 1))\n    50.0\n\n\n    >>> exact_match = datasets.load_metric("exact_match")\n    >>> refs = ["the cat", "theater", "YELLING", "agent007"]\n    >>> preds = ["cat?", "theater", "yelling", "agent"]\n    >>> results = exact_match.compute(references=refs, predictions=preds, regexes_to_ignore=["the ", "yell", "YELL"], ignore_case=True, ignore_punctuation=True)\n    >>> print(round(results["exact_match"], 1))\n    75.0\n\n    >>> exact_match = datasets.load_metric("exact_match")\n    >>> refs = ["the cat", "theater", "YELLING", "agent007"]\n    >>> preds = ["cat?", "theater", "yelling", "agent"]\n    >>> results = exact_match.compute(references=refs, predictions=preds, regexes_to_ignore=["the ", "yell", "YELL"], ignore_case=True, ignore_punctuation=True, ignore_numbers=True)\n    >>> print(round(results["exact_match"], 1))\n    100.0\n\n    >>> exact_match = datasets.load_metric("exact_match")\n    >>> refs = ["The cat sat on the mat.", "Theaters are great.", "It\'s like comparing oranges and apples."]\n    >>> preds = ["The cat sat on the mat?", "Theaters are great.", "It\'s like comparing apples and oranges."]\n    >>> results = exact_match.compute(references=refs, predictions=preds)\n    >>> print(round(results["exact_match"], 1))\n    33.3\n\n'

snake_case_           =      '\n'

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


      def        a			(self	:				Optional[Any]				):





         """simple docstring"""



         return datasets.MetricInfo(
             description=_DESCRIPTION    ,      citation=_CITATION    ,      inputs_description=_KWARGS_DESCRIPTION    ,      features=datasets.Features(
                 {
                     '''predictions''': datasets.Value('''string'''    ,      id='''sequence'''				),
                     '''references''': datasets.Value('''string'''    ,      id='''sequence'''				),
                 }				)    ,      reference_urls=[]    ,      )


      def        a			(self	:				List[Any]    ,      a__	:				int    ,      a__	:				Optional[int]    ,      a__	:				Optional[Any]=None    ,      a__	:				Any=False    ,      a__	:				Dict=False    ,      a__	:				Tuple=False    ,      ):





         """simple docstring"""



         if regexes_to_ignore is not None:
            for s in regexes_to_ignore:
               __snake_case								=       np.array([re.sub(a__    ,      ''''''    ,      a__				) for x in predictions]				)
               __snake_case								=       np.array([re.sub(a__    ,      ''''''    ,      a__				) for x in references]				)
         else:
            __snake_case								=       np.asarray(a__				)
            __snake_case								=       np.asarray(a__				)

         if ignore_case:
            __snake_case								=       np.char.lower(a__				)
            __snake_case								=       np.char.lower(a__				)

         if ignore_punctuation:
            __snake_case								=       string.punctuation.maketrans(''''''    ,      ''''''    ,      string.punctuation				)
            __snake_case								=       np.char.translate(a__    ,      table=a__				)
            __snake_case								=       np.char.translate(a__    ,      table=a__				)

         if ignore_numbers:
            __snake_case								=       string.digits.maketrans(''''''    ,      ''''''    ,      string.digits				)
            __snake_case								=       np.char.translate(a__    ,      table=a__				)
            __snake_case								=       np.char.translate(a__    ,      table=a__				)

         __snake_case								=       predictions == references

         return {"exact_match": np.mean(a__				) * 100}

