


'''simple docstring'''


import os
import re
import warnings
from shutil import copyfile
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple

import sentencepiece as spm

from ...tokenization_utils import PreTrainedTokenizer


if TYPE_CHECKING:
						from ...tokenization_utils_base import TextInput
from ...utils import logging


lowerCamelCase__					=   logging.get_logger(__name__)

lowerCamelCase__					=   {'vocab_file': 'spiece.model'}

lowerCamelCase__					=   {
    'vocab_file': {
        't5-small': 'https://huggingface.co/t5-small/resolve/main/spiece.model',
        't5-base': 'https://huggingface.co/t5-base/resolve/main/spiece.model',
        't5-large': 'https://huggingface.co/t5-large/resolve/main/spiece.model',
        't5-3b': 'https://huggingface.co/t5-3b/resolve/main/spiece.model',
        't5-11b': 'https://huggingface.co/t5-11b/resolve/main/spiece.model',
    }
}


# TODO(PVP) - this should be removed in Transformers v5
lowerCamelCase__					=   {
    't5-small': 512,
    't5-base': 512,
    't5-large': 512,
    't5-3b': 512,
    't5-11b': 512,
}

lowerCamelCase__					=   '▁'






class 		lowerCAmelCase__	(	UpperCAmelCase__		):
			lowerCAmelCase       :		str				      =   VOCAB_FILES_NAMES
			lowerCAmelCase       :		str				      =   PRETRAINED_VOCAB_FILES_MAP
			lowerCAmelCase       :		Any				      =   PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
			lowerCAmelCase       :		Tuple				      =   ["input_ids", "attention_mask"]




			def __init__(						self     :   List[str]		,     lowerCamelCase__     :   List[str]		,     lowerCamelCase__     :   int="</s>"		,     lowerCamelCase__     :   List[Any]="<unk>"		,     lowerCamelCase__     :   str="<pad>"		,     lowerCamelCase__     :   Tuple=1_00		,     lowerCamelCase__     :   int=None		,     lowerCamelCase__     :   Optional[Dict[str, Any]] = None		,     lowerCamelCase__     :   Optional[Any]=True		,     **lowerCamelCase__     :   List[str]		,     )      ->None:



					'''simple docstring'''







					if extra_ids > 0 and additional_special_tokens is None:
							_UpperCAmelCase						:				Dict	    =					[F"""<extra_id_{i}>""" for i in range(lowerCamelCase__       )]
					elif extra_ids > 0 and additional_special_tokens is not None:
							# Check that we have the right number of extra_id special tokens
							_UpperCAmelCase						:				List[str]	    =					len(set(filter(lambda lowerCamelCase__       : bool("extra_id" in str(lowerCamelCase__       )       )		,     lowerCamelCase__       )       )       )
							if extra_tokens != extra_ids:
									raise ValueError(
									    F"""Both extra_ids ({extra_ids}) and additional_special_tokens ({additional_special_tokens}) are"""
									    " provided to T5Tokenizer. In this case the additional_special_tokens must include the extra_ids"
									    " tokens"       )
					if legacy:
							logger.warning_once(
							    F"""You are using the legacy behaviour of the {self.__class__}. This means that tokens that come after special tokens will not be properly handled. We recommend you to"""
							    " read the related pull request available at https://github.com/huggingface/transformers/pull/24565"       )

					_UpperCAmelCase						:				List[str]	    =					legacy
					_UpperCAmelCase						:				Union[str, Any]	    =					{} if sp_model_kwargs is None else sp_model_kwargs

					super().__init__(
					    eos_token=lowerCamelCase__		,     unk_token=lowerCamelCase__		,     pad_token=lowerCamelCase__		,     extra_ids=lowerCamelCase__		,     additional_special_tokens=lowerCamelCase__		,     sp_model_kwargs=self.sp_model_kwargs		,     legacy=lowerCamelCase__		,     **lowerCamelCase__		,     )

					_UpperCAmelCase						:				str	    =					vocab_file
					_UpperCAmelCase						:				str	    =					extra_ids

					_UpperCAmelCase						:				Tuple	    =					spm.SentencePieceProcessor(**self.sp_model_kwargs       )
					self.sp_model.Load(lowerCamelCase__       )




			@staticmethod
			def 							lowerCAmelCase__			(						lowerCamelCase__     :   Any		,     lowerCamelCase__     :   Tuple		,     lowerCamelCase__     :   List[Any]       )      ->Optional[int]:



					'''simple docstring'''







					if pretrained_model_name_or_path in TaTokenizer.max_model_input_sizes:
							_UpperCAmelCase						:				Any	    =					TaTokenizer.max_model_input_sizes[pretrained_model_name_or_path]
							if init_max_model_length is not None and init_max_model_length != max_model_length:
									return init_max_model_length
							elif init_max_model_length is None:
									warnings.warn(
									    "This tokenizer was incorrectly instantiated with a model max length of"
									    F""" {deprecated_max_model_length} which will be corrected in Transformers v5.\nFor now, this"""
									    " behavior is kept to avoid breaking backwards compatibility when padding/encoding with"
									    " `truncation is True`.\n- Be aware that you SHOULD NOT rely on"
									    F""" {pretrained_model_name_or_path} automatically truncating your input to"""
									    F""" {deprecated_max_model_length} when padding/encoding.\n- If you want to encode/pad to sequences"""
									    F""" longer than {deprecated_max_model_length} you can either instantiate this tokenizer with"""
									    " `model_max_length` or pass `max_length` when encoding/padding.\n- To avoid this warning, please"
									    " instantiate this tokenizer with `model_max_length` set to your preferred value."		,     lowerCamelCase__		,     )

					return max_model_length




			@property
			def 							lowerCAmelCase__			(						self     :   Optional[int]       )      ->Optional[int]:



					'''simple docstring'''







					return self.sp_model.get_piece_size() + self._extra_ids




			def 							lowerCAmelCase__			(						self     :   Dict       )      ->Optional[int]:



					'''simple docstring'''







					_UpperCAmelCase						:				Optional[int]	    =					{self.convert_ids_to_tokens(lowerCamelCase__       ): i for i in range(self.vocab_size       )}
					vocab.update(self.added_tokens_encoder       )
					return vocab




			def 							lowerCAmelCase__			(						self     :   Optional[Any]		,     lowerCamelCase__     :   List[int]		,     lowerCamelCase__     :   Optional[List[int]] = None		,     lowerCamelCase__     :   bool = False       )      ->List[int]:



					'''simple docstring'''







					if already_has_special_tokens:
							return super().get_special_tokens_mask(
							    token_ids_a=lowerCamelCase__		,     token_ids_a=lowerCamelCase__		,     already_has_special_tokens=lowerCamelCase__       )

					# normal case: some special tokens
					if token_ids_a is None:
							return ([0] * len(lowerCamelCase__       )) + [1]
					return ([0] * len(lowerCamelCase__       )) + [1] + ([0] * len(lowerCamelCase__       )) + [1]




			def 							lowerCAmelCase__			(						self     :   Optional[Any]       )      ->Union[str, Any]:



					'''simple docstring'''







					return list(
					    set(filter(lambda lowerCamelCase__       : bool(re.search(R"<extra_id_\d+>"		,     lowerCamelCase__       )       ) is not None		,     self.additional_special_tokens       )       )       )




			def 							lowerCAmelCase__			(						self     :   Optional[int]       )      ->Optional[Any]:



					'''simple docstring'''







					return [self._convert_token_to_id(lowerCamelCase__       ) for token in self.get_sentinel_tokens()]




			def 							lowerCAmelCase__			(						self     :   Dict		,     lowerCamelCase__     :   List[int]       )      ->List[int]:



					'''simple docstring'''







					if len(lowerCamelCase__       ) > 0 and token_ids[-1] == self.eos_token_id:
							warnings.warn(
							    F"""This sequence already has {self.eos_token}. In future versions this behavior may lead to duplicated"""
							    " eos tokens being added."       )
							return token_ids
					else:
							return token_ids + [self.eos_token_id]




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



					'''simple docstring'''







					_UpperCAmelCase						:				Optional[int]	    =					[self.eos_token_id]

					if token_ids_a is None:
							return len(token_ids_a + eos       ) * [0]
					return len(token_ids_a + eos + token_ids_a + eos       ) * [0]




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



					'''simple docstring'''







					_UpperCAmelCase						:				Optional[Any]	    =					self._add_eos_if_not_present(lowerCamelCase__       )
					if token_ids_a is None:
							return token_ids_a
					else:
							_UpperCAmelCase						:				int	    =					self._add_eos_if_not_present(lowerCamelCase__       )
							return token_ids_a + token_ids_a




			def __getstate__(						self     :   int       )      ->Tuple:



					'''simple docstring'''







					_UpperCAmelCase						:				List[str]	    =					self.__dict__.copy()
					_UpperCAmelCase						:				Union[str, Any]	    =					None
					return state




			def __setstate__(						self     :   Union[str, Any]		,     lowerCamelCase__     :   Tuple       )      ->Union[str, Any]:



					'''simple docstring'''







					_UpperCAmelCase						:				Optional[Any]	    =					d

					# for backward compatibility
					if not hasattr(self		,     "sp_model_kwargs"       ):
							_UpperCAmelCase						:				Tuple	    =					{}

					_UpperCAmelCase						:				str	    =					spm.SentencePieceProcessor(**self.sp_model_kwargs       )
					self.sp_model.Load(self.vocab_file       )




			def 							lowerCAmelCase__			(						self     :   List[str]		,     lowerCamelCase__     :   "TextInput"		,     **lowerCamelCase__     :   str       )      ->List[str]:



					'''simple docstring'''







					if not self.legacy:
							_UpperCAmelCase						:				Tuple	    =					SPIECE_UNDERLINE + text.replace(lowerCamelCase__		,     " "       )
					return super().tokenize(lowerCamelCase__		,     **lowerCamelCase__       )




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



					'''simple docstring'''







					if not self.legacy:
							_UpperCAmelCase						:				int	    =					text.startswith(lowerCamelCase__       )
							if is_first:
									_UpperCAmelCase						:				List[Any]	    =					text[1:]

					_UpperCAmelCase						:				Optional[int]	    =					self.sp_model.encode(lowerCamelCase__		,     out_type=lowerCamelCase__       )

					if not self.legacy and not is_first and not text.startswith(" "       ) and tokens[0].startswith(lowerCamelCase__       ):
							_UpperCAmelCase						:				List[Any]	    =					([tokens[0][1:]] if len(tokens[0]       ) > 1 else []) + tokens[1:]
					return tokens




			def 							lowerCAmelCase__			(						self     :   Any		,     lowerCamelCase__     :   Any       )      ->Any:



					'''simple docstring'''







					if token.startswith("<extra_id_"       ):
							_UpperCAmelCase						:				int	    =					re.match(R"<extra_id_(\d+)>"		,     lowerCamelCase__       )
							_UpperCAmelCase						:				Tuple	    =					int(match.group(1       )       )
							return self.vocab_size - num - 1
					return self.sp_model.piece_to_id(lowerCamelCase__       )




			def 							lowerCAmelCase__			(						self     :   List[str]		,     lowerCamelCase__     :   Optional[int]       )      ->Any:



					'''simple docstring'''







					if index < self.sp_model.get_piece_size():
							_UpperCAmelCase						:				Dict	    =					self.sp_model.IdToPiece(lowerCamelCase__       )
					else:
							_UpperCAmelCase						:				List[Any]	    =					F"""<extra_id_{self.vocab_size - 1 - index}>"""
					return token




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



					'''simple docstring'''







					_UpperCAmelCase						:				Optional[Any]	    =					[]
					_UpperCAmelCase						:				int	    =					""
					_UpperCAmelCase						:				Union[str, Any]	    =					False
					for token in tokens:
							# make sure that special tokens are not decoded using sentencepiece model
							if token in self.all_special_tokens:
									if not prev_is_special:
											out_string += " "
									out_string += self.sp_model.decode(lowerCamelCase__       ) + token
									_UpperCAmelCase						:				Dict	    =					True
									_UpperCAmelCase						:				Optional[Any]	    =					[]
							else:
									current_sub_tokens.append(lowerCamelCase__       )
									_UpperCAmelCase						:				Dict	    =					False
					out_string += self.sp_model.decode(lowerCamelCase__       )
					return out_string.strip()






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



					'''simple docstring'''







					if not os.path.isdir(lowerCamelCase__       ):
							logger.error(F"""Vocabulary path ({save_directory}) should be a directory"""       )
							return
					_UpperCAmelCase						:				List[Any]	    =					os.path.join(
					    lowerCamelCase__		,     (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["vocab_file"]       )

					if os.path.abspath(self.vocab_file       ) != os.path.abspath(lowerCamelCase__       ) and os.path.isfile(self.vocab_file       ):
							copyfile(self.vocab_file		,     lowerCamelCase__       )
					elif not os.path.isfile(self.vocab_file       ):
							with open(lowerCamelCase__		,     "wb"       ) as fi:
									_UpperCAmelCase						:				str	    =					self.sp_model.serialized_model_proto()
									fi.write(lowerCamelCase__       )

					return (out_vocab_file,)







'''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)



