


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

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

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


def 				lowerCamelCase__				(	A       :      Dict				,       A       :      str					):




						'''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*).*'''				,       A					)[1]					)
						layer_number -= 3
						return f"""h.{layer_number}.""" + key


def 				lowerCamelCase__				(	A       :      str					):




						'''simple docstring'''






						if dtype == torch.bool:
												return 1 / 8
						UpperCAmelCase					    =     re.search(R'''[^\d](\d+)$'''				,       str(A					)					)
						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 				lowerCamelCase__				(	A       :      Any				,       A       :      Tuple				,       A       :      Dict				,       A       :      Optional[Any]				,       A       :      Union[str, Any]					):




						'''simple docstring'''






						if bloom_config_file == "":
												UpperCAmelCase					    =     BloomConfig()
						else:
												UpperCAmelCase					    =     BloomConfig.from_json_file(A					)

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

												UpperCAmelCase					    =     {'''weight_map''': {}, '''metadata''': {}}
												UpperCAmelCase					    =     0

												UpperCAmelCase					    =     None

												UpperCAmelCase					    =     BloomConfig()

												for j, file in enumerate(A					):
																		print('''Processing file: {}'''.format(A					)					)
																		UpperCAmelCase					    =     None

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

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

																								if tensors is None:
																														UpperCAmelCase					    =     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
																																										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=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					):
																														UpperCAmelCase					    =     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():
																								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(A					)					).zfill(5					)					)

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

												UpperCAmelCase					    =     os.listdir(A					)
												UpperCAmelCase					    =     sorted(filter(lambda A					: s.startswith('''layer'''					) and "model_00" in s				,       A					)					)

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

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

																								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(A					) 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=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					):
																														UpperCAmelCase					    =     tensors[key] / pretraining_tp

																		UpperCAmelCase					    =     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:
																								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(A				,       exist_ok=A					)
												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()				,       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__":
							_lowercase      :      Optional[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""",
							)
							_lowercase      :      List[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'''



import copy
import inspect
import unittest

from transformers import AutoBackbone
from transformers.configuration_utils import PretrainedConfig
from transformers.testing_utils import require_timm, require_torch, torch_device
from transformers.utils.import_utils import is_torch_available

from ...test_backbone_common import BackboneTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor


if is_torch_available():
	import torch

	from transformers import TimmBackbone, TimmBackboneConfig

from ...test_pipeline_mixin import PipelineTesterMixin






class 	UpperCamelCase__:




				def __init__(		self      :			List[str]				,   lowerCAmelCase      :			int				,   lowerCAmelCase      :			Optional[int]=None				,   lowerCAmelCase      :			Optional[int]=None				,   lowerCAmelCase      :			List[str]=None				,   lowerCAmelCase      :			int="resnet50"				,   lowerCAmelCase      :			Optional[int]=3				,   lowerCAmelCase      :			Union[str, Any]=32				,   lowerCAmelCase      :			List[Any]=3				,   lowerCAmelCase      :			Tuple=True				,   lowerCAmelCase      :			str=True				,   )->					Union[str, Any]:

										"""simple docstring"""




										UpperCAmelCase					    =     parent
										UpperCAmelCase					    =     out_indices if out_indices is not None else [4]
										UpperCAmelCase					    =     stage_names
										UpperCAmelCase					    =     out_features
										UpperCAmelCase					    =     backbone
										UpperCAmelCase					    =     batch_size
										UpperCAmelCase					    =     image_size
										UpperCAmelCase					    =     num_channels
										UpperCAmelCase					    =     use_pretrained_backbone
										UpperCAmelCase					    =     is_training




				def 					a__(		self      :			Optional[Any]				)->					Optional[int]:

										"""simple docstring"""




										UpperCAmelCase					    =     floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]				)
										UpperCAmelCase					    =     self.get_config()

										return config, pixel_values




				def 					a__(		self      :			int				)->					Optional[Any]:

										"""simple docstring"""




										return TimmBackboneConfig(
										    image_size=self.image_size				,   num_channels=self.num_channels				,   out_features=self.out_features				,   out_indices=self.out_indices				,   stage_names=self.stage_names				,   use_pretrained_backbone=self.use_pretrained_backbone				,   backbone=self.backbone				,   )




				def 					a__(		self      :			Dict				,   lowerCAmelCase      :			str				,   lowerCAmelCase      :			Dict				)->					Tuple:

										"""simple docstring"""




										UpperCAmelCase					    =     TimmBackbone(config=lowerCAmelCase				)
										model.to(lowerCAmelCase				)
										model.eval()
										with torch.no_grad():
																UpperCAmelCase					    =     model(lowerCAmelCase				)
										self.parent.assertEqual(
										    result.feature_map[-1].shape				,   (self.batch_size, model.channels[-1], 14, 14)				,   )





				def 					a__(		self      :			str				)->					Optional[Any]:

										"""simple docstring"""




										UpperCAmelCase					    =     self.prepare_config_and_inputs()
										UpperCAmelCase    ,    UpperCAmelCase					    =     config_and_inputs
										UpperCAmelCase					    =     {'''pixel_values''': pixel_values}
										return config, inputs_dict







@require_torch
@require_timm
class 	UpperCamelCase__(       lowerCAmelCase						,    lowerCAmelCase						,    lowerCAmelCase						,    unittest.TestCase	):
				__magic_name__    :				List[str]  		=	(TimmBackbone,) if is_torch_available() else ()
				__magic_name__    :				Any  		=	{"feature-extraction": TimmBackbone} if is_torch_available() else {}
				__magic_name__    :				Union[str, Any]  		=	False
				__magic_name__    :				int  		=	False
				__magic_name__    :				Tuple  		=	False
				__magic_name__    :				List[str]  		=	False




				def 					a__(		self      :			int				)->					str:

										"""simple docstring"""




										UpperCAmelCase					    =     TimmBackboneModelTester(self				)
										UpperCAmelCase					    =     ConfigTester(self				,   config_class=lowerCAmelCase				,   has_text_modality=lowerCAmelCase				)




				def 					a__(		self      :			List[str]				)->					Dict:

										"""simple docstring"""




										self.config_tester.create_and_test_config_to_json_string()
										self.config_tester.create_and_test_config_to_json_file()
										self.config_tester.create_and_test_config_from_and_save_pretrained()
										self.config_tester.create_and_test_config_with_num_labels()
										self.config_tester.check_config_can_be_init_without_params()
										self.config_tester.check_config_arguments_init()




				def 					a__(		self      :			List[Any]				)->					Optional[Any]:

										"""simple docstring"""




										UpperCAmelCase					    =     '''resnet18'''
										UpperCAmelCase					    =     '''microsoft/resnet-18'''

										UpperCAmelCase					    =     AutoBackbone.from_pretrained(lowerCAmelCase				,   use_timm_backbone=lowerCAmelCase				)
										UpperCAmelCase					    =     AutoBackbone.from_pretrained(lowerCAmelCase				)

										self.assertEqual(len(timm_model.out_features				)				,   len(transformers_model.out_features				)				)
										self.assertEqual(len(timm_model.stage_names				)				,   len(transformers_model.stage_names				)				)
										self.assertEqual(timm_model.channels				,   transformers_model.channels				)
										# Out indices are set to the last layer by default. For timm models, we don't know
										# the number of layers in advance, so we set it to (-1,), whereas for transformers
										# models, we set it to [len(stage_names) - 1] (kept for backward compatibility).
										self.assertEqual(timm_model.out_indices				,   (-1,)				)
										self.assertEqual(transformers_model.out_indices				,   [len(timm_model.stage_names				) - 1]				)

										UpperCAmelCase					    =     AutoBackbone.from_pretrained(lowerCAmelCase				,   use_timm_backbone=lowerCAmelCase				,   out_indices=[1, 2, 3]				)
										UpperCAmelCase					    =     AutoBackbone.from_pretrained(lowerCAmelCase				,   out_indices=[1, 2, 3]				)

										self.assertEqual(timm_model.out_indices				,   transformers_model.out_indices				)
										self.assertEqual(len(timm_model.out_features				)				,   len(transformers_model.out_features				)				)
										self.assertEqual(timm_model.channels				,   transformers_model.channels				)




				@unittest.skip('''TimmBackbone doesn\'t support feed forward chunking'''				)
				def 					a__(		self      :			Union[str, Any]				)->					Tuple:

										"""simple docstring"""




										pass




				@unittest.skip('''TimmBackbone doesn\'t have num_hidden_layers attribute'''				)
				def 					a__(		self      :			Dict				)->					Dict:

										"""simple docstring"""




										pass




				@unittest.skip('''TimmBackbone initialization is managed on the timm side'''				)
				def 					a__(		self      :			Tuple				)->					List[Any]:

										"""simple docstring"""




										pass




				@unittest.skip('''TimmBackbone models doesn\'t have inputs_embeds'''				)
				def 					a__(		self      :			List[str]				)->					Optional[Any]:

										"""simple docstring"""




										pass




				@unittest.skip('''TimmBackbone models doesn\'t have inputs_embeds'''				)
				def 					a__(		self      :			Optional[Any]				)->					Optional[Any]:

										"""simple docstring"""




										pass




				@unittest.skip('''TimmBackbone model cannot be created without specifying a backbone checkpoint'''				)
				def 					a__(		self      :			Any				)->					Optional[int]:

										"""simple docstring"""




										pass




				@unittest.skip('''Only checkpoints on timm can be loaded into TimmBackbone'''				)
				def 					a__(		self      :			Optional[int]				)->					Optional[int]:

										"""simple docstring"""




										pass




				@unittest.skip('''model weights aren\'t tied in TimmBackbone.'''				)
				def 					a__(		self      :			str				)->					str:

										"""simple docstring"""




										pass




				@unittest.skip('''model weights aren\'t tied in TimmBackbone.'''				)
				def 					a__(		self      :			Any				)->					int:

										"""simple docstring"""




										pass




				@unittest.skip('''Only checkpoints on timm can be loaded into TimmBackbone'''				)
				def 					a__(		self      :			List[Any]				)->					int:

										"""simple docstring"""




										pass




				@unittest.skip('''Only checkpoints on timm can be loaded into TimmBackbone'''				)
				def 					a__(		self      :			List[Any]				)->					List[str]:

										"""simple docstring"""




										pass




				@unittest.skip('''TimmBackbone doesn\'t have hidden size info in its configuration.'''				)
				def 					a__(		self      :			List[str]				)->					Dict:

										"""simple docstring"""




										pass




				@unittest.skip('''TimmBackbone doesn\'t support output_attentions.'''				)
				def 					a__(		self      :			str				)->					Tuple:

										"""simple docstring"""




										pass




				@unittest.skip('''Safetensors is not supported by timm.'''				)
				def 					a__(		self      :			List[Any]				)->					List[Any]:

										"""simple docstring"""




										pass




				@unittest.skip('''Will be fixed soon by reducing the size of the model used for common tests.'''				)
				def 					a__(		self      :			int				)->					Tuple:

										"""simple docstring"""




										pass




				def 					a__(		self      :			Dict				)->					Union[str, Any]:

										"""simple docstring"""




										UpperCAmelCase    ,    UpperCAmelCase					    =     self.model_tester.prepare_config_and_inputs_for_common()

										for model_class in self.all_model_classes:
																UpperCAmelCase					    =     model_class(lowerCAmelCase				)
																UpperCAmelCase					    =     inspect.signature(model.forward				)
																# signature.parameters is an OrderedDict => so arg_names order is deterministic
																UpperCAmelCase					    =     [*signature.parameters.keys()]

																UpperCAmelCase					    =     ['''pixel_values''']
																self.assertListEqual(arg_names[:1]				,   lowerCAmelCase				)




				def 					a__(		self      :			Union[str, Any]				)->					Optional[Any]:

										"""simple docstring"""




										UpperCAmelCase    ,    UpperCAmelCase					    =     self.model_tester.prepare_config_and_inputs_for_common()
										UpperCAmelCase					    =     True
										UpperCAmelCase					    =     self.has_attentions

										# no need to test all models as different heads yield the same functionality
										UpperCAmelCase					    =     self.all_model_classes[0]
										UpperCAmelCase					    =     model_class(lowerCAmelCase				)
										model.to(lowerCAmelCase				)

										UpperCAmelCase					    =     self._prepare_for_class(lowerCAmelCase				,   lowerCAmelCase				)
										UpperCAmelCase					    =     model(**lowerCAmelCase				)
										UpperCAmelCase					    =     outputs[0][-1]

										# Encoder-/Decoder-only models
										UpperCAmelCase					    =     outputs.hidden_states[0]
										hidden_states.retain_grad()

										if self.has_attentions:
																UpperCAmelCase					    =     outputs.attentions[0]
																attentions.retain_grad()

										output.flatten()[0].backward(retain_graph=lowerCAmelCase				)

										self.assertIsNotNone(hidden_states.grad				)

										if self.has_attentions:
																self.assertIsNotNone(attentions.grad				)





				def 					a__(		self      :			Optional[int]				)->					Optional[Any]:

										"""simple docstring"""




										UpperCAmelCase    ,    UpperCAmelCase					    =     self.model_tester.prepare_config_and_inputs_for_common()

										for model_class in self.all_model_classes:
																UpperCAmelCase					    =     model_class(lowerCAmelCase				)
																model.to(lowerCAmelCase				)
																model.eval()
																UpperCAmelCase					    =     model(**lowerCAmelCase				)

																self.assertEqual(len(result.feature_maps				)				,   len(config.out_indices				)				)
																self.assertEqual(len(model.channels				)				,   len(config.out_indices				)				)

																# Check output of last stage is taken if out_features=None, out_indices=None
																UpperCAmelCase					    =     copy.deepcopy(lowerCAmelCase				)
																UpperCAmelCase					    =     None
																UpperCAmelCase					    =     model_class(lowerCAmelCase				)
																model.to(lowerCAmelCase				)
																model.eval()
																UpperCAmelCase					    =     model(**lowerCAmelCase				)

																self.assertEqual(len(result.feature_maps				)				,   1				)
																self.assertEqual(len(model.channels				)				,   1				)

																# Check backbone can be initialized with fresh weights
																UpperCAmelCase					    =     copy.deepcopy(lowerCAmelCase				)
																UpperCAmelCase					    =     False
																UpperCAmelCase					    =     model_class(lowerCAmelCase				)
																model.to(lowerCAmelCase				)
																model.eval()
																UpperCAmelCase					    =     model(**lowerCAmelCase				)





