


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_				:			int   =	[
    '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_				:			Dict   =	[
    'mlp.dense_4h_to_h.weight',
    'self_attention.dense.weight',
]







def 	SCREAMING_SNAKE_CASE_						(      __A			:      Optional[Any]  ,   __A			:      Optional[int]    )     ->		str:

       """simple docstring"""

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







def 	SCREAMING_SNAKE_CASE_						(      __A			:      List[str]    )     ->		Any:

       """simple docstring"""

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







def 	SCREAMING_SNAKE_CASE_						(      __A			:      str  ,   __A			:      Union[str, Any]  ,   __A			:      Union[str, Any]  ,   __A			:      Optional[Any]  ,   __A			:      int    )     ->		Any:

       """simple docstring"""

       if bloom_config_file == "":
              a_ :		Dict							=  BloomConfig()
       else:
              a_ :		str							=  BloomConfig.from_json_file(__A    )

       if shard_model:
              a_ :		Optional[int]							=  os.listdir(__A    )
              a_ :		int							=  sorted(filter(lambda __A    : s.startswith('layer'    ) and "model_00" in s  ,   __A    )    )

              a_ :		Union[str, Any]							=  {'weight_map': {}, 'metadata': {}}
              a_ :		Optional[int]							=  0

              a_ :		str							=  None

              a_ :		List[Any]							=  BloomConfig()

              for j, file in enumerate(__A    ):
                     print('Processing file: {}'.format(__A    )    )
                     a_ :		Any							=  None

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

                            # Rename keys in the transformers names
                            a_ :		List[Any]							=  list(temp.keys()    )
                            for key in keys:
                                   a_ :		Union[str, Any]							=  temp.pop(__A    )

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

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

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

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

              a_ :		Tuple							=  os.listdir(__A    )
              a_ :		str							=  sorted(filter(lambda __A    : s.startswith('layer'    ) and "model_00" in s  ,   __A    )    )

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

                            # Rename keys in the transformers names
                            a_ :		int							=  list(temp.keys()    )
                            for key in keys:
                                   a_ :		Optional[int]							=  temp.pop(__A    )

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

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

                     a_ :		Optional[int]							=  model.load_state_dict(__A  ,   strict=__A    )
                     assert not other_keys.unexpected_keys, F"""The keys {other_keys.unexpected_keys} are unexpected"""
                     if missing_keys is None:
                            a_ :		Union[str, Any]							=  set(other_keys.missing_keys    )
                     else:
                            a_ :		Union[str, Any]							=  missing_keys.intersection(set(other_keys.missing_keys    )    )

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

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


if __name__ == "__main__":
   UpperCAmelCase_				:			Optional[int]   =	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_				:			Dict   =	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,
   )






def 	SCREAMING_SNAKE_CASE_						(      __A			:      int    )     ->		int:

       """simple docstring"""

       if not isinstance(__A  ,   __A    ):
              raise ValueError('Input must be an integer'    )
       if input_num <= 0:
              raise ValueError('Input must be positive'    )
       return sum(
           divisor for divisor in range(1  ,   input_num // 2 + 1    ) if input_num % divisor == 0    )


if __name__ == "__main__":
   import doctest

   doctest.testmod()


