


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

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

def _SCREAMING_SNAKE_CASE (A    ,					A			)			->			Tuple:






       """simple docstring"""







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

def _SCREAMING_SNAKE_CASE (A			)			->			Optional[Any]:






       """simple docstring"""







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

def _SCREAMING_SNAKE_CASE (A    ,					A    ,					A    ,					A    ,					A			)			->			List[str]:






       """simple docstring"""







       if bloom_config_file == "":
              lowercase__      =					BloomConfig()
       else:
              lowercase__      =					BloomConfig.from_json_file(A			)

       if shard_model:
              lowercase__      =					os.listdir(A			)
              lowercase__      =					sorted(filter(lambda A			: s.startswith('''layer'''			) and "model_00" in s    ,					A			)			)

              lowercase__      =					{'''weight_map''': {}, '''metadata''': {}}
              lowercase__      =					0

              lowercase__      =					None

              lowercase__      =					BloomConfig()

              for j, file in enumerate(A			):
                     print('''Processing file: {}'''.format(A			)			)
                     lowercase__      =					None

                     for i in range(A			):
                            # load all TP files
                            lowercase__      =					file.replace('''model_00'''    ,					f"model_0{i}"			)
                            lowercase__      =					torch.load(os.path.join(A    ,					A			)    ,					map_location='''cpu'''			)

                            # Rename keys in the transformers names
                            lowercase__      =					list(temp.keys()			)
                            for key in keys:
                                   lowercase__      =					temp.pop(A			)

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

              lowercase__      =					BloomConfig()
              lowercase__      =					pytorch_dump_folder_path + '''/''' + CONFIG_NAME
              lowercase__      =					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__      =					json.dumps(A    ,					indent=2    ,					sort_keys=A			) + '''\n'''
                     f.write(A			)
       else:
              lowercase__      =					BloomModel(A			)

              lowercase__      =					os.listdir(A			)
              lowercase__      =					sorted(filter(lambda A			: s.startswith('''layer'''			) and "model_00" in s    ,					A			)			)

              lowercase__      =					None
              for i, file in enumerate(A			):
                     lowercase__      =					None
                     for i in range(A			):
                            # load all TP files
                            lowercase__      =					file.replace('''model_00'''    ,					f"model_0{i}"			)
                            lowercase__      =					torch.load(os.path.join(A    ,					A			)    ,					map_location='''cpu'''			)

                            # Rename keys in the transformers names
                            lowercase__      =					list(temp.keys()			)
                            for key in keys:
                                   lowercase__      =					temp.pop(A			)

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

                     lowercase__      =					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__      =					set(other_keys.missing_keys			)
                     else:
                            lowercase__      =					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__      =					pytorch_dump_folder_path + '''/''' + WEIGHTS_NAME
              lowercase__      =					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__      =					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__":
 lowerCamelCase				:  str    		=      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				:  Tuple    		=      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 argparse import ArgumentParser

from .add_new_model import AddNewModelCommand
from .add_new_model_like import AddNewModelLikeCommand
from .convert import ConvertCommand
from .download import DownloadCommand
from .env import EnvironmentCommand
from .lfs import LfsCommands
from .pt_to_tf import PTtoTFCommand
from .run import RunCommand
from .serving import ServeCommand
from .user import UserCommands

def _SCREAMING_SNAKE_CASE ()			->			Dict:






       """simple docstring"""







       lowercase__      =					ArgumentParser('''Transformers CLI tool'''    ,					usage='''transformers-cli <command> [<args>]'''			)
       lowercase__      =					parser.add_subparsers(help='''transformers-cli command helpers'''			)

       # Register commands
       ConvertCommand.register_subcommand(A			)
       DownloadCommand.register_subcommand(A			)
       EnvironmentCommand.register_subcommand(A			)
       RunCommand.register_subcommand(A			)
       ServeCommand.register_subcommand(A			)
       UserCommands.register_subcommand(A			)
       AddNewModelCommand.register_subcommand(A			)
       AddNewModelLikeCommand.register_subcommand(A			)
       LfsCommands.register_subcommand(A			)
       PTtoTFCommand.register_subcommand(A			)

       # Let's go
       lowercase__      =					parser.parse_args()

       if not hasattr(A    ,					'''func'''			):
              parser.print_help()
              exit(1			)

       # Run
       lowercase__      =					args.func(A			)
       service.run()


if __name__ == "__main__":
 main()




