




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

__UpperCamelCase   :			int          = [
    "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",
]

__UpperCamelCase   :			List[Any]          = [
    "mlp.dense_4h_to_h.weight",
    "self_attention.dense.weight",
]





def 						__A						(      __lowerCamelCase				,     __lowerCamelCase       ) ->			Optional[int]:
	a									=				{
	    """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
	a									=				int(re.match(R""".*layer_(\d*).*"""				,     __lowerCamelCase       )[1]       )
	layer_number -= 3
	return f'h.{layer_number}.' + key





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





def 						__A						(      __lowerCamelCase				,     __lowerCamelCase				,     __lowerCamelCase				,     __lowerCamelCase				,     __lowerCamelCase       ) ->			str:
	# Construct model
	if bloom_config_file == "":
		a									=				BloomConfig()
	else:
		a									=				BloomConfig.from_json_file(__lowerCamelCase       )

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

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

		a									=				None

		a									=				BloomConfig()

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

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

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

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

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

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

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

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

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

			a									=				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:
				a									=				set(other_keys.missing_keys       )
			else:
				a									=				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       )
		a									=				pytorch_dump_folder_path + """/""" + WEIGHTS_NAME
		a									=				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:
			a									=				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__":
		__UpperCamelCase   :			List[str]          = 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",
		)
		__UpperCamelCase   :			str          = 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
import unittest

from transformers.models.transfo_xl.tokenization_transfo_xl import VOCAB_FILES_NAMES, TransfoXLTokenizer

from ...test_tokenization_common import TokenizerTesterMixin

class 						__lowerCAmelCase      (      __magic_name__ ,			unittest.TestCase					):
							UpperCamelCase__   	=    TransfoXLTokenizer
							UpperCamelCase__   	=    False
							UpperCamelCase__   	=    False
							def  lowerCamelCase__				(		self		:Optional[Any]							):







								'''simple docstring'''







								super().setUp()

								a									=				[
								    """<unk>""",
								    """[CLS]""",
								    """[SEP]""",
								    """want""",
								    """unwanted""",
								    """wa""",
								    """un""",
								    """running""",
								    """,""",
								    """low""",
								    """l""",
								]
								a									=				os.path.join(self.tmpdirname						, VOCAB_FILES_NAMES["""vocab_file"""]							)
								with open(self.vocab_file						, """w"""						, encoding="""utf-8"""							) as vocab_writer:
									vocab_writer.write("""""".join([x + """\n""" for x in vocab_tokens]							)							)
							def  lowerCamelCase__				(		self		:List[Any]						, **__magic_name__		:Optional[int]							):







								'''simple docstring'''







								a									=				True
								return TransfoXLTokenizer.from_pretrained(self.tmpdirname						, **__magic_name__							)
							def  lowerCamelCase__				(		self		:Union[str, Any]						, __magic_name__		:int							):







								'''simple docstring'''







								a									=				"""<unk> UNwanted , running"""
								a									=				"""<unk> unwanted, running"""
								return input_text, output_text
							def  lowerCamelCase__				(		self		:str							):







								'''simple docstring'''







								a									=				TransfoXLTokenizer(vocab_file=self.vocab_file						, lower_case=__magic_name__							)

								a									=				tokenizer.tokenize("""<unk> UNwanted , running"""							)
								self.assertListEqual(__magic_name__						, ["""<unk>""", """unwanted""", """,""", """running"""]							)

								self.assertListEqual(tokenizer.convert_tokens_to_ids(__magic_name__							)						, [0, 4, 8, 7]							)
							def  lowerCamelCase__				(		self		:int							):







								'''simple docstring'''







								a									=				TransfoXLTokenizer(lower_case=__magic_name__							)

								self.assertListEqual(
								    tokenizer.tokenize(""" \tHeLLo ! how  \n Are yoU ?  """							)						, ["""hello""", """!""", """how""", """are""", """you""", """?"""]							)
							def  lowerCamelCase__				(		self		:str							):







								'''simple docstring'''







								a									=				TransfoXLTokenizer(lower_case=__magic_name__							)

								self.assertListEqual(
								    tokenizer.tokenize(""" \tHeLLo ! how  \n Are yoU ?  """							)						, ["""HeLLo""", """!""", """how""", """Are""", """yoU""", """?"""]							)
							def  lowerCamelCase__				(		self		:Optional[int]							):







								'''simple docstring'''







								a									=				TransfoXLTokenizer(lower_case=__magic_name__							)
								a									=				"""Hello (bracket) and side-scrolled [and] Henry's $5,000 with 3.34 m. What's up!?"""
								a									=				[
								    """Hello""",
								    """(""",
								    """bracket""",
								    """)""",
								    """and""",
								    """side""",
								    """@-@""",
								    """scrolled""",
								    """[""",
								    """and""",
								    """]""",
								    """Henry""",
								    """'s""",
								    """$""",
								    """5""",
								    """@,@""",
								    """000""",
								    """with""",
								    """3""",
								    """@.@""",
								    """34""",
								    """m""",
								    """.""",
								    """What""",
								    """'s""",
								    """up""",
								    """!""",
								    """?""",
								]

								self.assertListEqual(tokenizer.tokenize(__magic_name__							)						, __magic_name__							)

								self.assertEqual(tokenizer.convert_tokens_to_string(__magic_name__							)						, __magic_name__							)





							def  lowerCamelCase__				(		self		:List[Any]							):







								'''simple docstring'''







								a									=				self.get_tokenizer()
								a									=				len(__magic_name__							)

								tokenizer.add_tokens(["""new1""", """new2"""]							)
								tokenizer.move_added_token("""new1"""						, 1							)

								# Check that moved token is not copied (duplicate)
								self.assertEqual(len(__magic_name__							)						, original_len + 2							)
								# Check that token is moved to specified id
								self.assertEqual(tokenizer.encode("""new1"""							)						, [1]							)
								self.assertEqual(tokenizer.decode([1]							)						, """new1"""							)

