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

_A   =     [
    "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",
]

_A   =     [
    "mlp.dense_4h_to_h.weight",
    "self_attention.dense.weight",
]






def    lowerCamelCase__				(				__lowerCAmelCase  :  Dict    ,							__lowerCAmelCase  :  Dict				):




			"""simple docstring"""


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






def    lowerCamelCase__				(				__lowerCAmelCase  :  Dict				):




			"""simple docstring"""


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






def    lowerCamelCase__				(				__lowerCAmelCase  :  List[str]    ,							__lowerCAmelCase  :  Tuple    ,							__lowerCAmelCase  :  str    ,							__lowerCAmelCase  :  List[str]    ,							__lowerCAmelCase  :  Dict				):




			"""simple docstring"""


			if bloom_config_file == "":
						lowerCAmelCase_						=	BloomConfig()
			else:
						lowerCAmelCase_						=	BloomConfig.from_json_file(__lowerCAmelCase				)

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

						lowerCAmelCase_						=	{"weight_map": {}, "metadata": {}}
						lowerCAmelCase_						=	0

						lowerCAmelCase_						=	None

						lowerCAmelCase_						=	BloomConfig()

						for j, file in enumerate(__lowerCAmelCase				):
									print("Processing file: {}".format(__lowerCAmelCase				)				)
									lowerCAmelCase_						=	None

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

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

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

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

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

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

						lowerCAmelCase_						=	os.listdir(__lowerCAmelCase				)
						lowerCAmelCase_						=	sorted(filter(lambda __lowerCAmelCase				: s.startswith("layer"				) and "model_00" in s    ,							__lowerCAmelCase				)				)

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

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

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

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

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

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

						# Save pytorch-model
						os.makedirs(__lowerCAmelCase    ,							exist_ok=__lowerCAmelCase				)
						lowerCAmelCase_						=	pytorch_dump_folder_path + "/" + WEIGHTS_NAME
						lowerCAmelCase_						=	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:
									lowerCAmelCase_						=	model.to(config.torch_dtype				)
						torch.save(model.state_dict()    ,							__lowerCAmelCase				)
						print(F"""Save configuration file to {pytorch_config_dump_path}"""				)
						with open(__lowerCAmelCase    ,							"w"    ,							encoding="utf-8"				) as f:
									f.write(config.to_json_string()				)


if __name__ == "__main__":
			_A   =     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",
			)
			_A   =     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 unittest

from datasets import load_dataset

from transformers import BloomTokenizerFast
from transformers.testing_utils import require_tokenizers

from ...test_tokenization_common import TokenizerTesterMixin





@require_tokenizers
class 	_lowerCAmelCase					(	__a					,      unittest.TestCase       ):
				_lowercase					=None
				_lowercase					=BloomTokenizerFast
				_lowercase					=BloomTokenizerFast
				_lowercase					=True
				_lowercase					=False
				_lowercase					='''tokenizer_file'''
				_lowercase					={'''bos_token''': '''<s>''', '''eos_token''': '''</s>''', '''unk_token''': '''<unk>''', '''pad_token''': '''<pad>'''}




				def      __a						(					self     )				->      Dict:
							super().setUp()
							lowerCAmelCase_						=	BloomTokenizerFast.from_pretrained("bigscience/tokenizer"     )
							tokenizer.save_pretrained(self.tmpdirname     )




				def      __a						(					self				,				**_UpperCamelCase     )				->      Tuple:
							kwargs.update(self.special_tokens_map     )
							return BloomTokenizerFast.from_pretrained(self.tmpdirname				,				**_UpperCamelCase     )




				def      __a						(					self     )				->      List[str]:
							lowerCAmelCase_						=	self.get_rust_tokenizer()

							lowerCAmelCase_						=	["The quick brown fox</s>", "jumps over the lazy dog</s>"]
							lowerCAmelCase_						=	[[2_175, 23_714, 73_173, 144_252, 2], [77, 132_619, 3_478, 368, 109_586, 35_433, 2]]

							lowerCAmelCase_						=	tokenizer.batch_encode_plus(_UpperCamelCase     )["input_ids"]
							self.assertListEqual(_UpperCamelCase				,				_UpperCamelCase     )

							lowerCAmelCase_						=	tokenizer.batch_decode(_UpperCamelCase     )
							self.assertListEqual(_UpperCamelCase				,				_UpperCamelCase     )




				def      __a						(					self				,				_UpperCamelCase=6     )				->      Tuple:
							for tokenizer, pretrained_name, kwargs in self.tokenizers_list:
										with self.subTest(f"""{tokenizer.__class__.__name__} ({pretrained_name})"""     ):
													lowerCAmelCase_						=	self.rust_tokenizer_class.from_pretrained(_UpperCamelCase				,				**_UpperCamelCase     )
													# tokenizer_r.pad_token = None # Hotfixing padding = None
													# Simple input
													lowerCAmelCase_						=	"This is a simple input"
													lowerCAmelCase_						=	["This is a simple input 1", "This is a simple input 2"]
													lowerCAmelCase_						=	("This is a simple input", "This is a pair")
													lowerCAmelCase_						=	[
													    ("This is a simple input 1", "This is a simple input 2"),
													    ("This is a simple pair 1", "This is a simple pair 2"),
													]

													# Simple input tests
													try:
																tokenizer_r.encode(_UpperCamelCase				,				max_length=_UpperCamelCase     )
																tokenizer_r.encode_plus(_UpperCamelCase				,				max_length=_UpperCamelCase     )

																tokenizer_r.batch_encode_plus(_UpperCamelCase				,				max_length=_UpperCamelCase     )
																tokenizer_r.encode(_UpperCamelCase				,				max_length=_UpperCamelCase     )
																tokenizer_r.batch_encode_plus(_UpperCamelCase				,				max_length=_UpperCamelCase     )
													except ValueError:
																self.fail("Bloom Tokenizer should be able to deal with padding"     )

													lowerCAmelCase_						=	None  # Hotfixing padding = None
													self.assertRaises(_UpperCamelCase				,				tokenizer_r.encode				,				_UpperCamelCase				,				max_length=_UpperCamelCase				,				padding="max_length"     )

													# Simple input
													self.assertRaises(_UpperCamelCase				,				tokenizer_r.encode_plus				,				_UpperCamelCase				,				max_length=_UpperCamelCase				,				padding="max_length"     )

													# Simple input
													self.assertRaises(
													    _UpperCamelCase				,				tokenizer_r.batch_encode_plus				,				_UpperCamelCase				,				max_length=_UpperCamelCase				,				padding="max_length"				,				)

													# Pair input
													self.assertRaises(_UpperCamelCase				,				tokenizer_r.encode				,				_UpperCamelCase				,				max_length=_UpperCamelCase				,				padding="max_length"     )

													# Pair input
													self.assertRaises(_UpperCamelCase				,				tokenizer_r.encode_plus				,				_UpperCamelCase				,				max_length=_UpperCamelCase				,				padding="max_length"     )

													# Pair input
													self.assertRaises(
													    _UpperCamelCase				,				tokenizer_r.batch_encode_plus				,				_UpperCamelCase				,				max_length=_UpperCamelCase				,				padding="max_length"				,				)




				def      __a						(					self     )				->      Any:
							lowerCAmelCase_						=	self.get_rust_tokenizer()
							lowerCAmelCase_						=	load_dataset("xnli"				,				"all_languages"				,				split="test"				,				streaming=_UpperCamelCase     )

							lowerCAmelCase_						=	next(iter(_UpperCamelCase     )     )["premise"]  # pick up one data
							lowerCAmelCase_						=	list(sample_data.values()     )

							lowerCAmelCase_						=	list(map(tokenizer.encode				,				_UpperCamelCase     )     )
							lowerCAmelCase_						=	[tokenizer.decode(_UpperCamelCase				,				clean_up_tokenization_spaces=_UpperCamelCase     ) for x in output_tokens]
							self.assertListEqual(_UpperCamelCase				,				_UpperCamelCase     )




				def      __a						(					self     )				->      List[Any]:
							# The test has to be overriden because BLOOM uses ALiBi positional embeddings that does not have
							# any sequence length constraints. This test of the parent class will fail since it relies on the
							# maximum sequence length of the positoonal embeddings.
							self.assertGreaterEqual(len(self.tokenizer_class.pretrained_vocab_files_map     )				,				1     )
							self.assertGreaterEqual(len(list(self.tokenizer_class.pretrained_vocab_files_map.values()     )[0]     )				,				1     )



