


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				:    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''',
]

snake_case				:    str       				=     [
    '''mlp.dense_4h_to_h.weight''',
    '''self_attention.dense.weight''',
]
def 				__lowerCamelCase						( UpperCAmelCase_						:		List[Any] ,  UpperCAmelCase_						:		Tuple						):




      """simple docstring"""





      a						:Union[str, 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
      a						:Tuple 							=       int(re.match(R'''.*layer_(\d*).*''' ,  UpperCAmelCase_						)[1]						)
      layer_number -= 3
      return F'''h.{layer_number}.''' + key
def 				__lowerCamelCase						( UpperCAmelCase_						:		List[str]						):




      """simple docstring"""





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




      """simple docstring"""





      if bloom_config_file == "":
            a						:Any 							=       BloomConfig()
      else:
            a						:List[str] 							=       BloomConfig.from_json_file(UpperCAmelCase_						)

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

            a						:Dict 							=       {'''weight_map''': {}, '''metadata''': {}}
            a						:Dict 							=       0

            a						:Tuple 							=       None

            a						:int 							=       BloomConfig()

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

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

                        # Rename keys in the transformers names
                        a						:Tuple 							=       list(temp.keys()						)
                        for key in keys:
                              a						:Any 							=       temp.pop(UpperCAmelCase_						)

                        if tensors is None:
                              a						: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
                                          a						:Any 							=       1 if any(text in key for text in WEIGHTS_WITH_ROW_PARALLELISM_CONTAIN						) else 0
                                          # We concatenate these weights accross TP ranks
                                          a						: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						):
                              a						:Optional[int] 							=       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():
                        a						:str 							=       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(UpperCAmelCase_						)						).zfill(5						)						)

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

            a						:Tuple 							=       os.listdir(UpperCAmelCase_						)
            a						:Dict 							=       sorted(filter(lambda UpperCAmelCase_						: s.startswith('''layer'''						) and "model_00" in s ,  UpperCAmelCase_						)						)

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

                        # Rename keys in the transformers names
                        a						:str 							=       list(temp.keys()						)
                        for key in keys:
                              a						:int 							=       temp.pop(UpperCAmelCase_						)

                        if tensors is None:
                              a						: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
                                          a						:Union[str, Any] 							=       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=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						):
                              a						:Union[str, Any] 							=       tensors[key] / pretraining_tp

                  a						:Optional[Any] 							=       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:
                        a						:List[str] 							=       set(other_keys.missing_keys						)
                  else:
                        a						: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_						)
            a						:int 							=       pytorch_dump_folder_path + '''/''' + WEIGHTS_NAME
            a						:List[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						: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__":
     snake_case				:    Optional[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''',
     )
     snake_case				:    List[Any]       				=     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,
     )






from __future__ import annotations





def    __lowercase     (						lowerCamelCase					:     Optional[Any]  ,  lowerCamelCase					:     Dict  ,  lowerCamelCase					:     Union[str, Any]  ,  lowerCamelCase					:     List[str]	):  # noqa: E741
 while r - l > 1:
  UpperCamelCase_					:						Union[str, Any] =					(l + r) // 2
  if v[m] >= key:
   UpperCamelCase_					:						str =					m
  else:
   UpperCamelCase_					:						List[Any] =					m  # noqa: E741
 return r





def    __lowercase     (						lowerCamelCase					:     list[int]	):
 if len(lowerCamelCase	) == 0:
  return 0

 UpperCamelCase_					:						Tuple =					[0] * len(lowerCamelCase	)
 UpperCamelCase_					:						int =					1

 UpperCamelCase_					:						Dict =					v[0]

 for i in range(1  ,  len(lowerCamelCase	)	):
  if v[i] < tail[0]:
   UpperCamelCase_					:						Any =					v[i]
  elif v[i] > tail[length - 1]:
   UpperCamelCase_					:						Dict =					v[i]
   length += 1
  else:
   UpperCamelCase_					:						List[str] =					v[i]

 return length


if __name__ == "__main__":
   import doctest

   doctest.testmod()



