


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

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






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






def       lowerCamelCase					(a_						:Tuple				,				a_						:Optional[int]				,				a_						:Dict				,				a_						:List[Any]				,				a_						:int)     ->							Tuple:
       # Construct model
       if bloom_config_file == "":
              lowercase		:str       							=				BloomConfig()
       else:
              lowercase		:List[str]       							=				BloomConfig.from_json_file(a_)

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

              lowercase		:Optional[Any]       							=				{'''weight_map''': {}, '''metadata''': {}}
              lowercase		:Dict       							=				0

              lowercase		:List[Any]       							=				None

              lowercase		:List[str]       							=				BloomConfig()

              for j, file in enumerate(a_):
                     print('''Processing file: {}'''.format(a_))
                     lowercase		:Union[str, Any]       							=				None

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

                            # Rename keys in the transformers names
                            lowercase		:Any       							=				list(temp.keys())
                            for key in keys:
                                   lowercase		:Dict       							=				temp.pop(a_)

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

              lowercase		:List[Any]       							=				BloomConfig()
              lowercase		:List[Any]       							=				pytorch_dump_folder_path + '''/''' + CONFIG_NAME
              lowercase		:Optional[int]       							=				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:
                     lowercase		:List[str]       							=				json.dumps(a_				,				indent=2				,				sort_keys=a_) + '''\n'''
                     f.write(a_)
       else:
              lowercase		:int       							=				BloomModel(a_)

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

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

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

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

                     lowercase		:Tuple       							=				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:
                            lowercase		:Any       							=				set(other_keys.missing_keys)
                     else:
                            lowercase		:List[str]       							=				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_)
              lowercase		:Dict       							=				pytorch_dump_folder_path + '''/''' + WEIGHTS_NAME
              lowercase		: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:
                     lowercase		: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      		=			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,
 )





"""simple docstring"""



import math






def       lowerCamelCase					(a_						:int)     ->							bool:
       assert isinstance(a_				,				a_) and (
           number >= 0
       ), "'number' must been an int and positive"

       if 1 < number < 4:
              # 2 and 3 are primes
              return True
       elif number < 2 or not number % 2:
              # Negatives, 0, 1 and all even numbers are not primes
              return False

       lowercase		:Any       							=				range(3				,				int(math.sqrt(a_) + 1)				,				2)
       return not any(not number % i for i in odd_numbers)






def       lowerCamelCase					(a_						:Optional[int]				,				a_						:Optional[int]=1				,				**a_						:List[str])     ->							Any:
       lowercase		:str       							=				factor * value
       lowercase		:int       							=				value

       while not is_prime(a_):
              value += 1 if not ("desc" in kwargs and kwargs["desc"] is True) else -1

       if value == first_value_val:
              return next_prime(value + 1				,				**a_)
       return value

