





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

_SCREAMING_SNAKE_CASE									=				[
    "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",
]

_SCREAMING_SNAKE_CASE									=				[
    "mlp.dense_4h_to_h.weight",
    "self_attention.dense.weight",
]


def 						__a(SCREAMING_SNAKE_CASE_	:       List[Any]		,							SCREAMING_SNAKE_CASE_	:       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*).*"		,							SCREAMING_SNAKE_CASE_     )[1]     )
    layer_number -= 3
    return F'''h.{layer_number}.''' + key


def 						__a(SCREAMING_SNAKE_CASE_	:       Dict     ):



    '''simple docstring'''





    if dtype == torch.bool:
        return 1 / 8
    _lowerCAmelCase						      =   re.search(R"[^\d](\d+)$"		,							str(SCREAMING_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 						__a(SCREAMING_SNAKE_CASE_	:       Optional[Any]		,							SCREAMING_SNAKE_CASE_	:       Dict		,							SCREAMING_SNAKE_CASE_	:       Union[str, Any]		,							SCREAMING_SNAKE_CASE_	:       Optional[int]		,							SCREAMING_SNAKE_CASE_	:       List[Any]     ):



    '''simple docstring'''





    if bloom_config_file == "":
        _lowerCAmelCase						      =   BloomConfig()
    else:
        _lowerCAmelCase						      =   BloomConfig.from_json_file(SCREAMING_SNAKE_CASE_     )

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

        _lowerCAmelCase						      =   {"weight_map": {}, "metadata": {}}
        _lowerCAmelCase						      =   0

        _lowerCAmelCase						      =   None

        _lowerCAmelCase						      =   BloomConfig()

        for j, file in enumerate(SCREAMING_SNAKE_CASE_     ):
            print("Processing file: {}".format(SCREAMING_SNAKE_CASE_     )     )
            _lowerCAmelCase						      =   None

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

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

                if tensors is None:
                    _lowerCAmelCase						      =   temp
                else:
                    for key in tensors.keys():
                        if any(key.endswith(SCREAMING_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=SCREAMING_SNAKE_CASE_     )

            # Divide by the number of TP the weights we want to average
            for key in tensors.keys():
                if any(key.endswith(SCREAMING_SNAKE_CASE_     ) for end in WEIGHTS_TO_AVERAGE_ENDSWITH     ):
                    _lowerCAmelCase						      =   tensors[key] / pretraining_tp
            torch.save(
                SCREAMING_SNAKE_CASE_		,							os.path.join(
                    SCREAMING_SNAKE_CASE_		,							"pytorch_model_{}-of-{}.bin".format(str(j + 1     ).zfill(5     )		,							str(len(SCREAMING_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(SCREAMING_SNAKE_CASE_     )     ).zfill(5     )     )

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

        _lowerCAmelCase						      =   os.listdir(SCREAMING_SNAKE_CASE_     )
        _lowerCAmelCase						      =   sorted(filter(lambda SCREAMING_SNAKE_CASE_     : s.startswith("layer"     ) and "model_00" in s		,							SCREAMING_SNAKE_CASE_     )     )

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

                # Rename keys in the transformers names
                _lowerCAmelCase						      =   list(temp.keys()     )
                for key in keys:
                    _lowerCAmelCase						      =   temp.pop(SCREAMING_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(SCREAMING_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=SCREAMING_SNAKE_CASE_     )

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

            _lowerCAmelCase						      =   model.load_state_dict(SCREAMING_SNAKE_CASE_		,							strict=SCREAMING_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(SCREAMING_SNAKE_CASE_		,							exist_ok=SCREAMING_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()		,							SCREAMING_SNAKE_CASE_     )
        print(F'''Save configuration file to {pytorch_config_dump_path}'''     )
        with open(SCREAMING_SNAKE_CASE_		,							"w"		,							encoding="utf-8"     ) as f:
            f.write(config.to_json_string()     )


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










import unittest

import numpy as np

from transformers.testing_utils import require_flax, require_tf, require_torch
from transformers.utils import (
    expand_dims,
    flatten_dict,
    is_flax_available,
    is_tf_available,
    is_torch_available,
    reshape,
    squeeze,
    transpose,
)


if is_flax_available():
					import jax.numpy as jnp

if is_tf_available():
					import tensorflow as tf

if is_torch_available():
					import torch
class     lowerCamelCase_					(   unittest.TestCase			):







				def       A						(    self  ):
								"""simple docstring"""
								__magic_name__		:List[Any]   =   {
								    '''task_specific_params''': {
								        '''summarization''': {'''length_penalty''': 1.0, '''max_length''': 1_2_8, '''min_length''': 1_2, '''num_beams''': 4},
								        '''summarization_cnn''': {'''length_penalty''': 2.0, '''max_length''': 1_4_2, '''min_length''': 5_6, '''num_beams''': 4},
								        '''summarization_xsum''': {'''length_penalty''': 1.0, '''max_length''': 6_2, '''min_length''': 1_1, '''num_beams''': 6},
								    }
								}
								__magic_name__		:List[str]   =   {
								    '''task_specific_params.summarization.length_penalty''': 1.0,
								    '''task_specific_params.summarization.max_length''': 1_2_8,
								    '''task_specific_params.summarization.min_length''': 1_2,
								    '''task_specific_params.summarization.num_beams''': 4,
								    '''task_specific_params.summarization_cnn.length_penalty''': 2.0,
								    '''task_specific_params.summarization_cnn.max_length''': 1_4_2,
								    '''task_specific_params.summarization_cnn.min_length''': 5_6,
								    '''task_specific_params.summarization_cnn.num_beams''': 4,
								    '''task_specific_params.summarization_xsum.length_penalty''': 1.0,
								    '''task_specific_params.summarization_xsum.max_length''': 6_2,
								    '''task_specific_params.summarization_xsum.min_length''': 1_1,
								    '''task_specific_params.summarization_xsum.num_beams''': 6,
								}

								self.assertEqual(flatten_dict(__lowerCAmelCase  )							,       __lowerCAmelCase  )







				def       A						(    self  ):
								"""simple docstring"""
								__magic_name__		:Optional[Any]   =   np.random.randn(3							,       4  )
								self.assertTrue(np.allclose(transpose(__lowerCAmelCase  )							,       x.transpose()  )  )

								__magic_name__		:List[Any]   =   np.random.randn(3							,       4							,       5  )
								self.assertTrue(np.allclose(transpose(__lowerCAmelCase							,       axes=(1, 2, 0)  )							,       x.transpose((1, 2, 0)  )  )  )







				@require_torch
				def       A						(    self  ):
								"""simple docstring"""
								__magic_name__		:Optional[Any]   =   np.random.randn(3							,       4  )
								__magic_name__		:Tuple   =   torch.tensor(__lowerCAmelCase  )
								self.assertTrue(np.allclose(transpose(__lowerCAmelCase  )							,       transpose(__lowerCAmelCase  ).numpy()  )  )

								__magic_name__		:int   =   np.random.randn(3							,       4							,       5  )
								__magic_name__		:Union[str, Any]   =   torch.tensor(__lowerCAmelCase  )
								self.assertTrue(np.allclose(transpose(__lowerCAmelCase							,       axes=(1, 2, 0)  )							,       transpose(__lowerCAmelCase							,       axes=(1, 2, 0)  ).numpy()  )  )







				@require_tf
				def       A						(    self  ):
								"""simple docstring"""
								__magic_name__		:int   =   np.random.randn(3							,       4  )
								__magic_name__		:Optional[Any]   =   tf.constant(__lowerCAmelCase  )
								self.assertTrue(np.allclose(transpose(__lowerCAmelCase  )							,       transpose(__lowerCAmelCase  ).numpy()  )  )

								__magic_name__		:List[str]   =   np.random.randn(3							,       4							,       5  )
								__magic_name__		:Optional[Any]   =   tf.constant(__lowerCAmelCase  )
								self.assertTrue(np.allclose(transpose(__lowerCAmelCase							,       axes=(1, 2, 0)  )							,       transpose(__lowerCAmelCase							,       axes=(1, 2, 0)  ).numpy()  )  )







				@require_flax
				def       A						(    self  ):
								"""simple docstring"""
								__magic_name__		:int   =   np.random.randn(3							,       4  )
								__magic_name__		:Dict   =   jnp.array(__lowerCAmelCase  )
								self.assertTrue(np.allclose(transpose(__lowerCAmelCase  )							,       np.asarray(transpose(__lowerCAmelCase  )  )  )  )

								__magic_name__		:Dict   =   np.random.randn(3							,       4							,       5  )
								__magic_name__		:Dict   =   jnp.array(__lowerCAmelCase  )
								self.assertTrue(np.allclose(transpose(__lowerCAmelCase							,       axes=(1, 2, 0)  )							,       np.asarray(transpose(__lowerCAmelCase							,       axes=(1, 2, 0)  )  )  )  )







				def       A						(    self  ):
								"""simple docstring"""
								__magic_name__		:Any   =   np.random.randn(3							,       4  )
								self.assertTrue(np.allclose(reshape(__lowerCAmelCase							,       (4, 3)  )							,       np.reshape(__lowerCAmelCase							,       (4, 3)  )  )  )

								__magic_name__		:Union[str, Any]   =   np.random.randn(3							,       4							,       5  )
								self.assertTrue(np.allclose(reshape(__lowerCAmelCase							,       (1_2, 5)  )							,       np.reshape(__lowerCAmelCase							,       (1_2, 5)  )  )  )







				@require_torch
				def       A						(    self  ):
								"""simple docstring"""
								__magic_name__		:Dict   =   np.random.randn(3							,       4  )
								__magic_name__		:Tuple   =   torch.tensor(__lowerCAmelCase  )
								self.assertTrue(np.allclose(reshape(__lowerCAmelCase							,       (4, 3)  )							,       reshape(__lowerCAmelCase							,       (4, 3)  ).numpy()  )  )

								__magic_name__		:Union[str, Any]   =   np.random.randn(3							,       4							,       5  )
								__magic_name__		:List[str]   =   torch.tensor(__lowerCAmelCase  )
								self.assertTrue(np.allclose(reshape(__lowerCAmelCase							,       (1_2, 5)  )							,       reshape(__lowerCAmelCase							,       (1_2, 5)  ).numpy()  )  )







				@require_tf
				def       A						(    self  ):
								"""simple docstring"""
								__magic_name__		:Dict   =   np.random.randn(3							,       4  )
								__magic_name__		:Union[str, Any]   =   tf.constant(__lowerCAmelCase  )
								self.assertTrue(np.allclose(reshape(__lowerCAmelCase							,       (4, 3)  )							,       reshape(__lowerCAmelCase							,       (4, 3)  ).numpy()  )  )

								__magic_name__		:List[Any]   =   np.random.randn(3							,       4							,       5  )
								__magic_name__		:Optional[int]   =   tf.constant(__lowerCAmelCase  )
								self.assertTrue(np.allclose(reshape(__lowerCAmelCase							,       (1_2, 5)  )							,       reshape(__lowerCAmelCase							,       (1_2, 5)  ).numpy()  )  )







				@require_flax
				def       A						(    self  ):
								"""simple docstring"""
								__magic_name__		:List[str]   =   np.random.randn(3							,       4  )
								__magic_name__		:Any   =   jnp.array(__lowerCAmelCase  )
								self.assertTrue(np.allclose(reshape(__lowerCAmelCase							,       (4, 3)  )							,       np.asarray(reshape(__lowerCAmelCase							,       (4, 3)  )  )  )  )

								__magic_name__		:List[Any]   =   np.random.randn(3							,       4							,       5  )
								__magic_name__		:List[str]   =   jnp.array(__lowerCAmelCase  )
								self.assertTrue(np.allclose(reshape(__lowerCAmelCase							,       (1_2, 5)  )							,       np.asarray(reshape(__lowerCAmelCase							,       (1_2, 5)  )  )  )  )







				def       A						(    self  ):
								"""simple docstring"""
								__magic_name__		:List[Any]   =   np.random.randn(1							,       3							,       4  )
								self.assertTrue(np.allclose(squeeze(__lowerCAmelCase  )							,       np.squeeze(__lowerCAmelCase  )  )  )

								__magic_name__		:Optional[Any]   =   np.random.randn(1							,       4							,       1							,       5  )
								self.assertTrue(np.allclose(squeeze(__lowerCAmelCase							,       axis=2  )							,       np.squeeze(__lowerCAmelCase							,       axis=2  )  )  )







				@require_torch
				def       A						(    self  ):
								"""simple docstring"""
								__magic_name__		:Dict   =   np.random.randn(1							,       3							,       4  )
								__magic_name__		:List[Any]   =   torch.tensor(__lowerCAmelCase  )
								self.assertTrue(np.allclose(squeeze(__lowerCAmelCase  )							,       squeeze(__lowerCAmelCase  ).numpy()  )  )

								__magic_name__		:List[str]   =   np.random.randn(1							,       4							,       1							,       5  )
								__magic_name__		:str   =   torch.tensor(__lowerCAmelCase  )
								self.assertTrue(np.allclose(squeeze(__lowerCAmelCase							,       axis=2  )							,       squeeze(__lowerCAmelCase							,       axis=2  ).numpy()  )  )







				@require_tf
				def       A						(    self  ):
								"""simple docstring"""
								__magic_name__		:int   =   np.random.randn(1							,       3							,       4  )
								__magic_name__		:Tuple   =   tf.constant(__lowerCAmelCase  )
								self.assertTrue(np.allclose(squeeze(__lowerCAmelCase  )							,       squeeze(__lowerCAmelCase  ).numpy()  )  )

								__magic_name__		:Tuple   =   np.random.randn(1							,       4							,       1							,       5  )
								__magic_name__		:Optional[int]   =   tf.constant(__lowerCAmelCase  )
								self.assertTrue(np.allclose(squeeze(__lowerCAmelCase							,       axis=2  )							,       squeeze(__lowerCAmelCase							,       axis=2  ).numpy()  )  )







				@require_flax
				def       A						(    self  ):
								"""simple docstring"""
								__magic_name__		:Tuple   =   np.random.randn(1							,       3							,       4  )
								__magic_name__		:Optional[Any]   =   jnp.array(__lowerCAmelCase  )
								self.assertTrue(np.allclose(squeeze(__lowerCAmelCase  )							,       np.asarray(squeeze(__lowerCAmelCase  )  )  )  )

								__magic_name__		:List[Any]   =   np.random.randn(1							,       4							,       1							,       5  )
								__magic_name__		:Optional[Any]   =   jnp.array(__lowerCAmelCase  )
								self.assertTrue(np.allclose(squeeze(__lowerCAmelCase							,       axis=2  )							,       np.asarray(squeeze(__lowerCAmelCase							,       axis=2  )  )  )  )







				def       A						(    self  ):
								"""simple docstring"""
								__magic_name__		:Any   =   np.random.randn(3							,       4  )
								self.assertTrue(np.allclose(expand_dims(__lowerCAmelCase							,       axis=1  )							,       np.expand_dims(__lowerCAmelCase							,       axis=1  )  )  )







				@require_torch
				def       A						(    self  ):
								"""simple docstring"""
								__magic_name__		:List[Any]   =   np.random.randn(3							,       4  )
								__magic_name__		:Any   =   torch.tensor(__lowerCAmelCase  )
								self.assertTrue(np.allclose(expand_dims(__lowerCAmelCase							,       axis=1  )							,       expand_dims(__lowerCAmelCase							,       axis=1  ).numpy()  )  )







				@require_tf
				def       A						(    self  ):
								"""simple docstring"""
								__magic_name__		:Union[str, Any]   =   np.random.randn(3							,       4  )
								__magic_name__		:Union[str, Any]   =   tf.constant(__lowerCAmelCase  )
								self.assertTrue(np.allclose(expand_dims(__lowerCAmelCase							,       axis=1  )							,       expand_dims(__lowerCAmelCase							,       axis=1  ).numpy()  )  )







				@require_flax
				def       A						(    self  ):
								"""simple docstring"""
								__magic_name__		:List[str]   =   np.random.randn(3							,       4  )
								__magic_name__		:Tuple   =   jnp.array(__lowerCAmelCase  )
								self.assertTrue(np.allclose(expand_dims(__lowerCAmelCase							,       axis=1  )							,       np.asarray(expand_dims(__lowerCAmelCase							,       axis=1  )  )  )  )






