






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

__UpperCAmelCase								=						[
    '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',
]

__UpperCAmelCase								=						[
    'mlp.dense_4h_to_h.weight',
    'self_attention.dense.weight',
]




def      lowercase__      (     __snake_case    :							Union[str, Any]     ,       __snake_case    :							Optional[Any]  ):


					'''simple docstring'''


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




def      lowercase__      (     __snake_case    :							str  ):


					'''simple docstring'''


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




def      lowercase__      (     __snake_case    :							Optional[int]     ,       __snake_case    :							Union[str, Any]     ,       __snake_case    :							Optional[int]     ,       __snake_case    :							int     ,       __snake_case    :							Tuple  ):


					'''simple docstring'''


					if bloom_config_file == "":
										UpperCAmelCase_    :      Optional[int]            =						BloomConfig()
					else:
										UpperCAmelCase_    :      Tuple            =						BloomConfig.from_json_file(__snake_case  )

					if shard_model:
										UpperCAmelCase_    :      Tuple            =						os.listdir(__snake_case  )
										UpperCAmelCase_    :      List[str]            =						sorted(filter(lambda __snake_case  : s.startswith('layer'  ) and "model_00" in s     ,       __snake_case  )  )

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

										UpperCAmelCase_    :      Dict            =						None

										UpperCAmelCase_    :      Optional[int]            =						BloomConfig()

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

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

																				# Rename keys in the transformers names
																				UpperCAmelCase_    :      int            =						list(temp.keys()  )
																				for key in keys:
																									UpperCAmelCase_    :      Tuple            =						temp.pop(__snake_case  )

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

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

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

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

										UpperCAmelCase_    :      int            =						os.listdir(__snake_case  )
										UpperCAmelCase_    :      Any            =						sorted(filter(lambda __snake_case  : s.startswith('layer'  ) and "model_00" in s     ,       __snake_case  )  )

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

																				# Rename keys in the transformers names
																				UpperCAmelCase_    :      List[str]            =						list(temp.keys()  )
																				for key in keys:
																									UpperCAmelCase_    :      List[str]            =						temp.pop(__snake_case  )

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

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

															UpperCAmelCase_    :      Optional[Any]            =						model.load_state_dict(__snake_case     ,       strict=__snake_case  )
															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_    :      Dict            =						missing_keys.intersection(set(other_keys.missing_keys  )  )

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

										# Save pytorch-model
										os.makedirs(__snake_case     ,       exist_ok=__snake_case  )
										UpperCAmelCase_    :      int            =						pytorch_dump_folder_path + '/' + WEIGHTS_NAME
										UpperCAmelCase_    :      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:
															UpperCAmelCase_    :      int            =						model.to(config.torch_dtype  )
										torch.save(model.state_dict()     ,       __snake_case  )
										print(F"Save configuration file to {pytorch_config_dump_path}"  )
										with open(__snake_case     ,       'w'     ,       encoding='utf-8'  ) as f:
															f.write(config.to_json_string()  )


if __name__ == "__main__":
		__UpperCAmelCase								=						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',
		)
		__UpperCAmelCase								=						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 typing import TYPE_CHECKING

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


a_    					=       {
    'configuration_efficientformer': [
        'EFFICIENTFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP',
        'EfficientFormerConfig',
    ]
}

try:
   if not is_vision_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   a_    					=       ['EfficientFormerImageProcessor']

try:
   if not is_torch_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   a_    					=       [
       'EFFICIENTFORMER_PRETRAINED_MODEL_ARCHIVE_LIST',
       'EfficientFormerForImageClassification',
       'EfficientFormerForImageClassificationWithTeacher',
       'EfficientFormerModel',
       'EfficientFormerPreTrainedModel',
   ]

try:
   if not is_tf_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   a_    					=       [
       'TF_EFFICIENTFORMER_PRETRAINED_MODEL_ARCHIVE_LIST',
       'TFEfficientFormerForImageClassification',
       'TFEfficientFormerForImageClassificationWithTeacher',
       'TFEfficientFormerModel',
       'TFEfficientFormerPreTrainedModel',
   ]

if TYPE_CHECKING:
   from .configuration_efficientformer import EFFICIENTFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP, EfficientFormerConfig

   try:
      if not is_vision_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .image_processing_efficientformer import EfficientFormerImageProcessor

   try:
      if not is_torch_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .modeling_efficientformer import (
          EFFICIENTFORMER_PRETRAINED_MODEL_ARCHIVE_LIST,
          EfficientFormerForImageClassification,
          EfficientFormerForImageClassificationWithTeacher,
          EfficientFormerModel,
          EfficientFormerPreTrainedModel,
      )
   try:
      if not is_tf_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .modeling_tf_efficientformer import (
          TF_EFFICIENTFORMER_PRETRAINED_MODEL_ARCHIVE_LIST,
          TFEfficientFormerForImageClassification,
          TFEfficientFormerForImageClassificationWithTeacher,
          TFEfficientFormerModel,
          TFEfficientFormerPreTrainedModel,
      )

else:
   import sys

   a_    					=       _LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)



