
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 							__A(lowerCAmelCase						,  lowerCAmelCase )		->				Any:

		"""simple docstring"""






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


def 							__A(lowerCAmelCase )		->				Optional[int]:

		"""simple docstring"""






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


def 							__A(lowerCAmelCase						,  lowerCAmelCase						,  lowerCAmelCase						,  lowerCAmelCase						,  lowerCAmelCase )		->				str:

		"""simple docstring"""






		if bloom_config_file == "":
				_UpperCamelCase						    =	BloomConfig()
		else:
				_UpperCamelCase						    =	BloomConfig.from_json_file(lowerCAmelCase )

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

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

				_UpperCamelCase						    =	None

				_UpperCamelCase						    =	BloomConfig()

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

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

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

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

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

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

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

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

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

						_UpperCamelCase						    =	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:
								_UpperCamelCase						    =	set(other_keys.missing_keys )
						else:
								_UpperCamelCase						    =	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 )
				_UpperCamelCase						    =	pytorch_dump_folder_path + """/""" + WEIGHTS_NAME
				_UpperCamelCase						    =	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						    =	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__":
	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,
	)





from typing import TYPE_CHECKING

from ....utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available


lowerCamelCase__	=							{
    "configuration_mctct": ["MCTCT_PRETRAINED_CONFIG_ARCHIVE_MAP", "MCTCTConfig"],
    "feature_extraction_mctct": ["MCTCTFeatureExtractor"],
    "processing_mctct": ["MCTCTProcessor"],
}


try:
		if not is_torch_available():
				raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
		pass
else:
		lowerCamelCase__	=							[
		    "MCTCT_PRETRAINED_MODEL_ARCHIVE_LIST",
		    "MCTCTForCTC",
		    "MCTCTModel",
		    "MCTCTPreTrainedModel",
		]


if TYPE_CHECKING:
		from .configuration_mctct import MCTCT_PRETRAINED_CONFIG_ARCHIVE_MAP, MCTCTConfig
		from .feature_extraction_mctct import MCTCTFeatureExtractor
		from .processing_mctct import MCTCTProcessor

		try:
				if not is_torch_available():
						raise OptionalDependencyNotAvailable()
		except OptionalDependencyNotAvailable:
				pass
		else:
				from .modeling_mctct import MCTCT_PRETRAINED_MODEL_ARCHIVE_LIST, MCTCTForCTC, MCTCTModel, MCTCTPreTrainedModel

else:
		import sys

		lowerCamelCase__	=							_LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)



