
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  :					List[Any]													=       [
    '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  :					int													=       [
    'mlp.dense_4h_to_h.weight',
    'self_attention.dense.weight',
]




def 	_SCREAMING_SNAKE_CASE					(		_lowerCamelCase :			Dict						,  _lowerCamelCase :			Dict)					->    Tuple:






							'''simple docstring'''





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




def 	_SCREAMING_SNAKE_CASE					(		_lowerCamelCase :			Optional[int])					->    Any:






							'''simple docstring'''





							if dtype == torch.bool:
														return 1 / 8
							__UpperCamelCase	:   str			   =					re.search(R"[^\d](\d+)$"						,  str(_lowerCamelCase))
							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 	_SCREAMING_SNAKE_CASE					(		_lowerCamelCase :			Optional[int]						,  _lowerCamelCase :			int						,  _lowerCamelCase :			str						,  _lowerCamelCase :			Optional[Any]						,  _lowerCamelCase :			Any)					->    List[str]:






							'''simple docstring'''





							if bloom_config_file == "":
														__UpperCamelCase	:   Optional[Any]			   =					BloomConfig()
							else:
														__UpperCamelCase	:   int			   =					BloomConfig.from_json_file(_lowerCamelCase)

							if shard_model:
														__UpperCamelCase	:   Dict			   =					os.listdir(_lowerCamelCase)
														__UpperCamelCase	:   int			   =					sorted(filter(lambda _lowerCamelCase: s.startswith("layer") and "model_00" in s						,  _lowerCamelCase))

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

														__UpperCamelCase	:   Dict			   =					None

														__UpperCamelCase	:   Optional[Any]			   =					BloomConfig()

														for j, file in enumerate(_lowerCamelCase):
																					print("Processing file: {}".format(_lowerCamelCase))
																					__UpperCamelCase	:   List[str]			   =					None

																					for i in range(_lowerCamelCase):
																												# load all TP files
																												__UpperCamelCase	:   int			   =					file.replace("model_00"						,  F'model_0{i}')
																												__UpperCamelCase	:   List[str]			   =					torch.load(os.path.join(_lowerCamelCase						,  _lowerCamelCase)						,  map_location="cpu")

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

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

														__UpperCamelCase	:   Any			   =					BloomConfig()
														__UpperCamelCase	:   Any			   =					pytorch_dump_folder_path + "/" + CONFIG_NAME
														__UpperCamelCase	:   Any			   =					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	:   List[str]			   =					json.dumps(_lowerCamelCase						,  indent=2						,  sort_keys=_lowerCamelCase) + "\n"
																					f.write(_lowerCamelCase)
							else:
														__UpperCamelCase	:   Tuple			   =					BloomModel(_lowerCamelCase)

														__UpperCamelCase	:   Optional[Any]			   =					os.listdir(_lowerCamelCase)
														__UpperCamelCase	:   str			   =					sorted(filter(lambda _lowerCamelCase: s.startswith("layer") and "model_00" in s						,  _lowerCamelCase))

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

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

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

																					__UpperCamelCase	:   Union[str, Any]			   =					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	:   Any			   =					set(other_keys.missing_keys)
																					else:
																												__UpperCamelCase	:   List[str]			   =					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	:   Tuple			   =					pytorch_dump_folder_path + "/" + WEIGHTS_NAME
														__UpperCamelCase	:   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:
																					__UpperCamelCase	:   Optional[int]			   =					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__":
		lowercase  :					str													=       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  :					Dict													=       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 __future__ import annotations

import inspect
import unittest

import numpy as np

from transformers import ResNetConfig
from transformers.testing_utils import require_tf, require_vision, slow
from transformers.utils import cached_property, is_tf_available, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
		import tensorflow as tf

		from transformers import TFResNetForImageClassification, TFResNetModel
		from transformers.models.resnet.modeling_tf_resnet import TF_RESNET_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
		from PIL import Image

		from transformers import AutoImageProcessor




class  lowerCamelCase__							:




	'''simple docstring'''




	def __init__(					self		:List[Any]     ,   a		:Dict     ,   a		:Any=3     ,   a		:Any=3_2     ,   a		:Optional[Any]=3     ,   a		:str=1_0     ,   a		:Union[str, Any]=[1_0, 2_0, 3_0, 4_0]     ,   a		:Optional[Any]=[1, 1, 2, 1]     ,   a		:Optional[Any]=True     ,   a		:Dict=True     ,   a		:Tuple="relu"     ,   a		:List[str]=3     ,   a		:Tuple=None     ,   )		-> Tuple:
								__UpperCamelCase	:   Optional[Any]			   =					parent
								__UpperCamelCase	:   Dict			   =					batch_size
								__UpperCamelCase	:   int			   =					image_size
								__UpperCamelCase	:   Dict			   =					num_channels
								__UpperCamelCase	:   Optional[int]			   =					embeddings_size
								__UpperCamelCase	:   List[Any]			   =					hidden_sizes
								__UpperCamelCase	:   Optional[Any]			   =					depths
								__UpperCamelCase	:   Optional[int]			   =					is_training
								__UpperCamelCase	:   Union[str, Any]			   =					use_labels
								__UpperCamelCase	:   Optional[int]			   =					hidden_act
								__UpperCamelCase	:   Tuple			   =					num_labels
								__UpperCamelCase	:   Tuple			   =					scope
								__UpperCamelCase	:   Dict			   =					len(a						)







	def    _lowerCamelCase				(					self		:Optional[int]						)		-> Any:
								__UpperCamelCase	:   List[str]			   =					floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]						)

								__UpperCamelCase	:   List[str]			   =					None
								if self.use_labels:
															__UpperCamelCase	:   Tuple			   =					ids_tensor([self.batch_size]     ,   self.num_labels						)

								__UpperCamelCase	:   List[Any]			   =					self.get_config()

								return config, pixel_values, labels







	def    _lowerCamelCase				(					self		:Union[str, Any]						)		-> int:
								return ResNetConfig(
								    num_channels=self.num_channels     ,   embeddings_size=self.embeddings_size     ,   hidden_sizes=self.hidden_sizes     ,   depths=self.depths     ,   hidden_act=self.hidden_act     ,   num_labels=self.num_labels     ,   image_size=self.image_size     ,   )







	def    _lowerCamelCase				(					self		:List[Any]     ,   a		:Dict     ,   a		:int     ,   a		:Optional[Any]						)		-> Tuple:
								__UpperCamelCase	:   str			   =					TFResNetModel(config=a						)
								__UpperCamelCase	:   Union[str, Any]			   =					model(a						)
								# expected last hidden states: B, C, H // 32, W // 32
								self.parent.assertEqual(
								    result.last_hidden_state.shape     ,   (self.batch_size, self.hidden_sizes[-1], self.image_size // 3_2, self.image_size // 3_2)     ,   )







	def    _lowerCamelCase				(					self		:Union[str, Any]     ,   a		:Optional[int]     ,   a		:List[str]     ,   a		:Optional[Any]						)		-> Any:
								__UpperCamelCase	:   str			   =					self.num_labels
								__UpperCamelCase	:   Optional[int]			   =					TFResNetForImageClassification(a						)
								__UpperCamelCase	:   List[str]			   =					model(a     ,   labels=a						)
								self.parent.assertEqual(result.logits.shape     ,   (self.batch_size, self.num_labels)						)







	def    _lowerCamelCase				(					self		:Optional[int]						)		-> List[str]:
								__UpperCamelCase	:   Optional[int]			   =					self.prepare_config_and_inputs()
								__UpperCamelCase   ,						__UpperCamelCase   ,						__UpperCamelCase	:   Union[str, Any]			   =					config_and_inputs
								__UpperCamelCase	:   List[str]			   =					{"pixel_values": pixel_values}
								return config, inputs_dict




@require_tf
class  lowerCamelCase__							(				__lowercase     ,						__lowercase     ,						unittest.TestCase):




	'''simple docstring'''




	_A								=     (TFResNetModel, TFResNetForImageClassification) if is_tf_available() else ()
	_A								=     (
	    {'feature-extraction': TFResNetModel, 'image-classification': TFResNetForImageClassification}
	    if is_tf_available()
	    else {}
	)

	_A								=     False
	_A								=     False
	_A								=     False
	_A								=     False
	_A								=     False







	def    _lowerCamelCase				(					self		:int						)		-> List[str]:
								__UpperCamelCase	:   Union[str, Any]			   =					TFResNetModelTester(self						)
								__UpperCamelCase	:   List[Any]			   =					ConfigTester(self     ,   config_class=a     ,   has_text_modality=a						)







	def    _lowerCamelCase				(					self		:int						)		-> Dict:
								self.create_and_test_config_common_properties()
								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    _lowerCamelCase				(					self		:str						)		-> Optional[Any]:
								return







	@unittest.skip(reason="ResNet does not use inputs_embeds"						)
	def    _lowerCamelCase				(					self		:Tuple						)		-> Tuple:
								pass







	@unittest.skip(reason="ResNet does not support input and output embeddings"						)
	def    _lowerCamelCase				(					self		:List[Any]						)		-> List[str]:
								pass







	def    _lowerCamelCase				(					self		:Optional[int]						)		-> Tuple:
								__UpperCamelCase   ,						__UpperCamelCase	:   Union[str, Any]			   =					self.model_tester.prepare_config_and_inputs_for_common()

								for model_class in self.all_model_classes:
															__UpperCamelCase	:   Dict			   =					model_class(a						)
															__UpperCamelCase	:   Union[str, Any]			   =					inspect.signature(model.call						)
															# signature.parameters is an OrderedDict => so arg_names order is deterministic
															__UpperCamelCase	:   Dict			   =					[*signature.parameters.keys()]

															__UpperCamelCase	:   Union[str, Any]			   =					["pixel_values"]
															self.assertListEqual(arg_names[:1]     ,   a						)







	def    _lowerCamelCase				(					self		:List[str]						)		-> List[str]:
								__UpperCamelCase	:   Union[str, Any]			   =					self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_model(*a						)







	def    _lowerCamelCase				(					self		:Optional[Any]						)		-> Tuple:
								def check_hidden_states_output(a		:Optional[Any]     ,   a		:Optional[int]     ,   a		:List[str]						):
															__UpperCamelCase	:   int			   =					model_class(a						)
															__UpperCamelCase	:   int			   =					model(**self._prepare_for_class(a     ,   a						)						)

															__UpperCamelCase	:   Optional[Any]			   =					outputs.encoder_hidden_states if config.is_encoder_decoder else outputs.hidden_states

															__UpperCamelCase	:   Optional[int]			   =					self.model_tester.num_stages
															self.assertEqual(len(a						)     ,   expected_num_stages + 1						)

															# ResNet's feature maps are of shape (batch_size, num_channels, height, width)
															self.assertListEqual(
															    list(hidden_states[0].shape[-2:]						)     ,   [self.model_tester.image_size // 4, self.model_tester.image_size // 4]     ,   )

								__UpperCamelCase   ,						__UpperCamelCase	:   Optional[Any]			   =					self.model_tester.prepare_config_and_inputs_for_common()
								__UpperCamelCase	:   str			   =					["basic", "bottleneck"]
								for model_class in self.all_model_classes:
															for layer_type in layers_type:
																						__UpperCamelCase	:   int			   =					layer_type
																						__UpperCamelCase	:   int			   =					True
																						check_hidden_states_output(a     ,   a     ,   a						)

																						# check that output_hidden_states also work using config
																						del inputs_dict["output_hidden_states"]
																						__UpperCamelCase	:   int			   =					True

																						check_hidden_states_output(a     ,   a     ,   a						)







	def    _lowerCamelCase				(					self		:Union[str, Any]						)		-> Dict:
								__UpperCamelCase	:   int			   =					self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_image_classification(*a						)







	@slow
	def    _lowerCamelCase				(					self		:Dict						)		-> Dict:
								for model_name in TF_RESNET_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
															__UpperCamelCase	:   Optional[Any]			   =					TFResNetModel.from_pretrained(a						)
															self.assertIsNotNone(a						)




def 	_SCREAMING_SNAKE_CASE					(		)					->    int:






							'''simple docstring'''





							__UpperCamelCase	:   Tuple			   =					Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png")
							return image




@require_tf
@require_vision
class  lowerCamelCase__							(				unittest.TestCase):




	'''simple docstring'''




	@cached_property
	def    _lowerCamelCase				(					self		:Optional[Any]						)		-> Tuple:
								return (
								    AutoImageProcessor.from_pretrained(TF_RESNET_PRETRAINED_MODEL_ARCHIVE_LIST[0]						)
								    if is_vision_available()
								    else None
								)







	@slow
	def    _lowerCamelCase				(					self		:Optional[int]						)		-> Optional[int]:
								__UpperCamelCase	:   int			   =					TFResNetForImageClassification.from_pretrained(TF_RESNET_PRETRAINED_MODEL_ARCHIVE_LIST[0]						)

								__UpperCamelCase	:   List[Any]			   =					self.default_image_processor
								__UpperCamelCase	:   List[str]			   =					prepare_img()
								__UpperCamelCase	:   List[str]			   =					image_processor(images=a     ,   return_tensors="tf"						)

								# forward pass
								__UpperCamelCase	:   Dict			   =					model(**a						)

								# verify the logits
								__UpperCamelCase	:   Dict			   =					tf.TensorShape((1, 1_0_0_0)						)
								self.assertEqual(outputs.logits.shape     ,   a						)

								__UpperCamelCase	:   Union[str, Any]			   =					tf.constant([-11.1069, -9.7877, -8.3777]						)

								self.assertTrue(np.allclose(outputs.logits[0, :3].numpy()     ,   a     ,   atol=1E-4						)						)