


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

__snake_case    						=[
    """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""",
]

__snake_case    						=[
    """mlp.dense_4h_to_h.weight""",
    """self_attention.dense.weight""",
]





def 	a_					( lowerCamelCase	:     List[Any]	, lowerCamelCase	:     Tuple				):
   lowerCAmelCase     =			{
       '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
   lowerCAmelCase     =			int(re.match(R'.*layer_(\d*).*'	, lowerCamelCase				)[1]				)
   layer_number -= 3
   return f'''h.{layer_number}.''' + key





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





def 	a_					( lowerCamelCase	:     int	, lowerCamelCase	:     str	, lowerCamelCase	:     Any	, lowerCamelCase	:     List[Any]	, lowerCamelCase	:     Union[str, Any]				):
   # Construct model
   if bloom_config_file == "":
      lowerCAmelCase     =			BloomConfig()
   else:
      lowerCAmelCase     =			BloomConfig.from_json_file(lowerCamelCase				)

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

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

      lowerCAmelCase     =			None

      lowerCAmelCase     =			BloomConfig()

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

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

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

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

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

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

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

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

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

         lowerCAmelCase     =			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:
            lowerCAmelCase     =			set(other_keys.missing_keys				)
         else:
            lowerCAmelCase     =			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				)
      lowerCAmelCase     =			pytorch_dump_folder_path + '/' + WEIGHTS_NAME
      lowerCAmelCase     =			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:
         lowerCAmelCase     =			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__":
       __snake_case    						=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""",
       )
       __snake_case    						=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,
       )









class 			_lowercase :  # Public class to implement a graph


       """simple docstring"""




       def __init__(   self							:			Dict					,      __lowerCamelCase							:			int					,      __lowerCamelCase							:			int					,      __lowerCamelCase							:			list[list[bool]]       ):


             '''simple docstring'''



             lowerCamelCase__		:				int    =  row
             lowerCamelCase__		:				Optional[Any]    =  col
             lowerCamelCase__		:				Union[str, Any]    =  graph





       def 						lowerCAmelCase     (   self							:			Optional[int]					,      __lowerCamelCase							:			int					,      __lowerCamelCase							:			int					,      __lowerCamelCase							:			list[list[bool]]       ):


             '''simple docstring'''



             return (
                 0 <= i < self.ROW
                 and 0 <= j < self.COL
                 and not visited[i][j]
                 and self.graph[i][j]
             )





       def 						lowerCAmelCase     (   self							:			Dict					,      __lowerCamelCase							:			int					,      __lowerCamelCase							:			int					,      __lowerCamelCase							:			list[list[bool]]       ):


             '''simple docstring'''



             lowerCamelCase__		:				str    =  [-1, -1, -1, 0, 0, 1, 1, 1]  # Coordinate order
             lowerCamelCase__		:				Tuple    =  [-1, 0, 1, -1, 1, -1, 0, 1]
             lowerCamelCase__		:				List[Any]    =  True  # Make those cells visited
             for k in range(8       ):
                   if self.is_safe(i + row_nbr[k]					,      j + col_nbr[k]					,      __lowerCamelCase       ):
                         self.diffs(i + row_nbr[k]					,      j + col_nbr[k]					,      __lowerCamelCase       )





       def 						lowerCAmelCase     (   self							:			List[Any]       ):  # And finally, count all islands.


             '''simple docstring'''



             lowerCamelCase__		:				List[Any]    =  [[False for j in range(self.COL       )] for i in range(self.ROW       )]
             lowerCamelCase__		:				Optional[Any]    =  0
             for i in range(self.ROW       ):
                   for j in range(self.COL       ):
                         if visited[i][j] is False and self.graph[i][j] == 1:
                               self.diffs(__lowerCamelCase					,      __lowerCamelCase					,      __lowerCamelCase       )
                               count += 1
             return count


