


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					:		Union[str, 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""",
]

lowerCamelCase					:		Tuple            =						[
    """mlp.dense_4h_to_h.weight""",
    """self_attention.dense.weight""",
]





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





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





def 	lowercase__(		A     ,			A     ,			A     ,			A     ,			A      ):
       # Construct model
       if bloom_config_file == "":
              snake_case__ :       Union[str, Any]	   =							BloomConfig()
       else:
              snake_case__ :       List[str]	   =							BloomConfig.from_json_file(UpperCAmelCase_      )

       if shard_model:
              snake_case__ :       Dict	   =							os.listdir(UpperCAmelCase_      )
              snake_case__ :       Union[str, Any]	   =							sorted(filter(lambda A      : s.startswith('layer'      ) and "model_00" in s     ,			UpperCAmelCase_      )      )

              snake_case__ :       Any	   =							{'weight_map': {}, 'metadata': {}}
              snake_case__ :       int	   =							0

              snake_case__ :       int	   =							None

              snake_case__ :       Dict	   =							BloomConfig()

              for j, file in enumerate(UpperCAmelCase_      ):
                     print('Processing file: {}'.format(UpperCAmelCase_      )      )
                     snake_case__ :       Optional[Any]	   =							None

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

                            # Rename keys in the transformers names
                            snake_case__ :       Dict	   =							list(temp.keys()      )
                            for key in keys:
                                   snake_case__ :       Optional[Any]	   =							temp.pop(UpperCAmelCase_      )

                            if tensors is None:
                                   snake_case__ :       List[str]	   =							temp
                            else:
                                   for key in tensors.keys():
                                          if any(key.endswith(UpperCAmelCase_      ) 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__ :       Tuple	   =							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__ :       str	   =							torch.cat([tensors[key], temp[key]]     ,			dim=UpperCAmelCase_      )

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

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

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

              snake_case__ :       List[Any]	   =							os.listdir(UpperCAmelCase_      )
              snake_case__ :       Tuple	   =							sorted(filter(lambda A      : s.startswith('layer'      ) and "model_00" in s     ,			UpperCAmelCase_      )      )

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

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

                            if tensors is None:
                                   snake_case__ :       List[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(UpperCAmelCase_      ) for end in WEIGHTS_TO_AVERAGE_ENDSWITH      ):
                                                 tensors[key] += temp[key]
                                          else:
                                                 # Some weights are RowParallelLinear in Megatron-Deepspeed, others are ColumnParallel
                                                 snake_case__ :       Any	   =							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__ :       int	   =							torch.cat([tensors[key], temp[key]]     ,			dim=UpperCAmelCase_      )

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

                     snake_case__ :       int	   =							model.load_state_dict(UpperCAmelCase_     ,			strict=UpperCAmelCase_      )
                     assert not other_keys.unexpected_keys, f'''The keys {other_keys.unexpected_keys} are unexpected'''
                     if missing_keys is None:
                            snake_case__ :       str	   =							set(other_keys.missing_keys      )
                     else:
                            snake_case__ :       int	   =							missing_keys.intersection(set(other_keys.missing_keys      )      )

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

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


if __name__ == "__main__":
 lowerCamelCase					:		Union[str, Any]            =						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					:		List[str]            =						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,
 )









'''simple docstring'''



from collections.abc import Generator
from math import sin



def      UpperCAmelCase__   (      UpperCAmelCase_       :							bytes				)	->						bytes:
     if len(UpperCAmelCase_				) != 32:
          raise ValueError('Input must be of length 32'				)

     __lowerCamelCase  :      Dict														=						B''
     for i in [3, 2, 1, 0]:
          little_endian += string_aa[8 * i : 8 * i + 8]
     return little_endian



def      UpperCAmelCase__   (      UpperCAmelCase_       :							int				)	->						bytes:
     if i < 0:
          raise ValueError('Input must be non-negative'				)

     __lowerCamelCase  :      Union[str, Any]														=						format(UpperCAmelCase_ ,       '08x'				)[-8:]
     __lowerCamelCase  :      str														=						B''
     for i in [3, 2, 1, 0]:
          little_endian_hex += hex_rep[2 * i : 2 * i + 2].encode('utf-8'				)
     return little_endian_hex



def      UpperCAmelCase__   (      UpperCAmelCase_       :							bytes				)	->						bytes:
     __lowerCamelCase  :      Optional[Any]														=						B''
     for char in message:
          bit_string += format(UpperCAmelCase_ ,       '08b'				).encode('utf-8'				)
     __lowerCamelCase  :      List[str]														=						format(len(UpperCAmelCase_				) ,       '064b'				).encode('utf-8'				)

     # Pad bit_string to a multiple of 512 chars
     bit_string += b"1"
     while len(UpperCAmelCase_				) % 5_12 != 4_48:
          bit_string += b"0"
     bit_string += to_little_endian(start_len[32:]				) + to_little_endian(start_len[:32]				)

     return bit_string



def      UpperCAmelCase__   (      UpperCAmelCase_       :							bytes				)	->						Generator[list[int], None, None]:
     if len(UpperCAmelCase_				) % 5_12 != 0:
          raise ValueError('Input must have length that\'s a multiple of 512'				)

     for pos in range(0 ,       len(UpperCAmelCase_				) ,       5_12				):
          __lowerCamelCase  :      Any														=						bit_string[pos : pos + 5_12]
          __lowerCamelCase  :      Optional[int]														=						[]
          for i in range(0 ,       5_12 ,       32				):
               block_words.append(int(to_little_endian(block[i : i + 32]				) ,       2				)				)
          yield block_words



def      UpperCAmelCase__   (      UpperCAmelCase_       :							int				)	->						int:
     if i < 0:
          raise ValueError('Input must be non-negative'				)

     __lowerCamelCase  :      List[Any]														=						format(UpperCAmelCase_ ,       '032b'				)
     __lowerCamelCase  :      Optional[int]														=						''
     for c in i_str:
          new_str += "1" if c == "0" else "0"
     return int(UpperCAmelCase_ ,       2				)



def      UpperCAmelCase__   (      UpperCAmelCase_       :							int ,       UpperCAmelCase_       :							int				)	->						int:
     return (a + b) % 2**32



def      UpperCAmelCase__   (      UpperCAmelCase_       :							int ,       UpperCAmelCase_       :							int				)	->						int:
     if i < 0:
          raise ValueError('Input must be non-negative'				)
     if shift < 0:
          raise ValueError('Shift must be non-negative'				)
     return ((i << shift) ^ (i >> (32 - shift))) % 2**32



def      UpperCAmelCase__   (      UpperCAmelCase_       :							bytes				)	->						bytes:
     __lowerCamelCase  :      Optional[Any]														=						preprocess(UpperCAmelCase_				)

     __lowerCamelCase  :      Union[str, Any]														=						[int(2**32 * abs(sin(i + 1				)				)				) for i in range(64				)]

     # Starting states
     __lowerCamelCase  :      Dict														=						0x67_45_23_01
     __lowerCamelCase  :      Union[str, Any]														=						0xef_cd_ab_89
     __lowerCamelCase  :      Optional[Any]														=						0x98_ba_dc_fe
     __lowerCamelCase  :      Union[str, Any]														=						0x10_32_54_76

     __lowerCamelCase  :      List[str]														=						[
         7,
         12,
         17,
         22,
         7,
         12,
         17,
         22,
         7,
         12,
         17,
         22,
         7,
         12,
         17,
         22,
         5,
         9,
         14,
         20,
         5,
         9,
         14,
         20,
         5,
         9,
         14,
         20,
         5,
         9,
         14,
         20,
         4,
         11,
         16,
         23,
         4,
         11,
         16,
         23,
         4,
         11,
         16,
         23,
         4,
         11,
         16,
         23,
         6,
         10,
         15,
         21,
         6,
         10,
         15,
         21,
         6,
         10,
         15,
         21,
         6,
         10,
         15,
         21,
     ]

     # Process bit string in chunks, each with 16 32-char words
     for block_words in get_block_words(UpperCAmelCase_				):
          __lowerCamelCase  :      Dict														=						aa
          __lowerCamelCase  :      Tuple														=						ba
          __lowerCamelCase  :      List[Any]														=						ca
          __lowerCamelCase  :      Dict														=						da

          # Hash current chunk
          for i in range(64				):
               if i <= 15:
                    # f = (b & c) | (not_32(b) & d)     # Alternate definition for f
                    __lowerCamelCase  :      List[str]														=						d ^ (b & (c ^ d))
                    __lowerCamelCase  :      Optional[int]														=						i
               elif i <= 31:
                    # f = (d & b) | (not_32(d) & c)     # Alternate definition for f
                    __lowerCamelCase  :      Optional[int]														=						c ^ (d & (b ^ c))
                    __lowerCamelCase  :      Tuple														=						(5 * i + 1) % 16
               elif i <= 47:
                    __lowerCamelCase  :      str														=						b ^ c ^ d
                    __lowerCamelCase  :      Any														=						(3 * i + 5) % 16
               else:
                    __lowerCamelCase  :      Union[str, Any]														=						c ^ (b | not_aa(UpperCAmelCase_				))
                    __lowerCamelCase  :      int														=						(7 * i) % 16
               __lowerCamelCase  :      Optional[int]														=						(f + a + added_consts[i] + block_words[g]) % 2**32
               __lowerCamelCase  :      Optional[Any]														=						d
               __lowerCamelCase  :      Tuple														=						c
               __lowerCamelCase  :      Optional[int]														=						b
               __lowerCamelCase  :      List[str]														=						sum_aa(UpperCAmelCase_ ,       left_rotate_aa(UpperCAmelCase_ ,       shift_amounts[i]				)				)

          # Add hashed chunk to running total
          __lowerCamelCase  :      int														=						sum_aa(UpperCAmelCase_ ,       UpperCAmelCase_				)
          __lowerCamelCase  :      str														=						sum_aa(UpperCAmelCase_ ,       UpperCAmelCase_				)
          __lowerCamelCase  :      List[Any]														=						sum_aa(UpperCAmelCase_ ,       UpperCAmelCase_				)
          __lowerCamelCase  :      List[str]														=						sum_aa(UpperCAmelCase_ ,       UpperCAmelCase_				)

     __lowerCamelCase  :      Dict														=						reformat_hex(UpperCAmelCase_				) + reformat_hex(UpperCAmelCase_				) + reformat_hex(UpperCAmelCase_				) + reformat_hex(UpperCAmelCase_				)
     return digest


if __name__ == "__main__":
      import doctest

      doctest.testmod()






