





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 argparse
import os
import pickle
import sys

import torch

from transformers import TransfoXLConfig, TransfoXLLMHeadModel, load_tf_weights_in_transfo_xl
from transformers.models.transfo_xl import tokenization_transfo_xl as data_utils
from transformers.models.transfo_xl.tokenization_transfo_xl import CORPUS_NAME, VOCAB_FILES_NAMES
from transformers.utils import CONFIG_NAME, WEIGHTS_NAME, logging


logging.set_verbosity_info()

# We do this to be able to load python 2 datasets pickles
# See e.g. https://stackoverflow.com/questions/2121874/python-pickling-after-changing-a-modules-directory/2121918#2121918
snake_case_           =      data_utils.TransfoXLTokenizer
snake_case_           =      data_utils.TransfoXLCorpus
snake_case_           =      data_utils
snake_case_           =      data_utils






def lowerCamelCase__   (							snake_case_     :  Union[str, Any]	,				snake_case_     :  int	,				snake_case_     :  List[Any]	,				snake_case_     :  int	)				->    Dict:
   if transfo_xl_dataset_file:
      # Convert a pre-processed corpus (see original TensorFlow repo)
      with open(snake_case_	,				'''rb'''	) as fp:
         __snake_case								=       pickle.load(snake_case_	,				encoding='''latin1'''	)
      # Save vocabulary and dataset cache as Dictionaries (should be better than pickles for the long-term)
      __snake_case								=       pytorch_dump_folder_path + '''/''' + VOCAB_FILES_NAMES['''pretrained_vocab_file''']
      print(f"""Save vocabulary to {pytorch_vocab_dump_path}"""	)
      __snake_case								=       corpus.vocab.__dict__
      torch.save(snake_case_	,				snake_case_	)

      __snake_case								=       corpus.__dict__
      corpus_dict_no_vocab.pop('''vocab'''	,				snake_case_	)
      __snake_case								=       pytorch_dump_folder_path + '''/''' + CORPUS_NAME
      print(f"""Save dataset to {pytorch_dataset_dump_path}"""	)
      torch.save(snake_case_	,				snake_case_	)

   if tf_checkpoint_path:
      # Convert a pre-trained TensorFlow model
      __snake_case								=       os.path.abspath(snake_case_	)
      __snake_case								=       os.path.abspath(snake_case_	)

      print(f"""Converting Transformer XL checkpoint from {tf_path} with config at {config_path}."""	)
      # Initialise PyTorch model
      if transfo_xl_config_file == "":
         __snake_case								=       TransfoXLConfig()
      else:
         __snake_case								=       TransfoXLConfig.from_json_file(snake_case_	)
      print(f"""Building PyTorch model from configuration: {config}"""	)
      __snake_case								=       TransfoXLLMHeadModel(snake_case_	)

      __snake_case								=       load_tf_weights_in_transfo_xl(snake_case_	,				snake_case_	,				snake_case_	)
      # Save pytorch-model
      __snake_case								=       os.path.join(snake_case_	,				snake_case_	)
      __snake_case								=       os.path.join(snake_case_	,				snake_case_	)
      print(f"""Save PyTorch model to {os.path.abspath(snake_case_	)}"""	)
      torch.save(model.state_dict()	,				snake_case_	)
      print(f"""Save configuration file to {os.path.abspath(snake_case_	)}"""	)
      with open(snake_case_	,				'''w'''	,				encoding='''utf-8'''	) as f:
         f.write(config.to_json_string()	)


if __name__ == "__main__":
    snake_case_           =      argparse.ArgumentParser()
    parser.add_argument(
        '--pytorch_dump_folder_path',
        default=None,
        type=str,
        required=True,
        help='Path to the folder to store the PyTorch model or dataset/vocab.',
    )
    parser.add_argument(
        '--tf_checkpoint_path',
        default='',
        type=str,
        help='An optional path to a TensorFlow checkpoint path to be converted.',
    )
    parser.add_argument(
        '--transfo_xl_config_file',
        default='',
        type=str,
        help=(
            'An optional config json file corresponding to the pre-trained BERT model. \n'
            'This specifies the model architecture.'
        ),
    )
    parser.add_argument(
        '--transfo_xl_dataset_file',
        default='',
        type=str,
        help='An optional dataset file to be converted in a vocabulary.',
    )
    snake_case_           =      parser.parse_args()
    convert_transfo_xl_checkpoint_to_pytorch(
        args.tf_checkpoint_path,
        args.transfo_xl_config_file,
        args.pytorch_dump_folder_path,
        args.transfo_xl_dataset_file,
    )

