

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

_UpperCamelCase						:  str							       =				[
    "mlp.dense_4h_to_h.weight",
    "self_attention.dense.weight",
]


def       snake_case   (	snake_case      : Dict					,		snake_case      : List[Any]			) ->  List[Any]:


    """simple docstring"""


    lowerCAmelCase 						=							{
        '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
    lowerCAmelCase 						=							int(re.match(r'.*layer_(\d*).*'					,		snake_case			)[1]			)
    layer_number -= 3
    return F'h.{layer_number}.' + key


def       snake_case   (	snake_case      : Optional[Any]			) ->  Union[str, Any]:


    """simple docstring"""


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


def       snake_case   (	snake_case      : int					,		snake_case      : Union[str, Any]					,		snake_case      : Optional[Any]					,		snake_case      : Any					,		snake_case      : Optional[Any]			) ->  Tuple:


    """simple docstring"""


    if bloom_config_file == "":
        lowerCAmelCase 						=							BloomConfig()
    else:
        lowerCAmelCase 						=							BloomConfig.from_json_file(snake_case			)

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

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

        lowerCAmelCase 						=							None

        lowerCAmelCase 						=							BloomConfig()

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

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

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

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

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

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

        lowerCAmelCase 						=							BloomConfig()
        lowerCAmelCase 						=							pytorch_dump_folder_path + '/' + CONFIG_NAME
        lowerCAmelCase 						=							total_size
        with open(snake_case					,		'w'					,		encoding='utf-8'			) as f:
            f.write(config.to_json_string()			)
        with open(os.path.join(snake_case					,		WEIGHTS_NAME + '.index.json'			)					,		'w'					,		encoding='utf-8'			) as f:
            lowerCAmelCase 						=							json.dumps(snake_case					,		indent=2					,		sort_keys=snake_case			) + '\n'
            f.write(snake_case			)
    else:
        lowerCAmelCase 						=							BloomModel(snake_case			)

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

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

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

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

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

            lowerCAmelCase 						=							model.load_state_dict(snake_case					,		strict=snake_case			)
            assert not other_keys.unexpected_keys, F'The keys {other_keys.unexpected_keys} are unexpected'
            if missing_keys is None:
                lowerCAmelCase 						=							set(other_keys.missing_keys			)
            else:
                lowerCAmelCase 						=							missing_keys.intersection(set(other_keys.missing_keys			)			)

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

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


if __name__ == "__main__":
     _UpperCamelCase						:  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",
     )
     _UpperCamelCase						:  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,
     )








'''simple docstring'''



def       snake_case   (	snake_case      : int					,		snake_case      : int					,		snake_case      : int			) ->  int:


    """simple docstring"""


    if exponent == 1:
        return base
    if exponent % 2 == 0:
        lowerCAmelCase 						=							_modexpt(snake_case					,		exponent // 2					,		snake_case			) % modulo_value
        return (x * x) % modulo_value
    else:
        return (base * _modexpt(snake_case					,		exponent - 1					,		snake_case			)) % modulo_value


def       snake_case   (	snake_case      : int = 1777					,		snake_case      : int = 1855					,		snake_case      : int = 8			) ->  int:


    """simple docstring"""


    lowerCAmelCase 						=							base
    for _ in range(1					,		snake_case			):
        lowerCAmelCase 						=							_modexpt(snake_case					,		snake_case					,		10**digits			)

    return result


if __name__ == "__main__":
     print(F"""{solution() = }""")





