






"""simple docstring"""





import json
import os
from pathlib import Path
from shutil import copyfile
from typing import Any, Dict, List, Optional, Tuple, Union

import sentencepiece

from ...tokenization_utils import BatchEncoding, PreTrainedTokenizer
from ...utils import logging


__lowerCamelCase							=  logging.get_logger(__name__)

__lowerCamelCase							=  '▁'

__lowerCamelCase							=  {
    'vocab_file': 'vocab.json',
    'spm_file': 'sentencepiece.bpe.model',
    'tokenizer_config_file': 'tokenizer_config.json',
}

__lowerCamelCase							=  {
    'vocab_file': {
        'facebook/m2m100_418M': 'https://huggingface.co/facebook/m2m100_418M/resolve/main/vocab.json',
        'facebook/m2m100_1.2B': 'https://huggingface.co/facebook/m2m100_1.2B/resolve/main/vocab.json',
    },
    'spm_file': {
        'facebook/m2m100_418M': 'https://huggingface.co/facebook/m2m100_418M/resolve/main/sentencepiece.bpe.model',
        'facebook/m2m100_1.2B': 'https://huggingface.co/facebook/m2m100_1.2B/resolve/main/sentencepiece.bpe.model',
    },
    'tokenizer_config_file': {
        'facebook/m2m100_418M': 'https://huggingface.co/facebook/m2m100_418M/resolve/main/tokenizer_config.json',
        'facebook/m2m100_1.2B': 'https://huggingface.co/facebook/m2m100_1.2B/resolve/main/tokenizer_config.json',
    },
}

__lowerCamelCase							=  {
    'facebook/m2m100_418M': 10_24,
}

# fmt: off
__lowerCamelCase							=  {
    'm2m100': ['af', 'am', 'ar', 'ast', 'az', 'ba', 'be', 'bg', 'bn', 'br', 'bs', 'ca', 'ceb', 'cs', 'cy', 'da', 'de', 'el', 'en', 'es', 'et', 'fa', 'ff', 'fi', 'fr', 'fy', 'ga', 'gd', 'gl', 'gu', 'ha', 'he', 'hi', 'hr', 'ht', 'hu', 'hy', 'id', 'ig', 'ilo', 'is', 'it', 'ja', 'jv', 'ka', 'kk', 'km', 'kn', 'ko', 'lb', 'lg', 'ln', 'lo', 'lt', 'lv', 'mg', 'mk', 'ml', 'mn', 'mr', 'ms', 'my', 'ne', 'nl', 'no', 'ns', 'oc', 'or', 'pa', 'pl', 'ps', 'pt', 'ro', 'ru', 'sd', 'si', 'sk', 'sl', 'so', 'sq', 'sr', 'ss', 'su', 'sv', 'sw', 'ta', 'th', 'tl', 'tn', 'tr', 'uk', 'ur', 'uz', 'vi', 'wo', 'xh', 'yi', 'yo', 'zh', 'zu'],
    'wmt21': ['en', 'ha', 'is', 'ja', 'cs', 'ru', 'zh', 'de']
}



class       __A		(   SCREAMING_SNAKE_CASE_			):
					UpperCAmelCase__								=				VOCAB_FILES_NAMES
					UpperCAmelCase__								=				PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
					UpperCAmelCase__								=				PRETRAINED_VOCAB_FILES_MAP
					UpperCAmelCase__								=				["input_ids", "attention_mask"]

					UpperCAmelCase__								=				[]
					UpperCAmelCase__								=				[]

					def __init__(			self  :      List[str]     ,		__snake_case  :      Dict     ,		__snake_case  :      List[str]     ,		__snake_case  :      str=None     ,		__snake_case  :      Tuple=None     ,		__snake_case  :      Optional[int]="<s>"     ,		__snake_case  :      Optional[int]="</s>"     ,		__snake_case  :      int="</s>"     ,		__snake_case  :      Optional[int]="<pad>"     ,		__snake_case  :      Dict="<unk>"     ,		__snake_case  :      Optional[int]="m2m100"     ,		__snake_case  :      Optional[Dict[str, Any]] = None     ,		__snake_case  :      Union[str, Any]=8     ,		**__snake_case  :      Union[str, Any]     ,		)			->							None:
											__magic_name__:	Optional[Any]							=       {} if sp_model_kwargs is None else sp_model_kwargs

											__magic_name__:	Dict							=       language_codes
											__magic_name__:	List[Any]							=       FAIRSEQ_LANGUAGE_CODES[language_codes]
											__magic_name__:	Tuple							=       {lang_code: F'__{lang_code}__' for lang_code in fairseq_language_code}

											__magic_name__:	List[str]							=       kwargs.get("""additional_special_tokens"""     ,		[]					)
											kwargs["additional_special_tokens"] += [
											    self.get_lang_token(__snake_case					)
											    for lang_code in fairseq_language_code
											    if self.get_lang_token(__snake_case					) not in kwargs["additional_special_tokens"]
											]

											super().__init__(
											    src_lang=__snake_case     ,		tgt_lang=__snake_case     ,		bos_token=__snake_case     ,		eos_token=__snake_case     ,		sep_token=__snake_case     ,		unk_token=__snake_case     ,		pad_token=__snake_case     ,		language_codes=__snake_case     ,		sp_model_kwargs=self.sp_model_kwargs     ,		num_madeup_words=__snake_case     ,		**__snake_case     ,		)

											__magic_name__:	Union[str, Any]							=       vocab_file
											__magic_name__:	List[str]							=       load_json(__snake_case					)
											__magic_name__:	Optional[Any]							=       {v: k for k, v in self.encoder.items()}
											__magic_name__:	Union[str, Any]							=       spm_file
											__magic_name__:	Any							=       load_spm(__snake_case     ,		self.sp_model_kwargs					)

											__magic_name__:	Union[str, Any]							=       len(self.encoder					)

											__magic_name__:	List[Any]							=       {
											    self.get_lang_token(__snake_case					): self.encoder_size + i for i, lang_code in enumerate(__snake_case					)
											}
											__magic_name__:	Any							=       {lang_code: self.encoder_size + i for i, lang_code in enumerate(__snake_case					)}
											__magic_name__:	int							=       {v: k for k, v in self.lang_token_to_id.items()}

											__magic_name__:	Optional[Any]							=       src_lang if src_lang is not None else """en"""
											__magic_name__:	Union[str, Any]							=       tgt_lang
											__magic_name__:	List[Any]							=       self.get_lang_id(self._src_lang					)
											self.set_src_lang_special_tokens(self._src_lang					)

											__magic_name__:	str							=       num_madeup_words

					@property
					def   lowerCamelCase__   (			self  :      Dict					)			->							int:
											return len(self.encoder					) + len(self.lang_token_to_id					)

					@property
					def   lowerCamelCase__   (			self  :      Union[str, Any]					)			->							str:
											return self._src_lang

					@src_lang.setter
					def   lowerCamelCase__   (			self  :      Tuple     ,		__snake_case  :      str					)			->							None:
											__magic_name__:	Optional[int]							=       new_src_lang
											self.set_src_lang_special_tokens(self._src_lang					)

					def   lowerCamelCase__   (			self  :      Optional[int]     ,		__snake_case  :      str					)			->							List[str]:
											return self.sp_model.encode(__snake_case     ,		out_type=__snake_case					)

					def   lowerCamelCase__   (			self  :      List[Any]     ,		__snake_case  :      str					)			->							List[Any]:
											if token in self.lang_token_to_id:
																	return self.lang_token_to_id[token]
											return self.encoder.get(__snake_case     ,		self.encoder[self.unk_token]					)

					def   lowerCamelCase__   (			self  :      Tuple     ,		__snake_case  :      int					)			->							str:
											if index in self.id_to_lang_token:
																	return self.id_to_lang_token[index]
											return self.decoder.get(__snake_case     ,		self.unk_token					)

					def   lowerCamelCase__   (			self  :      List[str]     ,		__snake_case  :      int					)			->							str:
											__magic_name__:	List[str]							=       []
											__magic_name__:	List[Any]							=       """"""
											for token in tokens:
																	# make sure that special tokens are not decoded using sentencepiece model
																	if token in self.all_special_tokens:
																							out_string += self.sp_model.decode(__snake_case					) + token
																							__magic_name__:	Any							=       []
																	else:
																							current_sub_tokens.append(__snake_case					)
											out_string += self.sp_model.decode(__snake_case					)
											return out_string.strip()

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

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

											__magic_name__:	Dict							=       [1] * len(self.prefix_tokens					)
											__magic_name__:	Tuple							=       [1] * len(self.suffix_tokens					)
											if token_ids_a is None:
																	return prefix_ones + ([0] * len(__snake_case					)) + suffix_ones
											return prefix_ones + ([0] * len(__snake_case					)) + ([0] * len(__snake_case					)) + suffix_ones

					def   lowerCamelCase__   (			self  :      Union[str, Any]     ,		__snake_case  :      List[int]     ,		__snake_case  :      Optional[List[int]] = None					)			->							List[int]:
											if token_ids_a is None:
																	return self.prefix_tokens + token_ids_a + self.suffix_tokens
											# We don't expect to process pairs, but leave the pair logic for API consistency
											return self.prefix_tokens + token_ids_a + token_ids_a + self.suffix_tokens

					def   lowerCamelCase__   (			self  :      Union[str, Any]					)			->							Dict:
											__magic_name__:	Tuple							=       {self.convert_ids_to_tokens(__snake_case					): i for i in range(self.vocab_size					)}
											vocab.update(self.added_tokens_encoder					)
											return vocab

					def __getstate__(			self  :      Optional[int]					)			->							Dict:
											__magic_name__:	Dict							=       self.__dict__.copy()
											__magic_name__:	str							=       None
											return state

					def __setstate__(			self  :      Optional[Any]     ,		__snake_case  :      Dict					)			->							None:
											__magic_name__:	List[str]							=       d

											# for backward compatibility
											if not hasattr(self     ,		"""sp_model_kwargs"""					):
																	__magic_name__:	Union[str, Any]							=       {}

											__magic_name__:	Tuple							=       load_spm(self.spm_file     ,		self.sp_model_kwargs					)

					def   lowerCamelCase__   (			self  :      List[str]     ,		__snake_case  :      str     ,		__snake_case  :      Optional[str] = None					)			->							Tuple[str]:
											__magic_name__:	List[str]							=       Path(__snake_case					)
											if not save_dir.is_dir():
																	raise OSError(F'{save_directory} should be a directory'					)
											__magic_name__:	Optional[Any]							=       save_dir / (
											    (filename_prefix + """-""" if filename_prefix else """""") + self.vocab_files_names["""vocab_file"""]
											)
											__magic_name__:	List[Any]							=       save_dir / (
											    (filename_prefix + """-""" if filename_prefix else """""") + self.vocab_files_names["""spm_file"""]
											)

											save_json(self.encoder     ,		__snake_case					)

											if os.path.abspath(self.spm_file					) != os.path.abspath(__snake_case					) and os.path.isfile(self.spm_file					):
																	copyfile(self.spm_file     ,		__snake_case					)
											elif not os.path.isfile(self.spm_file					):
																	with open(__snake_case     ,		"""wb"""					) as fi:
																							__magic_name__:	Optional[Any]							=       self.sp_model.serialized_model_proto()
																							fi.write(__snake_case					)

											return (str(__snake_case					), str(__snake_case					))

					def   lowerCamelCase__   (			self  :      Optional[Any]     ,		__snake_case  :      List[str]     ,		__snake_case  :      str = "en"     ,		__snake_case  :      Optional[List[str]] = None     ,		__snake_case  :      str = "ro"     ,		**__snake_case  :      List[str]     ,		)			->							BatchEncoding:
											__magic_name__:	Dict							=       src_lang
											__magic_name__:	Union[str, Any]							=       tgt_lang
											self.set_src_lang_special_tokens(self.src_lang					)
											return super().prepare_seqaseq_batch(__snake_case     ,		__snake_case     ,		**__snake_case					)

					def   lowerCamelCase__   (			self  :      Any     ,		__snake_case  :      List[str]     ,		__snake_case  :      Optional[str]     ,		__snake_case  :      Optional[str]     ,		**__snake_case  :      Optional[int]					)			->							Optional[int]:
											if src_lang is None or tgt_lang is None:
																	raise ValueError("""Translation requires a `src_lang` and a `tgt_lang` for this model"""					)
											__magic_name__:	Any							=       src_lang
											__magic_name__:	Optional[int]							=       self(__snake_case     ,		add_special_tokens=__snake_case     ,		**__snake_case					)
											__magic_name__:	List[Any]							=       self.get_lang_id(__snake_case					)
											__magic_name__:	Tuple							=       tgt_lang_id
											return inputs

					def   lowerCamelCase__   (			self  :      List[str]					)			->							Optional[Any]:
											self.set_src_lang_special_tokens(self.src_lang					)

					def   lowerCamelCase__   (			self  :      Optional[int]					)			->							int:
											self.set_tgt_lang_special_tokens(self.tgt_lang					)

					def   lowerCamelCase__   (			self  :      int     ,		__snake_case  :      str					)			->							None:
											__magic_name__:	List[Any]							=       self.get_lang_token(__snake_case					)
											__magic_name__:	Optional[int]							=       self.lang_token_to_id[lang_token]
											__magic_name__:	Tuple							=       [self.cur_lang_id]
											__magic_name__:	Tuple							=       [self.eos_token_id]

					def   lowerCamelCase__   (			self  :      int     ,		__snake_case  :      str					)			->							None:
											__magic_name__:	List[Any]							=       self.get_lang_token(__snake_case					)
											__magic_name__:	List[Any]							=       self.lang_token_to_id[lang_token]
											__magic_name__:	Optional[Any]							=       [self.cur_lang_id]
											__magic_name__:	int							=       [self.eos_token_id]

					def   lowerCamelCase__   (			self  :      Dict     ,		__snake_case  :      str					)			->							str:
											return self.lang_code_to_token[lang]






					def   lowerCamelCase__   (			self  :      Union[str, Any]     ,		__snake_case  :      str					)			->							int:
											__magic_name__:	Optional[Any]							=       self.get_lang_token(__snake_case					)
											return self.lang_token_to_id[lang_token]



def   a    (     __UpperCAmelCase      :    str      ,   __UpperCAmelCase      :    Dict[str, Any]  )    ->							sentencepiece.SentencePieceProcessor:
						__magic_name__:	Union[str, Any]							=       sentencepiece.SentencePieceProcessor(**__UpperCAmelCase  )
						spm.Load(str(__UpperCAmelCase  )  )
						return spm



def   a    (     __UpperCAmelCase      :    str  )    ->							Union[Dict, List]:
						with open(__UpperCAmelCase      ,   """r"""  ) as f:
												return json.load(__UpperCAmelCase  )



def   a    (     __UpperCAmelCase      :    Optional[Any]      ,   __UpperCAmelCase      :    str  )    ->							None:
						with open(__UpperCAmelCase      ,   """w"""  ) as f:
												json.dump(__UpperCAmelCase      ,   __UpperCAmelCase      ,   indent=2  )











"""simple docstring"""





import argparse
import json
import os
import re

import torch

from transformers import BloomConfig, BloomModel
from transformers.file_utils import CONFIG_NAME, WEIGHTS_NAME
from transformers.utils import logging


logging.set_verbosity_info()

__lowerCamelCase							=  [
    'word_embeddings_layernorm.weight',
    'word_embeddings_layernorm.bias',
    'input_layernorm.weight',
    'input_layernorm.bias',
    'post_attention_layernorm.weight',
    'post_attention_layernorm.bias',
    'self_attention.dense.bias',
    'mlp.dense_4h_to_h.bias',
    'ln_f.weight',
    'ln_f.bias',
]

__lowerCamelCase							=  [
    'mlp.dense_4h_to_h.weight',
    'self_attention.dense.weight',
]



def   a    (     __UpperCAmelCase      :    Any      ,   __UpperCAmelCase      :    Any  )    ->							Any:
						__magic_name__:	List[str]							=       {
						    """word_embeddings.weight""": """word_embeddings.weight""",
						    """word_embeddings.norm.weight""": """word_embeddings_layernorm.weight""",
						    """word_embeddings.norm.bias""": """word_embeddings_layernorm.bias""",
						    """weight""": """ln_f.weight""",
						    """bias""": """ln_f.bias""",
						}

						if key in layer_rename_map:
												return layer_rename_map[key]

						# Handle transformer blocks
						__magic_name__:	Any							=       int(re.match(R""".*layer_(\d*).*"""      ,   __UpperCAmelCase  )[1]  )
						layer_number -= 3
						return f'h.{layer_number}.' + key



def   a    (     __UpperCAmelCase      :    List[str]  )    ->							Union[str, Any]:
						if dtype == torch.bool:
												return 1 / 8
						__magic_name__:	Optional[int]							=       re.search(R"""[^\d](\d+)$"""      ,   str(__UpperCAmelCase  )  )
						if bit_search is None:
												raise ValueError(f'`dtype` is not a valid dtype: {dtype}.'  )
						__magic_name__:	Dict							=       int(bit_search.groups()[0]  )
						return bit_size // 8



def   a    (     __UpperCAmelCase      :    Union[str, Any]      ,   __UpperCAmelCase      :    Optional[int]      ,   __UpperCAmelCase      :    List[Any]      ,   __UpperCAmelCase      :    Optional[int]      ,   __UpperCAmelCase      :    Tuple  )    ->							List[str]:
						# Construct model
						if bloom_config_file == "":
												__magic_name__:	Optional[Any]							=       BloomConfig()
						else:
												__magic_name__:	Tuple							=       BloomConfig.from_json_file(__UpperCAmelCase  )

						if shard_model:
												__magic_name__:	Union[str, Any]							=       os.listdir(__UpperCAmelCase  )
												__magic_name__:	Optional[int]							=       sorted(filter(lambda __UpperCAmelCase  : s.startswith("""layer"""  ) and "model_00" in s      ,   __UpperCAmelCase  )  )

												__magic_name__:	str							=       {"""weight_map""": {}, """metadata""": {}}
												__magic_name__:	Optional[Any]							=       0

												__magic_name__:	Dict							=       None

												__magic_name__:	Optional[int]							=       BloomConfig()

												for j, file in enumerate(__UpperCAmelCase  ):
																		print("""Processing file: {}""".format(__UpperCAmelCase  )  )
																		__magic_name__:	Tuple							=       None

																		for i in range(__UpperCAmelCase  ):
																								# load all TP files
																								__magic_name__:	int							=       file.replace("""model_00"""      ,   f'model_0{i}'  )
																								__magic_name__:	Optional[Any]							=       torch.load(os.path.join(__UpperCAmelCase      ,   __UpperCAmelCase  )      ,   map_location="""cpu"""  )

																								# Rename keys in the transformers names
																								__magic_name__:	str							=       list(temp.keys()  )
																								for key in keys:
																														__magic_name__:	Any							=       temp.pop(__UpperCAmelCase  )

																								if tensors is None:
																														__magic_name__:	Union[str, Any]							=       temp
																								else:
																														for key in tensors.keys():
																																				if any(key.endswith(__UpperCAmelCase  ) for end in WEIGHTS_TO_AVERAGE_ENDSWITH  ):
																																										# We average (sum and then divide) some weights accross TP ranks (see https://github.com/bigscience-workshop/Megatron-DeepSpeed/blob/olruwase/sync_layer_norms/megatron/training.py#L425)
																																										tensors[key] += temp[key]
																																				else:
																																										# Some weights are RowParallelLinear in Megatron-Deepspeed, others are ColumnParallel
																																										__magic_name__:	int							=       1 if any(text in key for text in WEIGHTS_WITH_ROW_PARALLELISM_CONTAIN  ) else 0
																																										# We concatenate these weights accross TP ranks
																																										__magic_name__:	int							=       torch.cat([tensors[key], temp[key]]      ,   dim=__UpperCAmelCase  )

            # Divide by the number of TP the weights we want to average
																		for key in tensors.keys():
																								if any(key.endswith(__UpperCAmelCase  ) for end in WEIGHTS_TO_AVERAGE_ENDSWITH  ):
																														__magic_name__:	Union[str, Any]							=       tensors[key] / pretraining_tp
																		torch.save(
																		    __UpperCAmelCase      ,   os.path.join(
																		        __UpperCAmelCase      ,   """pytorch_model_{}-of-{}.bin""".format(str(j + 1  ).zfill(5  )      ,   str(len(__UpperCAmelCase  )  ).zfill(5  )  )      ,   )      ,   )

																		for key in tensors.keys():
																								__magic_name__:	Optional[Any]							=       tensors[key]
																								total_size += value.numel() * get_dtype_size(value.dtype  )
																								if key not in index_dict["weight_map"]:
																														__magic_name__:	Union[str, Any]							=       """pytorch_model_{}-of-{}.bin""".format(
																														    str(j + 1  ).zfill(5  )      ,   str(len(__UpperCAmelCase  )  ).zfill(5  )  )

												__magic_name__:	Tuple							=       BloomConfig()
												__magic_name__:	Optional[int]							=       pytorch_dump_folder_path + """/""" + CONFIG_NAME
												__magic_name__:	Optional[int]							=       total_size
												with open(__UpperCAmelCase      ,   """w"""      ,   encoding="""utf-8"""  ) as f:
																		f.write(config.to_json_string()  )
												with open(os.path.join(__UpperCAmelCase      ,   WEIGHTS_NAME + """.index.json"""  )      ,   """w"""      ,   encoding="""utf-8"""  ) as f:
																		__magic_name__:	str							=       json.dumps(__UpperCAmelCase      ,   indent=2      ,   sort_keys=__UpperCAmelCase  ) + """\n"""
																		f.write(__UpperCAmelCase  )
						else:
												__magic_name__:	Optional[int]							=       BloomModel(__UpperCAmelCase  )

												__magic_name__:	List[str]							=       os.listdir(__UpperCAmelCase  )
												__magic_name__:	Optional[Any]							=       sorted(filter(lambda __UpperCAmelCase  : s.startswith("""layer"""  ) and "model_00" in s      ,   __UpperCAmelCase  )  )

												__magic_name__:	Any							=       None
												for i, file in enumerate(__UpperCAmelCase  ):
																		__magic_name__:	List[Any]							=       None
																		for i in range(__UpperCAmelCase  ):
																								# load all TP files
																								__magic_name__:	Any							=       file.replace("""model_00"""      ,   f'model_0{i}'  )
																								__magic_name__:	Dict							=       torch.load(os.path.join(__UpperCAmelCase      ,   __UpperCAmelCase  )      ,   map_location="""cpu"""  )

																								# Rename keys in the transformers names
																								__magic_name__:	Optional[Any]							=       list(temp.keys()  )
																								for key in keys:
																														__magic_name__:	Optional[Any]							=       temp.pop(__UpperCAmelCase  )

																								if tensors is None:
																														__magic_name__:	Tuple							=       temp
																								else:
																														for key in tensors.keys():
																																				# We average (sum and then divide) some weights accross TP ranks (see https://github.com/bigscience-workshop/Megatron-DeepSpeed/blob/olruwase/sync_layer_norms/megatron/training.py#L425)
																																				if any(key.endswith(__UpperCAmelCase  ) for end in WEIGHTS_TO_AVERAGE_ENDSWITH  ):
																																										tensors[key] += temp[key]
																																				else:
																																										# Some weights are RowParallelLinear in Megatron-Deepspeed, others are ColumnParallel
																																										__magic_name__:	Tuple							=       1 if any(text in key for text in WEIGHTS_WITH_ROW_PARALLELISM_CONTAIN  ) else 0
																																										# We concatenate these weights accross TP ranks
																																										__magic_name__:	Dict							=       torch.cat([tensors[key], temp[key]]      ,   dim=__UpperCAmelCase  )

            # Divide by the number of TP the weights we want to average
																		for key in tensors.keys():
																								if any(key.endswith(__UpperCAmelCase  ) for end in WEIGHTS_TO_AVERAGE_ENDSWITH  ):
																														__magic_name__:	Optional[int]							=       tensors[key] / pretraining_tp

																		__magic_name__:	Union[str, Any]							=       model.load_state_dict(__UpperCAmelCase      ,   strict=__UpperCAmelCase  )
																		assert not other_keys.unexpected_keys, f'The keys {other_keys.unexpected_keys} are unexpected'
																		if missing_keys is None:
																								__magic_name__:	Optional[Any]							=       set(other_keys.missing_keys  )
																		else:
																								__magic_name__:	Tuple							=       missing_keys.intersection(set(other_keys.missing_keys  )  )

												assert not missing_keys, f'The keys {missing_keys} are missing'

												# Save pytorch-model
												os.makedirs(__UpperCAmelCase      ,   exist_ok=__UpperCAmelCase  )
												__magic_name__:	int							=       pytorch_dump_folder_path + """/""" + WEIGHTS_NAME
												__magic_name__:	List[str]							=       pytorch_dump_folder_path + """/""" + CONFIG_NAME
												print(f'Save PyTorch model to {pytorch_weights_dump_path} with dtype {config.torch_dtype}'  )
												if config.torch_dtype is not None:
																		__magic_name__:	Optional[Any]							=       model.to(config.torch_dtype  )
												torch.save(model.state_dict()      ,   __UpperCAmelCase  )
												print(f'Save configuration file to {pytorch_config_dump_path}'  )
												with open(__UpperCAmelCase      ,   """w"""      ,   encoding="""utf-8"""  ) as f:
																		f.write(config.to_json_string()  )


if __name__ == "__main__":
							__lowerCamelCase							=  argparse.ArgumentParser()
							# Required parameters
							parser.add_argument(
							    '--bloom_checkpoint_path',
							    default=None,
							    type=str,
							    required=True,
							    help='Path to the Megatron-LM checkpoint path.',
							)
							parser.add_argument(
							    '--pytorch_dump_folder_path', default=None, type=str, required=True, help='Path to the output PyTorch model.'
							)
							parser.add_argument(
							    '--bloom_config_file',
							    default='',
							    type=str,
							    help=(
							        'An optional config json file corresponding to the pre-trained model. \n'
							        'This specifies the model architecture.'
							    ),
							)
							parser.add_argument(
							    '--shard_model',
							    action='store_true',
							    help='An optional setting to shard the output model \nThis enables sharding the converted checkpoint',
							)
							parser.add_argument(
							    '--pretraining_tp',
							    default=4,
							    type=int,
							    help='Pretraining TP rank that has been used when training the model in Megatron-LM \n',
							)
							__lowerCamelCase							=  parser.parse_args()
							convert_bloom_checkpoint_to_pytorch(
							    args.bloom_checkpoint_path,
							    args.bloom_config_file,
							    args.pytorch_dump_folder_path,
							    args.shard_model,
							    args.pretraining_tp,
							)



