
"""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__ :       Tuple						     =			[
    "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__ :       Dict						     =			[
    "mlp.dense_4h_to_h.weight",
    "self_attention.dense.weight",
]





def __A						(   a_ :							str		,		a_ :							Tuple     )->							Optional[int]:




		'''simple docstring'''






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





def __A						(   a_ :							Optional[Any]     )->							int:




		'''simple docstring'''






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





def __A						(   a_ :							str		,		a_ :							Dict		,		a_ :							Dict		,		a_ :							List[str]		,		a_ :							int     )->							Dict:




		'''simple docstring'''






		if bloom_config_file == "":
				SCREAMING_SNAKE_CASE       :     List[str]				= BloomConfig()
		else:
				SCREAMING_SNAKE_CASE       :     Union[str, Any]				= BloomConfig.from_json_file(a_     )

		if shard_model:
				SCREAMING_SNAKE_CASE       :     Union[str, Any]				= os.listdir(a_     )
				SCREAMING_SNAKE_CASE       :     List[Any]				= sorted(filter(lambda a_     : s.startswith('''layer'''     ) and "model_00" in s		,		a_     )     )

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

				SCREAMING_SNAKE_CASE       :     List[str]				= None

				SCREAMING_SNAKE_CASE       :     List[Any]				= BloomConfig()

				for j, file in enumerate(a_     ):
						print('''Processing file: {}'''.format(a_     )     )
						SCREAMING_SNAKE_CASE       :     List[str]				= None

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

								# Rename keys in the transformers names
								SCREAMING_SNAKE_CASE       :     Dict				= list(temp.keys()     )
								for key in keys:
										SCREAMING_SNAKE_CASE       :     Any				= temp.pop(a_     )

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

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

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

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

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

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

								# Rename keys in the transformers names
								SCREAMING_SNAKE_CASE       :     Dict				= list(temp.keys()     )
								for key in keys:
										SCREAMING_SNAKE_CASE       :     int				= temp.pop(a_     )

								if tensors is None:
										SCREAMING_SNAKE_CASE       :     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(a_     ) for end in WEIGHTS_TO_AVERAGE_ENDSWITH     ):
														tensors[key] += temp[key]
												else:
														# Some weights are RowParallelLinear in Megatron-Deepspeed, others are ColumnParallel
														SCREAMING_SNAKE_CASE       :     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
														SCREAMING_SNAKE_CASE       :     List[Any]				= torch.cat([tensors[key], temp[key]]		,		dim=a_     )

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

						SCREAMING_SNAKE_CASE       :     Dict				= model.load_state_dict(a_		,		strict=a_     )
						assert not other_keys.unexpected_keys, F"The keys {other_keys.unexpected_keys} are unexpected"
						if missing_keys is None:
								SCREAMING_SNAKE_CASE       :     List[str]				= set(other_keys.missing_keys     )
						else:
								SCREAMING_SNAKE_CASE       :     List[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(a_		,		exist_ok=a_     )
				SCREAMING_SNAKE_CASE       :     int				= pytorch_dump_folder_path + '''/''' + WEIGHTS_NAME
				SCREAMING_SNAKE_CASE       :     List[str]				= 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:
						SCREAMING_SNAKE_CASE       :     Dict				= model.to(config.torch_dtype     )
				torch.save(model.state_dict()		,		a_     )
				print(F"Save configuration file to {pytorch_config_dump_path}"     )
				with open(a_		,		'''w'''		,		encoding='''utf-8'''     ) as f:
						f.write(config.to_json_string()     )


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







"""simple docstring"""






from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_tf_available,
    is_tokenizers_available,
    is_torch_available,
)


lowerCamelCase__ :       Optional[Any]						     =			{
    "configuration_funnel": ["FUNNEL_PRETRAINED_CONFIG_ARCHIVE_MAP", "FunnelConfig"],
    "convert_funnel_original_tf_checkpoint_to_pytorch": [],
    "tokenization_funnel": ["FunnelTokenizer"],
}

try:
					if not is_tokenizers_available():
										raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
					pass
else:
					lowerCamelCase__ :       Optional[Any]						     =			["FunnelTokenizerFast"]

try:
					if not is_torch_available():
										raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
					pass
else:
					lowerCamelCase__ :       Dict						     =			[
					    "FUNNEL_PRETRAINED_MODEL_ARCHIVE_LIST",
					    "FunnelBaseModel",
					    "FunnelForMaskedLM",
					    "FunnelForMultipleChoice",
					    "FunnelForPreTraining",
					    "FunnelForQuestionAnswering",
					    "FunnelForSequenceClassification",
					    "FunnelForTokenClassification",
					    "FunnelModel",
					    "FunnelPreTrainedModel",
					    "load_tf_weights_in_funnel",
					]

try:
					if not is_tf_available():
										raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
					pass
else:
					lowerCamelCase__ :       Tuple						     =			[
					    "TF_FUNNEL_PRETRAINED_MODEL_ARCHIVE_LIST",
					    "TFFunnelBaseModel",
					    "TFFunnelForMaskedLM",
					    "TFFunnelForMultipleChoice",
					    "TFFunnelForPreTraining",
					    "TFFunnelForQuestionAnswering",
					    "TFFunnelForSequenceClassification",
					    "TFFunnelForTokenClassification",
					    "TFFunnelModel",
					    "TFFunnelPreTrainedModel",
					]


if TYPE_CHECKING:
					from .configuration_funnel import FUNNEL_PRETRAINED_CONFIG_ARCHIVE_MAP, FunnelConfig
					from .tokenization_funnel import FunnelTokenizer

					try:
										if not is_tokenizers_available():
															raise OptionalDependencyNotAvailable()
					except OptionalDependencyNotAvailable:
										pass
					else:
										from .tokenization_funnel_fast import FunnelTokenizerFast

					try:
										if not is_torch_available():
															raise OptionalDependencyNotAvailable()
					except OptionalDependencyNotAvailable:
										pass
					else:
										from .modeling_funnel import (
										    FUNNEL_PRETRAINED_MODEL_ARCHIVE_LIST,
										    FunnelBaseModel,
										    FunnelForMaskedLM,
										    FunnelForMultipleChoice,
										    FunnelForPreTraining,
										    FunnelForQuestionAnswering,
										    FunnelForSequenceClassification,
										    FunnelForTokenClassification,
										    FunnelModel,
										    FunnelPreTrainedModel,
										    load_tf_weights_in_funnel,
										)

					try:
										if not is_tf_available():
															raise OptionalDependencyNotAvailable()
					except OptionalDependencyNotAvailable:
										pass
					else:
										from .modeling_tf_funnel import (
										    TF_FUNNEL_PRETRAINED_MODEL_ARCHIVE_LIST,
										    TFFunnelBaseModel,
										    TFFunnelForMaskedLM,
										    TFFunnelForMultipleChoice,
										    TFFunnelForPreTraining,
										    TFFunnelForQuestionAnswering,
										    TFFunnelForSequenceClassification,
										    TFFunnelForTokenClassification,
										    TFFunnelModel,
										    TFFunnelPreTrainedModel,
										)

else:
					import sys

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





