






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

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

__a      :       Optional[Any]				  =   [
    "mlp.dense_4h_to_h.weight",
    "self_attention.dense.weight",
]



def 			_SCREAMING_SNAKE_CASE						( __lowercase   :			int			,				__lowercase   :			Union[str, Any]    )						->      Dict:





    """simple docstring"""

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



def 			_SCREAMING_SNAKE_CASE						( __lowercase   :			Tuple    )						->      Union[str, Any]:





    """simple docstring"""

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



def 			_SCREAMING_SNAKE_CASE						( __lowercase   :			Optional[int]			,				__lowercase   :			Optional[int]			,				__lowercase   :			Union[str, Any]			,				__lowercase   :			Optional[int]			,				__lowercase   :			Any    )						->      Dict:





    """simple docstring"""

    if bloom_config_file == "":
        __A         =	BloomConfig()
    else:
        __A         =	BloomConfig.from_json_file(__lowercase    )

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

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

        __A         =	None

        __A         =	BloomConfig()

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

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

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

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

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

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

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

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

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

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

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

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

            __A         =	model.load_state_dict(__lowercase			,				strict=__lowercase    )
            assert not other_keys.unexpected_keys, f"The keys {other_keys.unexpected_keys} are unexpected"
            if missing_keys is None:
                __A         =	set(other_keys.missing_keys    )
            else:
                __A         =	missing_keys.intersection(set(other_keys.missing_keys    )    )

        assert not missing_keys, f"The keys {missing_keys} are missing"

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


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

from transformers.testing_utils import CaptureStdout
from transformers.tools.python_interpreter import evaluate



def 			_SCREAMING_SNAKE_CASE						( __lowercase   :			Optional[int]    )						->      List[str]:





    """simple docstring"""

    return x + 2






class       __lowercase      (	unittest.TestCase  ):







     '''simple docstring'''




     def 		lowerCAmelCase_     (				self       :			List[Any]				):


         """simple docstring"""



         __A         =	"""x = 3"""
         __A         =	{}
         __A         =	evaluate(UpperCamelCase_ ,	{} ,	state=UpperCamelCase_				)
         assert result == 3
         self.assertDictEqual(UpperCamelCase_ ,	{"""x""": 3}				)

         __A         =	"""x = y"""
         __A         =	{"""y""": 5}
         __A         =	evaluate(UpperCamelCase_ ,	{} ,	state=UpperCamelCase_				)
         # evaluate returns the value of the last assignment.
         assert result == 5
         self.assertDictEqual(UpperCamelCase_ ,	{"""x""": 5, """y""": 5}				)




     def 		lowerCAmelCase_     (				self       :			List[Any]				):


         """simple docstring"""



         __A         =	"""y = add_two(x)"""
         __A         =	{"""x""": 3}
         __A         =	evaluate(UpperCamelCase_ ,	{"""add_two""": add_two} ,	state=UpperCamelCase_				)
         assert result == 5
         self.assertDictEqual(UpperCamelCase_ ,	{"""x""": 3, """y""": 5}				)

         # Won't work without the tool
         with CaptureStdout() as out:
             __A         =	evaluate(UpperCamelCase_ ,	{} ,	state=UpperCamelCase_				)
         assert result is None
         assert "tried to execute add_two" in out.out




     def 		lowerCAmelCase_     (				self       :			Optional[Any]				):


         """simple docstring"""



         __A         =	"""x = 3"""
         __A         =	{}
         __A         =	evaluate(UpperCamelCase_ ,	{} ,	state=UpperCamelCase_				)
         assert result == 3
         self.assertDictEqual(UpperCamelCase_ ,	{"""x""": 3}				)




     def 		lowerCAmelCase_     (				self       :			Any				):


         """simple docstring"""



         __A         =	"""test_dict = {'x': x, 'y': add_two(x)}"""
         __A         =	{"""x""": 3}
         __A         =	evaluate(UpperCamelCase_ ,	{"""add_two""": add_two} ,	state=UpperCamelCase_				)
         self.assertDictEqual(UpperCamelCase_ ,	{"""x""": 3, """y""": 5}				)
         self.assertDictEqual(UpperCamelCase_ ,	{"""x""": 3, """test_dict""": {"""x""": 3, """y""": 5}}				)




     def 		lowerCAmelCase_     (				self       :			str				):


         """simple docstring"""



         __A         =	"""x = 3\ny = 5"""
         __A         =	{}
         __A         =	evaluate(UpperCamelCase_ ,	{} ,	state=UpperCamelCase_				)
         # evaluate returns the value of the last assignment.
         assert result == 5
         self.assertDictEqual(UpperCamelCase_ ,	{"""x""": 3, """y""": 5}				)




     def 		lowerCAmelCase_     (				self       :			Optional[int]				):


         """simple docstring"""



         __A         =	"""text = f'This is x: {x}.'"""
         __A         =	{"""x""": 3}
         __A         =	evaluate(UpperCamelCase_ ,	{} ,	state=UpperCamelCase_				)
         # evaluate returns the value of the last assignment.
         assert result == "This is x: 3."
         self.assertDictEqual(UpperCamelCase_ ,	{"""x""": 3, """text""": """This is x: 3."""}				)




     def 		lowerCAmelCase_     (				self       :			List[str]				):


         """simple docstring"""



         __A         =	"""if x <= 3:\n    y = 2\nelse:\n    y = 5"""
         __A         =	{"""x""": 3}
         __A         =	evaluate(UpperCamelCase_ ,	{} ,	state=UpperCamelCase_				)
         # evaluate returns the value of the last assignment.
         assert result == 2
         self.assertDictEqual(UpperCamelCase_ ,	{"""x""": 3, """y""": 2}				)

         __A         =	{"""x""": 8}
         __A         =	evaluate(UpperCamelCase_ ,	{} ,	state=UpperCamelCase_				)
         # evaluate returns the value of the last assignment.
         assert result == 5
         self.assertDictEqual(UpperCamelCase_ ,	{"""x""": 8, """y""": 5}				)




     def 		lowerCAmelCase_     (				self       :			Optional[int]				):


         """simple docstring"""



         __A         =	"""test_list = [x, add_two(x)]"""
         __A         =	{"""x""": 3}
         __A         =	evaluate(UpperCamelCase_ ,	{"""add_two""": add_two} ,	state=UpperCamelCase_				)
         self.assertListEqual(UpperCamelCase_ ,	[3, 5]				)
         self.assertDictEqual(UpperCamelCase_ ,	{"""x""": 3, """test_list""": [3, 5]}				)




     def 		lowerCAmelCase_     (				self       :			Dict				):


         """simple docstring"""



         __A         =	"""y = x"""
         __A         =	{"""x""": 3}
         __A         =	evaluate(UpperCamelCase_ ,	{} ,	state=UpperCamelCase_				)
         assert result == 3
         self.assertDictEqual(UpperCamelCase_ ,	{"""x""": 3, """y""": 3}				)




     def 		lowerCAmelCase_     (				self       :			List[Any]				):


         """simple docstring"""



         __A         =	"""test_list = [x, add_two(x)]\ntest_list[1]"""
         __A         =	{"""x""": 3}
         __A         =	evaluate(UpperCamelCase_ ,	{"""add_two""": add_two} ,	state=UpperCamelCase_				)
         assert result == 5
         self.assertDictEqual(UpperCamelCase_ ,	{"""x""": 3, """test_list""": [3, 5]}				)

         __A         =	"""test_dict = {'x': x, 'y': add_two(x)}\ntest_dict['y']"""
         __A         =	{"""x""": 3}
         __A         =	evaluate(UpperCamelCase_ ,	{"""add_two""": add_two} ,	state=UpperCamelCase_				)
         assert result == 5
         self.assertDictEqual(UpperCamelCase_ ,	{"""x""": 3, """test_dict""": {"""x""": 3, """y""": 5}}				)





     def 		lowerCAmelCase_     (				self       :			Any				):


         """simple docstring"""



         __A         =	"""x = 0\nfor i in range(3):\n    x = i"""
         __A         =	{}
         __A         =	evaluate(UpperCamelCase_ ,	{"""range""": range} ,	state=UpperCamelCase_				)
         assert result == 2
         self.assertDictEqual(UpperCamelCase_ ,	{"""x""": 2, """i""": 2}				)




