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

_lowercase						: List[str]						=						[
    "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",
]

_lowercase						: Tuple						=						[
    "mlp.dense_4h_to_h.weight",
    "self_attention.dense.weight",
]




def 							lowerCamelCase	(      UpperCAmelCase__       :					str      ,      UpperCAmelCase__       :					Dict )			->      str:
							lowercase_					:			int										=       {
							    """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
							lowercase_					:			Tuple										=       int(re.match(R""".*layer_(\d*).*"""      ,      UpperCAmelCase__ )[1] )
							layer_number -= 3
							return F'''h.{layer_number}.''' + key




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




def 							lowerCamelCase	(      UpperCAmelCase__       :					Any      ,      UpperCAmelCase__       :					Tuple      ,      UpperCAmelCase__       :					int      ,      UpperCAmelCase__       :					Tuple      ,      UpperCAmelCase__       :					Optional[int] )			->      List[Any]:
							# Construct model
							if bloom_config_file == "":
														lowercase_					:			Dict										=       BloomConfig()
							else:
														lowercase_					:			Dict										=       BloomConfig.from_json_file(UpperCAmelCase__ )

							if shard_model:
														lowercase_					:			Tuple										=       os.listdir(UpperCAmelCase__ )
														lowercase_					:			Any										=       sorted(filter(lambda UpperCAmelCase__ : s.startswith("""layer""" ) and "model_00" in s      ,      UpperCAmelCase__ ) )

														lowercase_					:			List[Any]										=       {"""weight_map""": {}, """metadata""": {}}
														lowercase_					:			int										=       0

														lowercase_					:			str										=       None

														lowercase_					:			Optional[Any]										=       BloomConfig()

														for j, file in enumerate(UpperCAmelCase__ ):
																					print("""Processing file: {}""".format(UpperCAmelCase__ ) )
																					lowercase_					:			Any										=       None

																					for i in range(UpperCAmelCase__ ):
																												# load all TP files
																												lowercase_					:			Any										=       file.replace("""model_00"""      ,      F'''model_0{i}''' )
																												lowercase_					:			List[str]										=       torch.load(os.path.join(UpperCAmelCase__      ,      UpperCAmelCase__ )      ,      map_location="""cpu""" )

																												# Rename keys in the transformers names
																												lowercase_					:			Union[str, Any]										=       list(temp.keys() )
																												for key in keys:
																																			lowercase_					:			Union[str, Any]										=       temp.pop(UpperCAmelCase__ )

																												if tensors is None:
																																			lowercase_					:			Tuple										=       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
																																																	lowercase_					:			Union[str, Any]										=       1 if any(text in key for text in WEIGHTS_WITH_ROW_PARALLELISM_CONTAIN ) else 0
																																																	# We concatenate these weights accross TP ranks
																																																	lowercase_					:			Union[str, Any]										=       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 ):
																																			lowercase_					:			Optional[int]										=       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():
																												lowercase_					:			int										=       tensors[key]
																												total_size += value.numel() * get_dtype_size(value.dtype )
																												if key not in index_dict["weight_map"]:
																																			lowercase_					:			str										=       """pytorch_model_{}-of-{}.bin""".format(
																																			    str(j + 1 ).zfill(5 )      ,      str(len(UpperCAmelCase__ ) ).zfill(5 ) )

														lowercase_					:			str										=       BloomConfig()
														lowercase_					:			Any										=       pytorch_dump_folder_path + """/""" + CONFIG_NAME
														lowercase_					:			Any										=       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:
																					lowercase_					:			Tuple										=       json.dumps(UpperCAmelCase__      ,      indent=2      ,      sort_keys=UpperCAmelCase__ ) + """\n"""
																					f.write(UpperCAmelCase__ )
							else:
														lowercase_					:			Any										=       BloomModel(UpperCAmelCase__ )

														lowercase_					:			List[str]										=       os.listdir(UpperCAmelCase__ )
														lowercase_					:			Dict										=       sorted(filter(lambda UpperCAmelCase__ : s.startswith("""layer""" ) and "model_00" in s      ,      UpperCAmelCase__ ) )

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

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

																												if tensors is None:
																																			lowercase_					:			Any										=       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
																																																	lowercase_					:			Tuple										=       1 if any(text in key for text in WEIGHTS_WITH_ROW_PARALLELISM_CONTAIN ) else 0
																																																	# We concatenate these weights accross TP ranks
																																																	lowercase_					:			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 ):
																																			lowercase_					:			Optional[Any]										=       tensors[key] / pretraining_tp

																					lowercase_					:			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:
																												lowercase_					:			List[Any]										=       set(other_keys.missing_keys )
																					else:
																												lowercase_					:			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__ )
														lowercase_					:			str										=       pytorch_dump_folder_path + """/""" + WEIGHTS_NAME
														lowercase_					:			Optional[int]										=       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:
																					lowercase_					:			Tuple										=       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__":
						_lowercase						: Any						=						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",
						)
						_lowercase						: Optional[int]						=						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,
						)






'''simple docstring'''




import os
from shutil import copyfile
from typing import List, Optional, Tuple

from ...tokenization_utils import AddedToken
from ...tokenization_utils_fast import PreTrainedTokenizerFast
from ...utils import is_sentencepiece_available, logging


if is_sentencepiece_available():
						from .tokenization_albert import AlbertTokenizer
else:
						_lowercase						: int						=						None

_lowercase						: Union[str, Any]						=						logging.get_logger(__name__)
_lowercase						: Tuple						=						{"vocab_file": "spiece.model", "tokenizer_file": "tokenizer.json"}

_lowercase						: Dict						=						{
    "vocab_file": {
        "albert-base-v1": "https://huggingface.co/albert-base-v1/resolve/main/spiece.model",
        "albert-large-v1": "https://huggingface.co/albert-large-v1/resolve/main/spiece.model",
        "albert-xlarge-v1": "https://huggingface.co/albert-xlarge-v1/resolve/main/spiece.model",
        "albert-xxlarge-v1": "https://huggingface.co/albert-xxlarge-v1/resolve/main/spiece.model",
        "albert-base-v2": "https://huggingface.co/albert-base-v2/resolve/main/spiece.model",
        "albert-large-v2": "https://huggingface.co/albert-large-v2/resolve/main/spiece.model",
        "albert-xlarge-v2": "https://huggingface.co/albert-xlarge-v2/resolve/main/spiece.model",
        "albert-xxlarge-v2": "https://huggingface.co/albert-xxlarge-v2/resolve/main/spiece.model",
    },
    "tokenizer_file": {
        "albert-base-v1": "https://huggingface.co/albert-base-v1/resolve/main/tokenizer.json",
        "albert-large-v1": "https://huggingface.co/albert-large-v1/resolve/main/tokenizer.json",
        "albert-xlarge-v1": "https://huggingface.co/albert-xlarge-v1/resolve/main/tokenizer.json",
        "albert-xxlarge-v1": "https://huggingface.co/albert-xxlarge-v1/resolve/main/tokenizer.json",
        "albert-base-v2": "https://huggingface.co/albert-base-v2/resolve/main/tokenizer.json",
        "albert-large-v2": "https://huggingface.co/albert-large-v2/resolve/main/tokenizer.json",
        "albert-xlarge-v2": "https://huggingface.co/albert-xlarge-v2/resolve/main/tokenizer.json",
        "albert-xxlarge-v2": "https://huggingface.co/albert-xxlarge-v2/resolve/main/tokenizer.json",
    },
}

_lowercase						: int						=						{
    "albert-base-v1": 512,
    "albert-large-v1": 512,
    "albert-xlarge-v1": 512,
    "albert-xxlarge-v1": 512,
    "albert-base-v2": 512,
    "albert-large-v2": 512,
    "albert-xlarge-v2": 512,
    "albert-xxlarge-v2": 512,
}

_lowercase						: Tuple						=						"▁"





class     __magic_name__						(    _UpperCAmelCase):
			UpperCamelCase__				=       VOCAB_FILES_NAMES
			UpperCamelCase__				=       PRETRAINED_VOCAB_FILES_MAP
			UpperCamelCase__				=       PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
			UpperCamelCase__				=       AlbertTokenizer







			def __init__(	self				:  Optional[Any]				,			lowercase_				:  Union[str, Any]=None				,			lowercase_				:  str=None				,			lowercase_				:  Any=True				,			lowercase_				:  Optional[int]=True				,			lowercase_				:  List[str]=False				,			lowercase_				:  Optional[int]="[CLS]"				,			lowercase_				:  Any="[SEP]"				,			lowercase_				:  int="<unk>"				,			lowercase_				:  Any="[SEP]"				,			lowercase_				:  int="<pad>"				,			lowercase_				:  Tuple="[CLS]"				,			lowercase_				:  Dict="[MASK]"				,			**lowercase_				:  Optional[int]				,			):
										# Mask token behave like a normal word, i.e. include the space before it and
										# is included in the raw text, there should be a match in a non-normalized sentence.
										lowercase_					:			Tuple										=       (
										    AddedToken(lowercase_				,			lstrip=lowercase_				,			rstrip=lowercase_				,			normalized=lowercase_    )
										    if isinstance(lowercase_				,			lowercase_    )
										    else mask_token
										)

										super().__init__(
										    lowercase_				,			tokenizer_file=lowercase_				,			do_lower_case=lowercase_				,			remove_space=lowercase_				,			keep_accents=lowercase_				,			bos_token=lowercase_				,			eos_token=lowercase_				,			unk_token=lowercase_				,			sep_token=lowercase_				,			pad_token=lowercase_				,			cls_token=lowercase_				,			mask_token=lowercase_				,			**lowercase_				,			)

										lowercase_					:			Optional[int]										=       do_lower_case
										lowercase_					:			Any										=       remove_space
										lowercase_					:			Dict										=       keep_accents
										lowercase_					:			Optional[int]										=       vocab_file
										lowercase_					:			Any										=       False if not self.vocab_file else True







			def 	SCREAMING_SNAKE_CASE_     (	self				:  Tuple				,			lowercase_				:  List[int]				,			lowercase_				:  Optional[List[int]] = None    ):
										lowercase_					:			Tuple										=       [self.sep_token_id]
										lowercase_					:			str										=       [self.cls_token_id]
										if token_ids_a is None:
																	return cls + token_ids_a + sep
										return cls + token_ids_a + sep + token_ids_a + sep







			def 	SCREAMING_SNAKE_CASE_     (	self				:  Any				,			lowercase_				:  List[int]				,			lowercase_				:  Optional[List[int]] = None    ):
										lowercase_					:			Union[str, Any]										=       [self.sep_token_id]
										lowercase_					:			int										=       [self.cls_token_id]

										if token_ids_a is None:
																	return len(cls + token_ids_a + sep    ) * [0]
										return len(cls + token_ids_a + sep    ) * [0] + len(token_ids_a + sep    ) * [1]







			def 	SCREAMING_SNAKE_CASE_     (	self				:  Union[str, Any]				,			lowercase_				:  str				,			lowercase_				:  Optional[str] = None    ):
										if not self.can_save_slow_tokenizer:
																	raise ValueError(
																	    """Your fast tokenizer does not have the necessary information to save the vocabulary for a slow """
																	    """tokenizer."""    )

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

										if os.path.abspath(self.vocab_file    ) != os.path.abspath(lowercase_    ):
																	copyfile(self.vocab_file				,			lowercase_    )

										return (out_vocab_file,)





