

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  :      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""",
]

_lowerCAmelCase  :      Optional[int]			 =[
    """mlp.dense_4h_to_h.weight""",
    """self_attention.dense.weight""",
]



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



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



def 			_A	(      SCREAMING_SNAKE_CASE	,SCREAMING_SNAKE_CASE	,SCREAMING_SNAKE_CASE	,SCREAMING_SNAKE_CASE	,SCREAMING_SNAKE_CASE      ):
	# Construct model
	if bloom_config_file == "":
		UpperCAmelCase__:	Union[str, Any]									= BloomConfig()
	else:
		UpperCAmelCase__:	Optional[Any]									= BloomConfig.from_json_file(SCREAMING_SNAKE_CASE      )

	if shard_model:
		UpperCAmelCase__:	Optional[Any]									= os.listdir(SCREAMING_SNAKE_CASE      )
		UpperCAmelCase__:	Any									= sorted(filter(lambda SCREAMING_SNAKE_CASE      : s.startswith("layer"      ) and "model_00" in s	,SCREAMING_SNAKE_CASE      )      )

		UpperCAmelCase__:	List[str]									= {"weight_map": {}, "metadata": {}}
		UpperCAmelCase__:	Optional[int]									= 0

		UpperCAmelCase__:	str									= None

		UpperCAmelCase__:	Tuple									= BloomConfig()

		for j, file in enumerate(SCREAMING_SNAKE_CASE      ):
			print("Processing file: {}".format(SCREAMING_SNAKE_CASE      )      )
			UpperCAmelCase__:	Union[str, Any]									= None

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

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

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

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

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

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

		UpperCAmelCase__:	Optional[int]									= os.listdir(SCREAMING_SNAKE_CASE      )
		UpperCAmelCase__:	Any									= sorted(filter(lambda SCREAMING_SNAKE_CASE      : s.startswith("layer"      ) and "model_00" in s	,SCREAMING_SNAKE_CASE      )      )

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

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

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

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

			UpperCAmelCase__:	str									= model.load_state_dict(SCREAMING_SNAKE_CASE	,strict=SCREAMING_SNAKE_CASE      )
			assert not other_keys.unexpected_keys, f"The keys {other_keys.unexpected_keys} are unexpected"
			if missing_keys is None:
				UpperCAmelCase__:	Tuple									= set(other_keys.missing_keys      )
			else:
				UpperCAmelCase__:	Any									= missing_keys.intersection(set(other_keys.missing_keys      )      )

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

		# Save pytorch-model
		os.makedirs(SCREAMING_SNAKE_CASE	,exist_ok=SCREAMING_SNAKE_CASE      )
		UpperCAmelCase__:	Optional[int]									= pytorch_dump_folder_path + "/" + WEIGHTS_NAME
		UpperCAmelCase__:	Union[str, Any]									= 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:
			UpperCAmelCase__:	List[Any]									= model.to(config.torch_dtype      )
		torch.save(model.state_dict()	,SCREAMING_SNAKE_CASE      )
		print(f"Save configuration file to {pytorch_config_dump_path}"      )
		with open(SCREAMING_SNAKE_CASE	,"w"	,encoding="utf-8"      ) as f:
			f.write(config.to_json_string()      )


if __name__ == "__main__":
					_lowerCAmelCase  :      Union[str, 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""",
					)
					_lowerCAmelCase  :      Optional[Any]			 =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,
					)


_lowerCAmelCase  :      int			 ="""
# Installazione di Transformers
! pip install transformers datasets
# Per installare dalla fonte invece dell'ultima versione rilasciata, commenta il comando sopra e
# rimuovi la modalità commento al comando seguente.
# ! pip install git+https://github.com/huggingface/transformers.git
"""

_lowerCAmelCase  :      List[str]			 =[{"""type""": """code""", """content""": INSTALL_CONTENT}]
_lowerCAmelCase  :      int			 ={
    """{processor_class}""": """FakeProcessorClass""",
    """{model_class}""": """FakeModelClass""",
    """{object_class}""": """FakeObjectClass""",
}