


'''simple docstring'''


import json
import re
from typing import TYPE_CHECKING, List, Optional, Tuple, Union

import numpy as np

from ...utils import is_tf_available, is_torch_available, logging


if TYPE_CHECKING:
						if is_torch_available():
												import torch
						if is_tf_available():
												import tensorflow as tf

from tokenizers import pre_tokenizers

from ...tokenization_utils_base import BatchEncoding
from ...tokenization_utils_fast import PreTrainedTokenizerFast
from .tokenization_codegen import CodeGenTokenizer


lowerCamelCase__					=   logging.get_logger(__name__)

lowerCamelCase__					=   {'vocab_file': 'vocab.json', 'merges_file': 'merges.txt', 'tokenizer_file': 'tokenizer.json'}

lowerCamelCase__					=   {
    'vocab_file': {
        'Salesforce/codegen-350M-mono': 'https://huggingface.co/Salesforce/codegen-350M-mono/resolve/main/vocab.json',
    },
    'merges_file': {
        'Salesforce/codegen-350M-mono': 'https://huggingface.co/Salesforce/codegen-350M-mono/resolve/main/merges.txt',
    },
    'tokenizer_file': {
        'Salesforce/codegen-350M-mono': (
            'https://huggingface.co/Salesforce/codegen-350M-mono/resolve/main/tokenizer.json'
        ),
    },
}

lowerCamelCase__					=   {
    'Salesforce/codegen-350M-mono': 2_048,
}






class 		lowerCAmelCase__	(	UpperCAmelCase__		):
			lowerCAmelCase       :		Optional[int]				      =   VOCAB_FILES_NAMES
			lowerCAmelCase       :		Any				      =   PRETRAINED_VOCAB_FILES_MAP
			lowerCAmelCase       :		Tuple				      =   PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
			lowerCAmelCase       :		List[Any]				      =   ["input_ids", "attention_mask"]
			lowerCAmelCase       :		Any				      =   CodeGenTokenizer




			def __init__(						self     :   Any		,     lowerCamelCase__     :   str=None		,     lowerCamelCase__     :   Union[str, Any]=None		,     lowerCamelCase__     :   List[str]=None		,     lowerCamelCase__     :   Dict="<|endoftext|>"		,     lowerCamelCase__     :   Optional[int]="<|endoftext|>"		,     lowerCamelCase__     :   Tuple="<|endoftext|>"		,     lowerCamelCase__     :   int=False		,     **lowerCamelCase__     :   Tuple		,     )      ->str:



					'''simple docstring'''







					super().__init__(
					    lowerCamelCase__		,     lowerCamelCase__		,     tokenizer_file=lowerCamelCase__		,     unk_token=lowerCamelCase__		,     bos_token=lowerCamelCase__		,     eos_token=lowerCamelCase__		,     add_prefix_space=lowerCamelCase__		,     **lowerCamelCase__		,     )

					if kwargs.pop("add_bos_token"		,     lowerCamelCase__       ):
							_UpperCAmelCase						:				Optional[Any]	    =					kwargs.pop("name_or_path"		,     ""       )
							raise ValueError(
							    "Currenty GPT2's fast tokenizer does NOT support adding a BOS token."
							    "Instead you should use GPT2's slow tokenizer class `CodeGenTokenizer` as follows: \n"
							    F"""`CodeGenTokenizer.from_pretrained('{model_id}')`\nor\n"""
							    F"""`AutoTokenizer.from_pretrained('{model_id}', use_fast=False)`\n"""
							    "This issue will be fixed soon, see: https://github.com/huggingface/tokenizers/pull/1005."
							    " so that the fast tokenizer works correctly."       )

					_UpperCAmelCase						:				Union[str, Any]	    =					json.loads(self.backend_tokenizer.pre_tokenizer.__getstate__()       )
					if pre_tok_state.get("add_prefix_space"		,     lowerCamelCase__       ) != add_prefix_space:
							_UpperCAmelCase						:				Optional[Any]	    =					getattr(lowerCamelCase__		,     pre_tok_state.pop("type"       )       )
							_UpperCAmelCase						:				int	    =					add_prefix_space
							_UpperCAmelCase						:				int	    =					pre_tok_class(**lowerCamelCase__       )

					_UpperCAmelCase						:				List[Any]	    =					add_prefix_space




			def 							lowerCAmelCase__			(						self     :   Optional[int]		,     *lowerCamelCase__     :   int		,     **lowerCamelCase__     :   Any       )      ->BatchEncoding:



					'''simple docstring'''







					_UpperCAmelCase						:				Tuple	    =					kwargs.get("is_split_into_words"		,     lowerCamelCase__       )
					assert self.add_prefix_space or not is_split_into_words, (
					    F"""You need to instantiate {self.__class__.__name__} with add_prefix_space=True """
					    "to use it with pretokenized inputs."
					)

					return super()._batch_encode_plus(*lowerCamelCase__		,     **lowerCamelCase__       )




			def 							lowerCAmelCase__			(						self     :   Optional[int]		,     *lowerCamelCase__     :   int		,     **lowerCamelCase__     :   Any       )      ->BatchEncoding:



					'''simple docstring'''







					_UpperCAmelCase						:				Any	    =					kwargs.get("is_split_into_words"		,     lowerCamelCase__       )

					assert self.add_prefix_space or not is_split_into_words, (
					    F"""You need to instantiate {self.__class__.__name__} with add_prefix_space=True """
					    "to use it with pretokenized inputs."
					)

					return super()._encode_plus(*lowerCamelCase__		,     **lowerCamelCase__       )




			def 							lowerCAmelCase__			(						self     :   Optional[int]		,     lowerCamelCase__     :   str		,     lowerCamelCase__     :   Optional[str] = None       )      ->Tuple[str]:



					'''simple docstring'''







					_UpperCAmelCase						:				Optional[int]	    =					self._tokenizer.model.save(lowerCamelCase__		,     name=lowerCamelCase__       )
					return tuple(lowerCamelCase__       )




			def 							lowerCAmelCase__			(						self     :   Optional[int]		,     lowerCamelCase__     :   Union[int, List[int], "np.ndarray", "torch.Tensor", "tf.Tensor"]		,     lowerCamelCase__     :   bool = False		,     lowerCamelCase__     :   bool = None		,     lowerCamelCase__     :   Optional[List[str]] = None		,     **lowerCamelCase__     :   Union[str, Any]		,     )      ->str:



					'''simple docstring'''







					_UpperCAmelCase						:				str	    =					super().decode(
					    token_ids=lowerCamelCase__		,     skip_special_tokens=lowerCamelCase__		,     clean_up_tokenization_spaces=lowerCamelCase__		,     **lowerCamelCase__		,     )

					if truncate_before_pattern is not None and len(lowerCamelCase__       ) > 0:
							_UpperCAmelCase						:				int	    =					self.truncate(lowerCamelCase__		,     lowerCamelCase__       )

					return decoded_text






			def 							lowerCAmelCase__			(						self     :   int		,     lowerCamelCase__     :   int		,     lowerCamelCase__     :   int       )      ->Tuple:



					'''simple docstring'''







					def find_re(lowerCamelCase__     :   Optional[Any]		,     lowerCamelCase__     :   Optional[Any]		,     lowerCamelCase__     :   Tuple       ):
							_UpperCAmelCase						:				Optional[int]	    =					pattern.search(lowerCamelCase__		,     lowerCamelCase__       )
							return m.start() if m else -1

					_UpperCAmelCase						:				Dict	    =					[re.compile(lowerCamelCase__		,     re.MULTILINE       ) for pattern in truncate_before_pattern]

					_UpperCAmelCase						:				Any	    =					list(re.finditer("^print"		,     lowerCamelCase__		,     re.MULTILINE       )       )

					if len(lowerCamelCase__       ) > 1:
							_UpperCAmelCase						:				Tuple	    =					completion[: prints[1].start()]

					_UpperCAmelCase						:				Dict	    =					list(re.finditer("^def"		,     lowerCamelCase__		,     re.MULTILINE       )       )

					if len(lowerCamelCase__       ) > 1:
							_UpperCAmelCase						:				Optional[Any]	    =					completion[: defs[1].start()]

					_UpperCAmelCase						:				List[str]	    =					0

					_UpperCAmelCase						:				Tuple	    =					[
					    pos for pos in [find_re(lowerCamelCase__		,     lowerCamelCase__		,     lowerCamelCase__       ) for terminal in terminals] if pos != -1
					]

					if len(lowerCamelCase__       ) > 0:
							return completion[: min(lowerCamelCase__       )]
					else:
							return completion







'''simple docstring'''


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       (__lowerCAmelCase      ):
		_UpperCAmelCase						:				List[str]	    =					[
		    "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(__lowerCAmelCase					,       __lowerCAmelCase      )






def        __lowerCAmelCase       (__lowerCAmelCase      ):
		_UpperCAmelCase    ,			_UpperCAmelCase						:				Optional[Any]	    =					emb.weight.shape
		_UpperCAmelCase						:				str	    =					nn.Linear(__lowerCAmelCase					,       __lowerCAmelCase					,       bias=__lowerCAmelCase      )
		_UpperCAmelCase						:				Optional[int]	    =					emb.weight.data
		return lin_layer






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






def        __lowerCAmelCase       (__lowerCAmelCase					,       __lowerCAmelCase					,       __lowerCAmelCase					,       __lowerCAmelCase					,       __lowerCAmelCase = WEIGHTS_NAME      ):
		_UpperCAmelCase						:				Optional[int]	    =					[]
		_UpperCAmelCase						:				Optional[Any]	    =					0
		os.makedirs(__lowerCAmelCase					,       exist_ok=__lowerCAmelCase      )

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

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

    # Add the metadata
		_UpperCAmelCase						:				Any	    =					{"total_size": total_size}
		_UpperCAmelCase						:				List[str]	    =					{"metadata": metadata, "weight_map": weight_map}

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

		return metadata, index


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

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



