


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        _lowerCamelCase(	lowercase__   , lowercase__	)						->	str:

		'''simple docstring'''







		__lowercase=			{
		    '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=			int(re.match(R'.*layer_(\d*).*'   , lowercase__	)[1]	)
		layer_number -= 3
		return F'h.{layer_number}.' + key







def        _lowerCamelCase(	lowercase__	)						->	List[str]:

		'''simple docstring'''







		if dtype == torch.bool:
				return 1 / 8
		__lowercase=			re.search(R'[^\d](\d+)$'   , str(lowercase__	)	)
		if bit_search is None:
				raise ValueError(F'`dtype` is not a valid dtype: {dtype}.'	)
		__lowercase=			int(bit_search.groups()[0]	)
		return bit_size // 8







def        _lowerCamelCase(	lowercase__   , lowercase__   , lowercase__   , lowercase__   , lowercase__	)						->	Optional[Any]:

		'''simple docstring'''







		if bloom_config_file == "":
				__lowercase=			BloomConfig()
		else:
				__lowercase=			BloomConfig.from_json_file(lowercase__	)

		if shard_model:
				__lowercase=			os.listdir(lowercase__	)
				__lowercase=			sorted(filter(lambda lowercase__	: s.startswith('layer'	) and "model_00" in s   , lowercase__	)	)

				__lowercase=			{'weight_map': {}, 'metadata': {}}
				__lowercase=			0

				__lowercase=			None

				__lowercase=			BloomConfig()

				for j, file in enumerate(lowercase__	):
						print('Processing file: {}'.format(lowercase__	)	)
						__lowercase=			None

						for i in range(lowercase__	):
								# load all TP files
								__lowercase=			file.replace('model_00'   , F'model_0{i}'	)
								__lowercase=			torch.load(os.path.join(lowercase__   , lowercase__	)   , map_location='cpu'	)

								# Rename keys in the transformers names
								__lowercase=			list(temp.keys()	)
								for key in keys:
										__lowercase=			temp.pop(lowercase__	)

								if tensors is None:
										__lowercase=			temp
								else:
										for key in tensors.keys():
												if any(key.endswith(lowercase__	) 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=			1 if any(text in key for text in WEIGHTS_WITH_ROW_PARALLELISM_CONTAIN	) else 0
														# We concatenate these weights accross TP ranks
														__lowercase=			torch.cat([tensors[key], temp[key]]   , dim=lowercase__	)

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

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

				__lowercase=			BloomConfig()
				__lowercase=			pytorch_dump_folder_path + '/' + CONFIG_NAME
				__lowercase=			total_size
				with open(lowercase__   , 'w'   , encoding='utf-8'	) as f:
						f.write(config.to_json_string()	)
				with open(os.path.join(lowercase__   , WEIGHTS_NAME + '.index.json'	)   , 'w'   , encoding='utf-8'	) as f:
						__lowercase=			json.dumps(lowercase__   , indent=2   , sort_keys=lowercase__	) + '\n'
						f.write(lowercase__	)
		else:
				__lowercase=			BloomModel(lowercase__	)

				__lowercase=			os.listdir(lowercase__	)
				__lowercase=			sorted(filter(lambda lowercase__	: s.startswith('layer'	) and "model_00" in s   , lowercase__	)	)

				__lowercase=			None
				for i, file in enumerate(lowercase__	):
						__lowercase=			None
						for i in range(lowercase__	):
								# load all TP files
								__lowercase=			file.replace('model_00'   , F'model_0{i}'	)
								__lowercase=			torch.load(os.path.join(lowercase__   , lowercase__	)   , map_location='cpu'	)

								# Rename keys in the transformers names
								__lowercase=			list(temp.keys()	)
								for key in keys:
										__lowercase=			temp.pop(lowercase__	)

								if tensors is None:
										__lowercase=			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(lowercase__	) for end in WEIGHTS_TO_AVERAGE_ENDSWITH	):
														tensors[key] += temp[key]
												else:
														# Some weights are RowParallelLinear in Megatron-Deepspeed, others are ColumnParallel
														__lowercase=			1 if any(text in key for text in WEIGHTS_WITH_ROW_PARALLELISM_CONTAIN	) else 0
														# We concatenate these weights accross TP ranks
														__lowercase=			torch.cat([tensors[key], temp[key]]   , dim=lowercase__	)

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

						__lowercase=			model.load_state_dict(lowercase__   , strict=lowercase__	)
						assert not other_keys.unexpected_keys, F'The keys {other_keys.unexpected_keys} are unexpected'
						if missing_keys is None:
								__lowercase=			set(other_keys.missing_keys	)
						else:
								__lowercase=			missing_keys.intersection(set(other_keys.missing_keys	)	)

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

				# Save pytorch-model
				os.makedirs(lowercase__   , exist_ok=lowercase__	)
				__lowercase=			pytorch_dump_folder_path + '/' + WEIGHTS_NAME
				__lowercase=			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=			model.to(config.torch_dtype	)
				torch.save(model.state_dict()   , lowercase__	)
				print(F'Save configuration file to {pytorch_config_dump_path}'	)
				with open(lowercase__   , '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,
	)










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

import sentencepiece as spm

from ...tokenization_utils import PreTrainedTokenizer
from ...utils import logging


lowerCAmelCase           =		'''▁'''

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

lowerCAmelCase           =		{
    '''vocab_file''': {'''google/pegasus-xsum''': '''https://huggingface.co/google/pegasus-xsum/resolve/main/spiece.model'''}
}

lowerCAmelCase           =		{
    '''google/pegasus-xsum''': 5_1_2,
}


lowerCAmelCase           =		logging.get_logger(__name__)





class       A    (   A_ ):
							UpperCamelCase_   :							Union[str, Any]       =VOCAB_FILES_NAMES

							UpperCamelCase_   :							List[Any]       =VOCAB_FILES_NAMES
							UpperCamelCase_   :							int       =PRETRAINED_VOCAB_FILES_MAP
							UpperCamelCase_   :							Tuple       =PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
							UpperCamelCase_   :							int       =['''input_ids''', '''attention_mask''']


							def __init__(self      ,				lowerCAmelCase      ,				lowerCAmelCase="<pad>"      ,				lowerCAmelCase="</s>"      ,				lowerCAmelCase="<unk>"      ,				lowerCAmelCase="<mask_2>"      ,				lowerCAmelCase="<mask_1>"      ,				lowerCAmelCase=None      ,				lowerCAmelCase=1_0_3      ,				lowerCAmelCase = None      ,				**lowerCAmelCase      ,				):
									__lowercase=			offset
									if additional_special_tokens is not None:
											if not isinstance(lowerCAmelCase      ,				lowerCAmelCase					):
													raise TypeError(
													    f'additional_special_tokens should be of type {type(lowerCAmelCase					)}, but is'
													    f' {type(lowerCAmelCase					)}'					)

											__lowercase=			(
											    ([mask_token_sent] + additional_special_tokens)
											    if mask_token_sent not in additional_special_tokens and mask_token_sent is not None
											    else additional_special_tokens
											)
											# fill additional tokens with ..., <unk_token_102> in case not all additional tokens are already taken
											additional_special_tokens_extended += [
											    f'<unk_{i}>' for i in range(len(lowerCAmelCase					)      ,				self.offset - 1					)
											]

											if len(set(lowerCAmelCase					)					) != len(lowerCAmelCase					):
													raise ValueError(
													    'Please make sure that the provided additional_special_tokens do not contain an incorrectly'
													    f' shifted list of <unk_x> tokens. Found {additional_special_tokens_extended}.'					)
											__lowercase=			additional_special_tokens_extended
									else:
											__lowercase=			[mask_token_sent] if mask_token_sent is not None else []
											additional_special_tokens += [f'<unk_{i}>' for i in range(2      ,				self.offset					)]

									__lowercase=			{} if sp_model_kwargs is None else sp_model_kwargs

									super().__init__(
									    eos_token=lowerCAmelCase      ,				unk_token=lowerCAmelCase      ,				mask_token=lowerCAmelCase      ,				pad_token=lowerCAmelCase      ,				mask_token_sent=lowerCAmelCase      ,				offset=lowerCAmelCase      ,				additional_special_tokens=lowerCAmelCase      ,				sp_model_kwargs=self.sp_model_kwargs      ,				**lowerCAmelCase      ,				)
									__lowercase=			mask_token_sent
									__lowercase=			vocab_file
									__lowercase=			spm.SentencePieceProcessor(**self.sp_model_kwargs					)
									self.sp_model.Load(lowerCAmelCase					)

									# add special tokens to encoder dict
									__lowercase=			{
									    0: self.pad_token,
									    1: self.eos_token,
									}

									if self.mask_token_sent is not None:
											self.encoder.update(
											    {
											        2: self.mask_token_sent,
											        3: self.mask_token,
											    }					)

									if self.offset > 0:
											# entries 2-104 are only used for pretraining and called <mask_1>, <mask_2>, unk_2, ...unk_102
											# mask_token_sent is already added to list -> so start at 1
											self.encoder.update({i + 3: additional_special_tokens[i] for i in range(1      ,				self.offset - 1					)}					)

									__lowercase=			{v: k for k, v in self.encoder.items()}


							@property
							def 			_A       (self					):
									return len(self.sp_model					) + self.offset


							def 			_A       (self					):
									__lowercase=			{self.convert_ids_to_tokens(lowerCAmelCase					): i for i in range(self.vocab_size					)}
									vocab.update(self.added_tokens_encoder					)
									return vocab


							def __getstate__(self					):
									__lowercase=			self.__dict__.copy()
									__lowercase=			None
									return state


							def __setstate__(self      ,				lowerCAmelCase					):
									__lowercase=			d

									# for backward compatibility
									if not hasattr(self      ,				'sp_model_kwargs'					):
											__lowercase=			{}

									__lowercase=			spm.SentencePieceProcessor(**self.sp_model_kwargs					)
									self.sp_model.Load(self.vocab_file					)


							def 			_A       (self      ,				lowerCAmelCase					):
									return self.sp_model.encode(lowerCAmelCase      ,				out_type=lowerCAmelCase					)


							def 			_A       (self      ,				lowerCAmelCase					):
									if token in self.decoder:
											return self.decoder[token]
									elif token in self.added_tokens_decoder:
											return self.added_tokens_decoder[token]
									__lowercase=			self.sp_model.piece_to_id(lowerCAmelCase					)
									return sp_id + self.offset


							def 			_A       (self      ,				lowerCAmelCase					):
									if index in self.encoder:
											return self.encoder[index]
									elif index in self.added_tokens_encoder:
											return self.added_tokens_encoder[index]
									else:
											__lowercase=			self.sp_model.IdToPiece(index - self.offset					)
									return token


							def 			_A       (self      ,				lowerCAmelCase					):
									__lowercase=			[]
									__lowercase=			''
									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(lowerCAmelCase					) + token
													__lowercase=			[]
											else:
													current_sub_tokens.append(lowerCAmelCase					)
									out_string += self.sp_model.decode(lowerCAmelCase					)
									return out_string.strip()


							def 			_A       (self      ,				lowerCAmelCase=False					):
									return 1


							def 			_A       (self      ,				lowerCAmelCase					):
									__lowercase=			set(self.all_special_ids					)  # call it once instead of inside list comp
									all_special_ids.remove(self.unk_token_id					)  # <unk> is only sometimes special

									return [1 if x in all_special_ids else 0 for x in seq]


							def 			_A       (self      ,				lowerCAmelCase      ,				lowerCAmelCase = None      ,				lowerCAmelCase = False					):
									if already_has_special_tokens:
											return self._special_token_mask(lowerCAmelCase					)
									elif token_ids_a is None:
											return self._special_token_mask(lowerCAmelCase					) + [1]
									else:
											return self._special_token_mask(token_ids_a + token_ids_a					) + [1]


							def 			_A       (self      ,				lowerCAmelCase      ,				lowerCAmelCase=None					):
									if token_ids_a is None:
											return token_ids_a + [self.eos_token_id]
									# We don't expect to process pairs, but leave the pair logic for API consistency
									return token_ids_a + token_ids_a + [self.eos_token_id]







							def 			_A       (self      ,				lowerCAmelCase      ,				lowerCAmelCase = None					):
									if not os.path.isdir(lowerCAmelCase					):
											logger.error(f'Vocabulary path ({save_directory}) should be a directory'					)
											return
									__lowercase=			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:
													__lowercase=			self.sp_model.serialized_model_proto()
													fi.write(lowerCAmelCase					)

									return (out_vocab_file,)






