
"""simple docstring"""







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       (_UpperCamelCase ,						_UpperCamelCase	):
	__lowerCAmelCase    :		List[Any]		    =   {
	    '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    :		Optional[Any]		    =   int(re.match(r'.*layer_(\d*).*' ,						__UpperCAmelCase	)[1]	)
	layer_number -= 3
	return F"h.{layer_number}." + key
def 					__lowerCAmelCase       (_UpperCamelCase	):







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







	if bloom_config_file == "":
		__lowerCAmelCase    :		Optional[Any]		    =   BloomConfig()
	else:
		__lowerCAmelCase    :		Tuple		    =   BloomConfig.from_json_file(__UpperCAmelCase	)

	if shard_model:
		__lowerCAmelCase    :		Tuple		    =   os.listdir(__UpperCAmelCase	)
		__lowerCAmelCase    :		Dict		    =   sorted(filter(lambda _UpperCamelCase	: s.startswith('layer'	) and "model_00" in s ,						__UpperCAmelCase	)	)

		__lowerCAmelCase    :		Union[str, Any]		    =   {'weight_map': {}, 'metadata': {}}
		__lowerCAmelCase    :		Any		    =   0

		__lowerCAmelCase    :		Optional[Any]		    =   None

		__lowerCAmelCase    :		Tuple		    =   BloomConfig()

		for j, file in enumerate(__UpperCAmelCase	):
			print('Processing file: {}'.format(__UpperCAmelCase	)	)
			__lowerCAmelCase    :		int		    =   None

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

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

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

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

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

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

		__lowerCAmelCase    :		List[Any]		    =   os.listdir(__UpperCAmelCase	)
		__lowerCAmelCase    :		int		    =   sorted(filter(lambda _UpperCamelCase	: s.startswith('layer'	) and "model_00" in s ,						__UpperCAmelCase	)	)

		__lowerCAmelCase    :		List[str]		    =   None
		for i, file in enumerate(__UpperCAmelCase	):
			__lowerCAmelCase    :		Tuple		    =   None
			for i in range(__UpperCAmelCase	):
				# load all TP files
				__lowerCAmelCase    :		List[Any]		    =   file.replace('model_00' ,						F"model_0{i}"	)
				__lowerCAmelCase    :		Optional[Any]		    =   torch.load(os.path.join(__UpperCAmelCase ,						__UpperCAmelCase	) ,						map_location='cpu'	)

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

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

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

			__lowerCAmelCase    :		Union[str, Any]		    =   model.load_state_dict(__UpperCAmelCase ,						strict=__UpperCAmelCase	)
			assert not other_keys.unexpected_keys, F"The keys {other_keys.unexpected_keys} are unexpected"
			if missing_keys is None:
				__lowerCAmelCase    :		Any		    =   set(other_keys.missing_keys	)
			else:
				__lowerCAmelCase    :		Union[str, 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(__UpperCAmelCase ,						exist_ok=__UpperCAmelCase	)
		__lowerCAmelCase    :		int		    =   pytorch_dump_folder_path + '/' + WEIGHTS_NAME
		__lowerCAmelCase    :		Tuple		    =   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    :		Dict		    =   model.to(config.torch_dtype	)
		torch.save(model.state_dict() ,						__UpperCAmelCase	)
		print(F"Save configuration file to {pytorch_config_dump_path}"	)
		with open(__UpperCAmelCase ,						'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,
					)

"""simple docstring"""







from ....configuration_utils import PretrainedConfig
from ....utils import logging


lowerCamelCase__           =	logging.get_logger(__name__)

lowerCamelCase__           =	{
    """speechbrain/m-ctc-t-large""": """https://huggingface.co/speechbrain/m-ctc-t-large/resolve/main/config.json""",
    # See all M-CTC-T models at https://huggingface.co/models?filter=mctct
}




class 				A__ (			_lowerCamelCase):
		A_		: Any            =			'mctct'

		def __init__(   self ,	_SCREAMING_SNAKE_CASE=80_65 ,	_SCREAMING_SNAKE_CASE=15_36 ,	_SCREAMING_SNAKE_CASE=36 ,	_SCREAMING_SNAKE_CASE=61_44 ,	_SCREAMING_SNAKE_CASE=4 ,	_SCREAMING_SNAKE_CASE=3_84 ,	_SCREAMING_SNAKE_CASE=9_20 ,	_SCREAMING_SNAKE_CASE=1E-5 ,	_SCREAMING_SNAKE_CASE=0.3 ,	_SCREAMING_SNAKE_CASE="relu" ,	_SCREAMING_SNAKE_CASE=0.02 ,	_SCREAMING_SNAKE_CASE=0.3 ,	_SCREAMING_SNAKE_CASE=0.3 ,	_SCREAMING_SNAKE_CASE=1 ,	_SCREAMING_SNAKE_CASE=0 ,	_SCREAMING_SNAKE_CASE=2 ,	_SCREAMING_SNAKE_CASE=1 ,	_SCREAMING_SNAKE_CASE=0.3 ,	_SCREAMING_SNAKE_CASE=1 ,	_SCREAMING_SNAKE_CASE=(7,) ,	_SCREAMING_SNAKE_CASE=(3,) ,	_SCREAMING_SNAKE_CASE=80 ,	_SCREAMING_SNAKE_CASE=1 ,	_SCREAMING_SNAKE_CASE=None ,	_SCREAMING_SNAKE_CASE="sum" ,	_SCREAMING_SNAKE_CASE=False ,	**_SCREAMING_SNAKE_CASE ,	):
			super().__init__(**_SCREAMING_SNAKE_CASE ,	pad_token_id=_SCREAMING_SNAKE_CASE ,	bos_token_id=_SCREAMING_SNAKE_CASE ,	eos_token_id=_SCREAMING_SNAKE_CASE				)
			__lowerCAmelCase    :		str		    =   vocab_size
			__lowerCAmelCase    :		Any		    =   hidden_size
			__lowerCAmelCase    :		int		    =   num_hidden_layers
			__lowerCAmelCase    :		Any		    =   intermediate_size
			__lowerCAmelCase    :		Optional[int]		    =   num_attention_heads
			__lowerCAmelCase    :		str		    =   attention_head_dim
			__lowerCAmelCase    :		int		    =   max_position_embeddings
			__lowerCAmelCase    :		List[Any]		    =   layer_norm_eps
			__lowerCAmelCase    :		Tuple		    =   layerdrop
			__lowerCAmelCase    :		Dict		    =   hidden_act
			__lowerCAmelCase    :		Dict		    =   initializer_range
			__lowerCAmelCase    :		Any		    =   hidden_dropout_prob
			__lowerCAmelCase    :		Dict		    =   attention_probs_dropout_prob
			__lowerCAmelCase    :		List[str]		    =   pad_token_id
			__lowerCAmelCase    :		Union[str, Any]		    =   bos_token_id
			__lowerCAmelCase    :		int		    =   eos_token_id
			__lowerCAmelCase    :		Optional[int]		    =   conv_glu_dim
			__lowerCAmelCase    :		List[Any]		    =   conv_dropout
			__lowerCAmelCase    :		Tuple		    =   num_conv_layers
			__lowerCAmelCase    :		Tuple		    =   input_feat_per_channel
			__lowerCAmelCase    :		Optional[int]		    =   input_channels
			__lowerCAmelCase    :		str		    =   conv_channels
			__lowerCAmelCase    :		str		    =   ctc_loss_reduction
			__lowerCAmelCase    :		Tuple		    =   ctc_zero_infinity

			# prevents config testing fail with exporting to json
			__lowerCAmelCase    :		List[str]		    =   list(_SCREAMING_SNAKE_CASE				)
			__lowerCAmelCase    :		Dict		    =   list(_SCREAMING_SNAKE_CASE				)

			if len(self.conv_kernel				) != self.num_conv_layers:
				raise ValueError(
				    'Configuration for convolutional module is incorrect. '
				    'It is required that `len(config.conv_kernel)` == `config.num_conv_layers` '
				    f"but is `len(config.conv_kernel) = {len(self.conv_kernel				)}`, "
				    f"`config.num_conv_layers = {self.num_conv_layers}`."				)