





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

def   a_   (			lowerCAmelCase_				:      Union[str, Any],     lowerCAmelCase_				:      Union[str, Any]     ):
   __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_				:      Tuple     ):
   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_				:      Any,     lowerCAmelCase_				:      Union[str, Any],     lowerCAmelCase_				:      Union[str, Any],     lowerCAmelCase_				:      int,     lowerCAmelCase_				:      Tuple     ):
   # 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  :   Tuple						   =  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  :   int						   =  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     _UpperCAmelCase   :

  """simple docstring"""


  def __init__(    self			:		Union[str, Any] ,  lowerCAmelCase_			:		List[str] ,  lowerCAmelCase_			:		int ,  lowerCAmelCase_			:		Dict					)							->						List[str]:
     __lowerCAmelCase							     =     name
     __lowerCAmelCase							     =     value
     __lowerCAmelCase							     =     weight






  def __repr__(    self			:		Union[str, Any]					)							->						List[str]:
     return f"""{self.__class__.__name__}({self.name}, {self.value}, {self.weight})"""






  def     lowercase		(    self			:		int					)							->						Optional[int]:
     return self.value






  def     lowercase		(    self			:		Optional[Any]					)							->						Union[str, Any]:
     return self.name






  def     lowercase		(    self			:		List[Any]					)							->						Tuple:
     return self.weight






  def     lowercase		(    self			:		int					)							->						Dict:
     return self.value / self.weight





def   a_   (			lowerCAmelCase_				:      Dict,     lowerCAmelCase_				:      int,     lowerCAmelCase_				:      Optional[int]     ):
   __lowerCAmelCase							     =     []
   for i in range(len(lowerCAmelCase_     )     ):
      menu.append(Things(name[i],     value[i],     weight[i]     )     )
   return menu

def   a_   (			lowerCAmelCase_				:      Optional[int],     lowerCAmelCase_				:      List[Any],     lowerCAmelCase_				:      Optional[Any]     ):
   __lowerCAmelCase							     =     sorted(lowerCAmelCase_,     key=lowerCAmelCase_,     reverse=lowerCAmelCase_     )
   __lowerCAmelCase							     =     []
   __lowerCAmelCase     ,     __lowerCAmelCase							     =     0.0, 0.0
   for i in range(len(lowerCAmelCase_     )     ):
      if (total_cost + items_copy[i].get_weight()) <= max_cost:
         result.append(items_copy[i]     )
         total_cost += items_copy[i].get_weight()
         total_value += items_copy[i].get_value()
   return (result, total_value)

def   a_   (			):
   pass


if __name__ == "__main__":
     import doctest

     doctest.testmod()

