
'''simple docstring'''
import argparse
import json
import os

import numpy as np
import PIL
import requests
import tensorflow.keras.applications.efficientnet as efficientnet
import torch
from huggingface_hub import hf_hub_download
from PIL import Image
from tensorflow.keras.preprocessing import image

from transformers import (
    EfficientNetConfig,
    EfficientNetForImageClassification,
    EfficientNetImageProcessor,
)
from transformers.utils import logging


logging.set_verbosity_info()
__lowerCAmelCase												=	logging.get_logger(__name__)

__lowerCAmelCase												=	{
    'b0': efficientnet.EfficientNetBa,
    'b1': efficientnet.EfficientNetBa,
    'b2': efficientnet.EfficientNetBa,
    'b3': efficientnet.EfficientNetBa,
    'b4': efficientnet.EfficientNetBa,
    'b5': efficientnet.EfficientNetBa,
    'b6': efficientnet.EfficientNetBa,
    'b7': efficientnet.EfficientNetBa,
}

__lowerCAmelCase												=	{
    'b0': {
        'hidden_dim': 1_280,
        'width_coef': 1.0,
        'depth_coef': 1.0,
        'image_size': 224,
        'dropout_rate': 0.2,
        'dw_padding': [],
    },
    'b1': {
        'hidden_dim': 1_280,
        'width_coef': 1.0,
        'depth_coef': 1.1,
        'image_size': 240,
        'dropout_rate': 0.2,
        'dw_padding': [16],
    },
    'b2': {
        'hidden_dim': 1_408,
        'width_coef': 1.1,
        'depth_coef': 1.2,
        'image_size': 260,
        'dropout_rate': 0.3,
        'dw_padding': [5, 8, 16],
    },
    'b3': {
        'hidden_dim': 1_536,
        'width_coef': 1.2,
        'depth_coef': 1.4,
        'image_size': 300,
        'dropout_rate': 0.3,
        'dw_padding': [5, 18],
    },
    'b4': {
        'hidden_dim': 1_792,
        'width_coef': 1.4,
        'depth_coef': 1.8,
        'image_size': 380,
        'dropout_rate': 0.4,
        'dw_padding': [6],
    },
    'b5': {
        'hidden_dim': 2_048,
        'width_coef': 1.6,
        'depth_coef': 2.2,
        'image_size': 456,
        'dropout_rate': 0.4,
        'dw_padding': [13, 27],
    },
    'b6': {
        'hidden_dim': 2_304,
        'width_coef': 1.8,
        'depth_coef': 2.6,
        'image_size': 528,
        'dropout_rate': 0.5,
        'dw_padding': [31],
    },
    'b7': {
        'hidden_dim': 2_560,
        'width_coef': 2.0,
        'depth_coef': 3.1,
        'image_size': 600,
        'dropout_rate': 0.5,
        'dw_padding': [18],
    },
}






def  __SCREAMING_SNAKE_CASE				(  _SCREAMING_SNAKE_CASE				):
							_snake_case       =  EfficientNetConfig()
							_snake_case       =  CONFIG_MAP[model_name]["""hidden_dim"""]
							_snake_case       =  CONFIG_MAP[model_name]["""width_coef"""]
							_snake_case       =  CONFIG_MAP[model_name]["""depth_coef"""]
							_snake_case       =  CONFIG_MAP[model_name]["""image_size"""]
							_snake_case       =  CONFIG_MAP[model_name]["""dropout_rate"""]
							_snake_case       =  CONFIG_MAP[model_name]["""dw_padding"""]

							_snake_case       =  """huggingface/label-files"""
							_snake_case       =  """imagenet-1k-id2label.json"""
							_snake_case       =  1000
							_snake_case       =  json.load(open(hf_hub_download(_lowerCAmelCase							, _lowerCAmelCase							, repo_type="""dataset"""				)							, """r"""				)				)
							_snake_case       =  {int(_lowerCAmelCase				): v for k, v in idalabel.items()}

							_snake_case       =  idalabel
							_snake_case       =  {v: k for k, v in idalabel.items()}
							return config






def  __SCREAMING_SNAKE_CASE				(  ):
							_snake_case       =  """http://images.cocodataset.org/val2017/000000039769.jpg"""
							_snake_case       =  Image.open(requests.get(_lowerCAmelCase							, stream=_lowerCAmelCase				).raw				)
							return im






def  __SCREAMING_SNAKE_CASE				(  _SCREAMING_SNAKE_CASE				):
							_snake_case       =  CONFIG_MAP[model_name]["""image_size"""]
							_snake_case       =  EfficientNetImageProcessor(
							    size={"""height""": size, """width""": size}							, image_mean=[0.485, 0.456, 0.406]							, image_std=[0.4785_3944, 0.473_2864, 0.4743_4163]							, do_center_crop=_lowerCAmelCase							, )
							return preprocessor






def  __SCREAMING_SNAKE_CASE				(  _SCREAMING_SNAKE_CASE				):
							_snake_case       =  [v.split("""_"""				)[0].split("""block"""				)[1] for v in original_param_names if v.startswith("""block"""				)]
							_snake_case       =  sorted(set(_lowerCAmelCase				)				)
							_snake_case       =  len(_lowerCAmelCase				)
							_snake_case       =  {b: str(_lowerCAmelCase				) for b, i in zip(_lowerCAmelCase							, range(_lowerCAmelCase				)				)}

							_snake_case       =  []
							rename_keys.append(("""stem_conv/kernel:0""", """embeddings.convolution.weight""")				)
							rename_keys.append(("""stem_bn/gamma:0""", """embeddings.batchnorm.weight""")				)
							rename_keys.append(("""stem_bn/beta:0""", """embeddings.batchnorm.bias""")				)
							rename_keys.append(("""stem_bn/moving_mean:0""", """embeddings.batchnorm.running_mean""")				)
							rename_keys.append(("""stem_bn/moving_variance:0""", """embeddings.batchnorm.running_var""")				)

							for b in block_names:
														_snake_case       =  block_name_mapping[b]
														rename_keys.append((f"""block{b}_expand_conv/kernel:0""", f"""encoder.blocks.{hf_b}.expansion.expand_conv.weight""")				)
														rename_keys.append((f"""block{b}_expand_bn/gamma:0""", f"""encoder.blocks.{hf_b}.expansion.expand_bn.weight""")				)
														rename_keys.append((f"""block{b}_expand_bn/beta:0""", f"""encoder.blocks.{hf_b}.expansion.expand_bn.bias""")				)
														rename_keys.append(
														    (f"""block{b}_expand_bn/moving_mean:0""", f"""encoder.blocks.{hf_b}.expansion.expand_bn.running_mean""")				)
														rename_keys.append(
														    (f"""block{b}_expand_bn/moving_variance:0""", f"""encoder.blocks.{hf_b}.expansion.expand_bn.running_var""")				)
														rename_keys.append(
														    (f"""block{b}_dwconv/depthwise_kernel:0""", f"""encoder.blocks.{hf_b}.depthwise_conv.depthwise_conv.weight""")				)
														rename_keys.append((f"""block{b}_bn/gamma:0""", f"""encoder.blocks.{hf_b}.depthwise_conv.depthwise_norm.weight""")				)
														rename_keys.append((f"""block{b}_bn/beta:0""", f"""encoder.blocks.{hf_b}.depthwise_conv.depthwise_norm.bias""")				)
														rename_keys.append(
														    (f"""block{b}_bn/moving_mean:0""", f"""encoder.blocks.{hf_b}.depthwise_conv.depthwise_norm.running_mean""")				)
														rename_keys.append(
														    (f"""block{b}_bn/moving_variance:0""", f"""encoder.blocks.{hf_b}.depthwise_conv.depthwise_norm.running_var""")				)

														rename_keys.append((f"""block{b}_se_reduce/kernel:0""", f"""encoder.blocks.{hf_b}.squeeze_excite.reduce.weight""")				)
														rename_keys.append((f"""block{b}_se_reduce/bias:0""", f"""encoder.blocks.{hf_b}.squeeze_excite.reduce.bias""")				)
														rename_keys.append((f"""block{b}_se_expand/kernel:0""", f"""encoder.blocks.{hf_b}.squeeze_excite.expand.weight""")				)
														rename_keys.append((f"""block{b}_se_expand/bias:0""", f"""encoder.blocks.{hf_b}.squeeze_excite.expand.bias""")				)
														rename_keys.append(
														    (f"""block{b}_project_conv/kernel:0""", f"""encoder.blocks.{hf_b}.projection.project_conv.weight""")				)
														rename_keys.append((f"""block{b}_project_bn/gamma:0""", f"""encoder.blocks.{hf_b}.projection.project_bn.weight""")				)
														rename_keys.append((f"""block{b}_project_bn/beta:0""", f"""encoder.blocks.{hf_b}.projection.project_bn.bias""")				)
														rename_keys.append(
														    (f"""block{b}_project_bn/moving_mean:0""", f"""encoder.blocks.{hf_b}.projection.project_bn.running_mean""")				)
														rename_keys.append(
														    (f"""block{b}_project_bn/moving_variance:0""", f"""encoder.blocks.{hf_b}.projection.project_bn.running_var""")				)

							rename_keys.append(("""top_conv/kernel:0""", """encoder.top_conv.weight""")				)
							rename_keys.append(("""top_bn/gamma:0""", """encoder.top_bn.weight""")				)
							rename_keys.append(("""top_bn/beta:0""", """encoder.top_bn.bias""")				)
							rename_keys.append(("""top_bn/moving_mean:0""", """encoder.top_bn.running_mean""")				)
							rename_keys.append(("""top_bn/moving_variance:0""", """encoder.top_bn.running_var""")				)

							_snake_case       =  {}
							for item in rename_keys:
														if item[0] in original_param_names:
																					_snake_case       =  """efficientnet.""" + item[1]

							_snake_case       =  """classifier.weight"""
							_snake_case       =  """classifier.bias"""
							return key_mapping






def  __SCREAMING_SNAKE_CASE				(  _SCREAMING_SNAKE_CASE							, _SCREAMING_SNAKE_CASE							, _SCREAMING_SNAKE_CASE				):



							for key, value in tf_params.items():
														if "normalization" in key:
																					continue

														_snake_case       =  key_mapping[key]
														if "_conv" in key and "kernel" in key:
																					_snake_case       =  torch.from_numpy(_lowerCAmelCase				).permute(3							, 2							, 0							, 1				)
														elif "depthwise_kernel" in key:
																					_snake_case       =  torch.from_numpy(_lowerCAmelCase				).permute(2							, 3							, 0							, 1				)
														elif "kernel" in key:
																					_snake_case       =  torch.from_numpy(np.transpose(_lowerCAmelCase				)				)
														else:
																					_snake_case       =  torch.from_numpy(_lowerCAmelCase				)

														# Replace HF parameters with original TF model parameters
														assert hf_params[hf_key].shape == new_hf_value.shape
														hf_params[hf_key].copy_(_lowerCAmelCase				)






@torch.no_grad()
def  __SCREAMING_SNAKE_CASE				(  _SCREAMING_SNAKE_CASE							, _SCREAMING_SNAKE_CASE							, _SCREAMING_SNAKE_CASE							, _SCREAMING_SNAKE_CASE				):
							_snake_case       =  model_classes[model_name](
							    include_top=_lowerCAmelCase							, weights="""imagenet"""							, input_tensor=_lowerCAmelCase							, input_shape=_lowerCAmelCase							, pooling=_lowerCAmelCase							, classes=1000							, classifier_activation="""softmax"""							, )

							_snake_case       =  original_model.trainable_variables
							_snake_case       =  original_model.non_trainable_variables
							_snake_case       =  {param.name: param.numpy() for param in tf_params}
							for param in tf_non_train_params:
														_snake_case       =  param.numpy()
							_snake_case       =  list(tf_params.keys()				)

							# Load HuggingFace model
							_snake_case       =  get_efficientnet_config(_lowerCAmelCase				)
							_snake_case       =  EfficientNetForImageClassification(_lowerCAmelCase				).eval()
							_snake_case       =  hf_model.state_dict()

							# Create src-to-dst parameter name mapping dictionary
							print("""Converting parameters..."""				)
							_snake_case       =  rename_keys(_lowerCAmelCase				)
							replace_params(_lowerCAmelCase							, _lowerCAmelCase							, _lowerCAmelCase				)

							# Initialize preprocessor and preprocess input image
							_snake_case       =  convert_image_processor(_lowerCAmelCase				)
							_snake_case       =  preprocessor(images=prepare_img()							, return_tensors="""pt"""				)

							# HF model inference
							hf_model.eval()
							with torch.no_grad():
														_snake_case       =  hf_model(**_lowerCAmelCase				)
							_snake_case       =  outputs.logits.detach().numpy()

							# Original model inference
							_snake_case       =  False
							_snake_case       =  CONFIG_MAP[model_name]["""image_size"""]
							_snake_case       =  prepare_img().resize((image_size, image_size)							, resample=PIL.Image.NEAREST				)
							_snake_case       =  image.img_to_array(_lowerCAmelCase				)
							_snake_case       =  np.expand_dims(_lowerCAmelCase							, axis=0				)
							_snake_case       =  original_model.predict(_lowerCAmelCase				)

							# Check whether original and HF model outputs match  -> np.allclose
							assert np.allclose(_lowerCAmelCase							, _lowerCAmelCase							, atol=1E-3				), "The predicted logits are not the same."
							print("""Model outputs match!"""				)

							if save_model:
														# Create folder to save model
														if not os.path.isdir(_lowerCAmelCase				):
																					os.mkdir(_lowerCAmelCase				)
														# Save converted model and image processor
														hf_model.save_pretrained(_lowerCAmelCase				)
														preprocessor.save_pretrained(_lowerCAmelCase				)

							if push_to_hub:
														# Push model and image processor to hub
														print(f"""Pushing converted {model_name} to the hub..."""				)
														_snake_case       =  f"""efficientnet-{model_name}"""
														preprocessor.push_to_hub(_lowerCAmelCase				)
														hf_model.push_to_hub(_lowerCAmelCase				)


if __name__ == "__main__":
			__lowerCAmelCase												=	argparse.ArgumentParser()
			# Required parameters
			parser.add_argument(
			    '--model_name',
			    default='b0',
			    type=str,
			    help='Version name of the EfficientNet model you want to convert, select from [b0, b1, b2, b3, b4, b5, b6, b7].',
			)
			parser.add_argument(
			    '--pytorch_dump_folder_path',
			    default='hf_model',
			    type=str,
			    help='Path to the output PyTorch model directory.',
			)
			parser.add_argument('--save_model', action='store_true', help='Save model to local')
			parser.add_argument('--push_to_hub', action='store_true', help='Push model and image processor to the hub')

			__lowerCAmelCase												=	parser.parse_args()
			convert_efficientnet_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.save_model, args.push_to_hub)

'''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  __SCREAMING_SNAKE_CASE				(  _SCREAMING_SNAKE_CASE							, _SCREAMING_SNAKE_CASE				):
							_snake_case       =  {
							    """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
							_snake_case       =  int(re.match(R""".*layer_(\d*).*"""							, _SCREAMING_SNAKE_CASE				)[1]				)
							layer_number -= 3
							return f"""h.{layer_number}.""" + key






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






def  __SCREAMING_SNAKE_CASE				(  _SCREAMING_SNAKE_CASE							, _SCREAMING_SNAKE_CASE							, _SCREAMING_SNAKE_CASE							, _SCREAMING_SNAKE_CASE							, _SCREAMING_SNAKE_CASE				):
							# Construct model
							if bloom_config_file == "":
														_snake_case       =  BloomConfig()
							else:
														_snake_case       =  BloomConfig.from_json_file(_SCREAMING_SNAKE_CASE				)

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

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

														_snake_case       =  None

														_snake_case       =  BloomConfig()

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

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

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

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

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

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

														_snake_case       =  BloomConfig()
														_snake_case       =  pytorch_dump_folder_path + """/""" + CONFIG_NAME
														_snake_case       =  total_size
														with open(_SCREAMING_SNAKE_CASE							, """w"""							, encoding="""utf-8"""				) as f:
																					f.write(config.to_json_string()				)
														with open(os.path.join(_SCREAMING_SNAKE_CASE							, WEIGHTS_NAME + """.index.json"""				)							, """w"""							, encoding="""utf-8"""				) as f:
																					_snake_case       =  json.dumps(_SCREAMING_SNAKE_CASE							, indent=2							, sort_keys=_SCREAMING_SNAKE_CASE				) + """\n"""
																					f.write(_SCREAMING_SNAKE_CASE				)
							else:
														_snake_case       =  BloomModel(_SCREAMING_SNAKE_CASE				)

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

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

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

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

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

																					_snake_case       =  model.load_state_dict(_SCREAMING_SNAKE_CASE							, strict=_SCREAMING_SNAKE_CASE				)
																					assert not other_keys.unexpected_keys, f"""The keys {other_keys.unexpected_keys} are unexpected"""
																					if missing_keys is None:
																												_snake_case       =  set(other_keys.missing_keys				)
																					else:
																												_snake_case       =  missing_keys.intersection(set(other_keys.missing_keys				)				)

														assert not missing_keys, f"""The keys {missing_keys} are missing"""

														# Save pytorch-model
														os.makedirs(_SCREAMING_SNAKE_CASE							, exist_ok=_SCREAMING_SNAKE_CASE				)
														_snake_case       =  pytorch_dump_folder_path + """/""" + WEIGHTS_NAME
														_snake_case       =  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:
																					_snake_case       =  model.to(config.torch_dtype				)
														torch.save(model.state_dict()							, _SCREAMING_SNAKE_CASE				)
														print(f"""Save configuration file to {pytorch_config_dump_path}"""				)
														with open(_SCREAMING_SNAKE_CASE							, """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,
				)